TsDB2/DB/core/pool/buffer_pool.go

package pool

import (
	"sync"
)

// pageEntry 缓冲池中的页条目
type pageEntry struct {
	data    []byte // 页数据副本
	dirty   bool   // 是否被修改（需要写回磁盘）
	pageLen int    // 页字节长度（用于计算文件偏移）
	ver     uint64 // 内容版本号: 每次内容被写入都换一个 (见 BufferPool.stampLocked)
}

// DirtyPage 脏页快照
//
// ver 是快照当时该页的版本号, 只在 pool 内部使用。提交成功后 MarkClean 靠它判断
// 「这页在提交窗口里有没有被再改过」: 改过就必须继续留脏, 不能被顺手标成干净页。
type DirtyPage struct {
	PageID  uint32
	Data    []byte
	PageLen int
	ver     uint64
}

// BufferPool 页面缓冲池
// 拦截 ReadPage/WritePage，将页缓存在内存中，减少磁盘 I/O
// 写回策略:
//   - Commit() → 通过 CollectDirty() 收集脏页，走 WAL 流程写回
//   - 脏页超过阈值时自动触发 onThreshold 回调 (由 DB 层执行 Commit)
type BufferPool struct {
	pages          map[uint32]*pageEntry // pageID → 页条目
	mu             sync.Mutex
	capacity       int                                   // 最大缓存页数
	dirtyThreshold int                                   // 脏页阈值 (0=不限制)
	onThreshold    func()                                // 脏页超阈值回调
	dirtyCount     int                                   // 当前脏页数 (实时追踪)
	verSeq         uint64                                // 版本号发放器 (见 stampLocked)
	diskRead       func(pageID uint32, buf []byte) error // 底层磁盘读
	diskWrite      func(pageID uint32, buf []byte) error // 底层磁盘写

	// 锁表机制
	writeLocked bool          // 写锁标志
	writeChan   chan struct{} // 用于通知解锁的 channel

	// BufferPool 容量控制
	capacityWait chan struct{} // 容量满时等待的 channel

	// 事务支持
	onPageWrite func(pageID uint32) // 页写入时的回调（用于事务缓存原始数据）
}

// NewBufferPool 创建页面缓冲池
//   - capacity: 最大缓存页数
//   - diskRead/diskWrite: 底层磁盘读写函数
func NewBufferPool(capacity int,
	diskRead func(uint32, []byte) error,
	diskWrite func(uint32, []byte) error,
) *BufferPool {
	bp := &BufferPool{
		pages:     make(map[uint32]*pageEntry, capacity),
		capacity:  capacity,
		diskRead:  diskRead,
		diskWrite: diskWrite,
	}
	return bp
}

// SetDirtyThreshold 设置脏页阈值和回调
// 当脏页数超过 threshold 时调用 onThreshold (由 DB 层触发 Commit)
func (bp *BufferPool) SetDirtyThreshold(threshold int, onThreshold func()) {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	bp.dirtyThreshold = threshold
	bp.onThreshold = onThreshold
}

// LockWrites 锁定写入（Commit/Flush 期间调用）
func (bp *BufferPool) LockWrites() {
	bp.mu.Lock()
	bp.writeLocked = true
	bp.writeChan = make(chan struct{}) // 创建新的 channel
	bp.mu.Unlock()
}

// UnlockWrites 解锁写入（Commit/Flush 完成后调用）
func (bp *BufferPool) UnlockWrites() {
	bp.mu.Lock()
	bp.writeLocked = false
	if bp.writeChan != nil {
		close(bp.writeChan) // 通知所有等待的 goroutine
		bp.writeChan = nil
	}
	bp.mu.Unlock()
}

// waitForWriteLock 等待写锁释放（在 WritePage 中调用）
func (bp *BufferPool) waitForWriteLock() {
	bp.mu.Lock()
	if !bp.writeLocked {
		bp.mu.Unlock()
		return // 未锁定，直接返回
	}

	// 保存 channel 引用
	ch := bp.writeChan
	bp.mu.Unlock()

	// 阻塞等待解锁（channel 关闭时会返回）
	<-ch
}

// CollectDirty 取出所有脏页快照 (由 DB.Commit() 调用，配合 WAL 流程统一写回)。
//
// 快照之后【不把页移出缓存】。移出会造出一个「既不在缓存、也还没落盘」的窗口: 从
// 收集完到 wal.Commit() 逐页回写完成之间要经过两次 fsync, 这期间任何缓存未命中的读
// 都只能回源磁盘, 于是拿到的是旧内容 —— 或者, 对一页本次会话刚分配、盘上从来没写过
// 的页, 拿到全零 (文件还没长到那个偏移时直接读出错)。btree 把全零页当成
// EntryLen==0 的节点, 随后在 btPID 里切片越界 panic; 读出错则表现为
// "insert: load failed"。页留在缓存里, 并发读就一定命中最新内容, 不存在两头都不在的时刻。
//
// 也不再清空 dirtyCount: 快照里的页此刻【尚未落盘】, 计数必须继续按"未持久化"算。
// 页的脏标记由提交方在 wal.Commit() 成功之后用 MarkClean 收掉 —— 只有快照后没被再改过
// (ver 未变) 的页才会转成干净页, 改过的继续留脏, 等下一次提交。
// 干净页照旧保留在池中继续命中 (旧实现无论脏净一律 delete, 等于每次提交都清空缓存,
// 提交后的读取全部回落到磁盘); 容量仍由 evictCleanForRoomLocked 约束。
func (bp *BufferPool) CollectDirty() []DirtyPage {
	// 锁定写入，防止新的写操作
	bp.LockWrites()
	defer bp.UnlockWrites()

	bp.mu.Lock()
	var dirty []DirtyPage
	for pid, entry := range bp.pages {
		if !entry.dirty {
			continue // 干净页: 与磁盘一致, 留在池中继续服务后续读
		}
		data := make([]byte, len(entry.data))
		copy(data, entry.data)
		dirty = append(dirty, DirtyPage{PageID: pid, Data: data, PageLen: entry.pageLen, ver: entry.ver})
	}
	// 这里没有释放任何页, 所以不唤醒等待容量的 goroutine (快照只是复制)。
	bp.mu.Unlock()
	return dirty
}

// MarkClean 把「确实已经落盘」的快照页转成干净页 (提交成功后由提交方调用)。
//
// 判据是版本号而不是"提交成功了": 提交窗口里写入线程可能已经把这一页又改了一遍
// (缓存里是更新的内容, 而本次 WAL 记的是快照那一份)。这种页必须继续留脏 —— 否则它
// 的新内容既不会被写盘, 脏标记还向引擎报告"已经持久化了", 该页一旦被淘汰就永久丢失。
//
// 快照里的页若在窗口里被别的路径丢弃 (回收 / 直写盘), 这里跳过即可: 那些路径自己
// 负责缓存一致性。
func (bp *BufferPool) MarkClean(pages []DirtyPage) {
	if len(pages) == 0 {
		return
	}
	bp.mu.Lock()
	defer bp.mu.Unlock()
	for _, p := range pages {
		entry, ok := bp.pages[p.PageID]
		if !ok || !entry.dirty || entry.ver != p.ver {
			continue
		}
		entry.dirty = false
		bp.dirtyCount--
	}
	if bp.dirtyCount < 0 {
		bp.dirtyCount = 0
	}
}

// stampLocked 给页条目盖一个新的版本号。调用方必须持有 bp.mu。
//
// 版本号只用于回答一个问题: "从我拍快照到现在, 这一页被改过吗"。
func (bp *BufferPool) stampLocked(e *pageEntry) {
	bp.verSeq++
	e.ver = bp.verSeq
}

// RestoreDirty 提交失败时兜底: 保证这批页在缓存里仍然是脏的。
//
// CollectDirty 现在不再把页移出缓存, 所以正常路径上这里无事可做 —— 留着是因为错误
// 处理期间别的动作 (页回收 / 直写盘) 也可能碰过这些页。只补【缺失】的条目: 页已经在
// 缓存里就原样不动, 免得用一份旧快照盖掉提交窗口里写入线程刚改出来的新内容。
//
// 这里刻意不触发阈值回调, 避免在错误处理路径上递归触发 Commit。
func (bp *BufferPool) RestoreDirty(pages []DirtyPage) {
	if len(pages) == 0 {
		return
	}
	bp.mu.Lock()
	defer bp.mu.Unlock()
	for _, p := range pages {
		if _, ok := bp.pages[p.PageID]; ok {
			continue // 已在缓存 (很可能还带着更新的内容), 不要覆盖
		}
		e := &pageEntry{data: p.Data, dirty: true, pageLen: p.PageLen}
		bp.stampLocked(e)
		bp.pages[p.PageID] = e
		bp.dirtyCount++
	}
}

// DiscardPages 丢弃指定页的缓存 (不落盘), 并释放其占用容量。
//
// 事务 Rollback 使用: 事务期间 autoCommit 被 currentTxn 拦住, 任何脏页都不会刷到磁盘,
// 所以「撤销」= 直接把本事务碰过的页从缓存里丢掉, 后续读取自然回到磁盘旧值。
// 只丢弃传入的页 —— 事务开始前遗留的脏页不属于本事务, 必须原样保留。
func (bp *BufferPool) DiscardPages(pageIDs map[uint32]struct{}) {
	if len(pageIDs) == 0 {
		return
	}
	bp.mu.Lock()
	defer bp.mu.Unlock()

	removedDirty := 0
	for pid := range pageIDs {
		entry, ok := bp.pages[pid]
		if !ok {
			continue
		}
		if entry.dirty {
			removedDirty++
		}
		delete(bp.pages, pid)
	}
	bp.dirtyCount -= removedDirty
	if bp.dirtyCount < 0 {
		bp.dirtyCount = 0
	}

	// 释放容量后唤醒等待空间的 goroutine
	if bp.capacityWait != nil {
		close(bp.capacityWait)
		bp.capacityWait = nil
	}
}

// DiscardPage 丢弃单页缓存 (不落盘), 并释放其占用容量。
//
// 用于「绕过缓冲池直写磁盘」(WritePageDirect) 之后, 让缓存中该页的旧副本立即失效,
// 避免随后 ReadPage 命中与磁盘不一致的陈旧内容。
func (bp *BufferPool) DiscardPage(pageID uint32) {
	bp.mu.Lock()
	entry, ok := bp.pages[pageID]
	if !ok {
		bp.mu.Unlock()
		return
	}
	if entry.dirty {
		bp.dirtyCount--
		if bp.dirtyCount < 0 {
			bp.dirtyCount = 0
		}
	}
	delete(bp.pages, pageID)
	if bp.capacityWait != nil {
		close(bp.capacityWait)
		bp.capacityWait = nil
	}
	bp.mu.Unlock()
}

// Flush 将所有脏页写入磁盘 (Close / 维护操作时使用), 并把落盘成功的页转成干净页。
//
// 收尾必须走 MarkClean: CollectDirty 不再移出缓存, 少了这一步, 刚写完的页会一直带着
// 脏标记, 下一次 Commit 会把同一批内容再写一遍 (并且写进 WAL)。
// 写盘错误在这里被忽略 (与调用方既有约定一致), 因此按"已落盘"收尾; ver 比对仍会放过
// 提交窗口里被改过的页, 它们继续留脏, 交给下一次 Flush / Commit。
func (bp *BufferPool) Flush() {
	dirty := bp.CollectDirty()
	for _, item := range dirty {
		_ = bp.diskWrite(item.PageID, item.Data)
	}
	bp.MarkClean(dirty)
}

// Close 关闭缓冲池 (刷盘后停止)
func (bp *BufferPool) Close() {
	bp.Flush()
}

// UpdatePage 更新缓存页 (WAL 提交后调用, 保持缓存与磁盘一致)
func (bp *BufferPool) UpdatePage(pageID uint32, data []byte) {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	cached := make([]byte, len(data))
	copy(cached, data)
	e := &pageEntry{data: cached, dirty: false, pageLen: len(data)}
	bp.stampLocked(e) // 内容换过了, 版本号必须往前走, 否则旧快照会误认为这页没变
	bp.pages[pageID] = e
	// 回填的页是干净页, 不受 dirtyCount 约束; 不加这一步, 大量回收页会把缓存
	// 无限撑大 (回收 N 页 → 缓存增加 N 页), 并让 WritePage 的容量判断形同虚设。
	bp.evictCleanForRoomLocked()
}

// evictCleanForRoomLocked 容量已满时淘汰干净页 (dirty=false), 直到降到容量以下。
//
// 干净页内容与磁盘一致, 丢弃后下次读取会重新从磁盘载入, 语义完全不受影响;
// 脏页绝不能被丢弃 (否则会静默丢掉尚未落盘的修改), 因此若缓存里全是脏页,
// 这里不做任何事 —— 允许缓存临时超出容量, 由脏页阈值触发的 Commit 尽快回落。
//
// 调用方必须持有 bp.mu。
func (bp *BufferPool) evictCleanForRoomLocked() {
	if bp.capacity <= 0 || len(bp.pages) < bp.capacity {
		return
	}
	for pid, entry := range bp.pages {
		if entry.dirty {
			continue
		}
		delete(bp.pages, pid)
		if len(bp.pages) < bp.capacity {
			return
		}
	}
}

// ClearAll 清除所有脏标记, 保留页在缓存中。
//
// ⚠️ 这是"无差别"操作, 只适用于确定没有并发写入的维护场景 (关闭 / 建表删表等)。
// 提交路径上【不能】用它: 提交窗口里被写入线程重新改脏的页并不在本笔 WAL 里, 一并
// 标干净就等于宣布"已持久化", 那份新内容以后既不会被写盘, 也会被淘汰逻辑当作可以
// 直接丢弃的干净页。提交成功要收脏标记请用 MarkClean (按版本号逐页判定)。
func (bp *BufferPool) ClearAll() {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	for _, entry := range bp.pages {
		entry.dirty = false
	}
	bp.dirtyCount = 0
}

// ReadPage 读页 — 先查缓存，未命中则读磁盘 (不缓存, 避免 BufferPool 膨胀)
func (bp *BufferPool) ReadPage(pageID uint32, buf []byte) error {
	bp.mu.Lock()
	if entry, ok := bp.pages[pageID]; ok {
		copy(buf, entry.data)
		bp.mu.Unlock()
		return nil
	}
	bp.mu.Unlock()

	// 缓存未命中 → 直接从磁盘读取, 不写入缓存
	return bp.diskRead(pageID, buf)
}

// WritePage 写页 — 写入缓存并标记脏，不立即落盘
func (bp *BufferPool) WritePage(pageID uint32, buf []byte) error {
	// 1. 等待写锁释放（Commit/Flush 期间阻塞）
	bp.waitForWriteLock()

	// 事务支持：通知 Transaction 缓存原始数据
	if bp.onPageWrite != nil {
		bp.onPageWrite(pageID)
	}

	cached := make([]byte, len(buf))
	copy(cached, buf)

	var triggerThreshold bool
	bp.mu.Lock()

	// 2. 容量控制: 满了先淘汰干净页, 不阻塞等待。
	//    旧实现阻塞在 capacityWait 上等 Commit/Discard 释放空间, 但这两者都不可能
	//    在"当前这个正在阻塞的 goroutine"内被调用 —— 单写入者模型下必然死锁
	//    (fatal error: all goroutines are asleep)。批量删除回收旧页时, realFreePage
	//    会把每个回收页经 SyncCache→UpdatePage 以干净页回填缓存(不计 dirtyCount),
	//    缓存被干净页迅速撑满, 紧接其后的写入即触发该死锁。
	bp.evictCleanForRoomLocked()

	existing, exists := bp.pages[pageID]
	if exists && existing.dirty {
		// 已是脏页, 直接覆盖, 不增加脏页计数
		existing.data = cached
		bp.stampLocked(existing) // 内容变了: 让提交窗口里的旧快照认不出这一页
	} else {
		// 新页或干净页, 增加脏页计数
		e := &pageEntry{data: cached, dirty: true, pageLen: len(buf)}
		bp.stampLocked(e)
		bp.pages[pageID] = e
		bp.dirtyCount++
	}

	// 3. 检查是否超过阈值 —— 只按「脏页数」判定。
	//
	// 不能再用 len(bp.pages): CollectDirty 现在会保留干净页, 缓存里长期驻留一批
	// 与磁盘一致的干净页, 总页数会持续偏高。若仍以总页数触发, 阈值会被这些干净页
	// 反复顶爆, 提交频率凭空放大数倍。脏页数才是真正"尚未落盘"的数据量, 也正是
	// 崩溃时可能丢失的上限; 干净页的容量由 evictCleanForRoomLocked 按 capacity 约束。
	if bp.dirtyThreshold > 0 && bp.onThreshold != nil && bp.dirtyCount >= bp.dirtyThreshold {
		triggerThreshold = true
	}
	bp.mu.Unlock()

	// 4. 在锁外触发回调, 避免死锁
	if triggerThreshold {
		bp.onThreshold()
	}

	return nil
}

// Len 当前缓存页数 (调试用)
func (bp *BufferPool) Len() int {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	return len(bp.pages)
}

// DirtyCount 脏页数 (调试用)
func (bp *BufferPool) DirtyCount() int {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	return bp.dirtyCount
}

// SetCapacity 设置 BufferPool 容量（页数）
func (bp *BufferPool) SetCapacity(capacity int) {
	bp.mu.Lock()
	bp.capacity = capacity
	bp.mu.Unlock()
}

// SetOnPageWriteCallback 设置页写入回调（用于事务缓存原始数据）
func (bp *BufferPool) SetOnPageWriteCallback(callback func(pageID uint32)) {
	bp.mu.Lock()
	bp.onPageWrite = callback
	bp.mu.Unlock()
}