TsDB2/DB/core/pool/buffer_pool.go

package pool

import (
	"container/list"
	"sync"
)

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

// 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)
//
// 缓存策略: 真正的 LRU —— 读/写命中都把页挪到链表尾, 容量满时从链表头淘汰最久
// 未用的【干净】页。读未命中也会把磁盘内容作为干净页载入缓存 (读缓存), 让重复
// 扫描/点查 (B+ 树下降每层、Hash 目录链逐页、堆页 Load) 不必每次都走系统调用。
// 脏页绝不淘汰 (宁可临时超容量, 由脏页阈值触发的 Commit 回落)。
type BufferPool struct {
	pages          map[uint32]*pageEntry // pageID → 页条目
	lru            *list.List            // 淘汰顺序: Front = 最久未用, Back = 最近使用; Value = pageID (uint32)
	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

	// 事务支持
	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),
		lru:       list.New(),
		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 中调用）。
//
// 醒来必须【回头复查】 writeLocked: 锁是成对替换的 (LockWrites 每次新建
// channel), 若本 waiter 持有旧 channel 被放行时, 恰逢下一笔提交已重新加锁,
// 不复查就会穿越进别人的提交窗口写页 —— 快照一致性被破坏。
func (bp *BufferPool) waitForWriteLock() {
	for {
		bp.mu.Lock()
		if !bp.writeLocked {
			bp.mu.Unlock()
			return // 未锁定，直接返回
		}

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

		// 阻塞等待解锁（channel 关闭时会返回）
		<-ch
		// 循环回头复查: 可能已被下一笔 LockWrites 重新锁住
	}
}

// touchLocked 把页标记为"刚被使用过": 挪到 LRU 链表尾。
// 新条目 (lruEl == nil) 入尾; 已在链中的移动到尾。调用方必须持有 bp.mu。
func (bp *BufferPool) touchLocked(e *pageEntry, pageID uint32) {
	if e.lruEl == nil {
		e.lruEl = bp.lru.PushBack(pageID)
		return
	}
	bp.lru.MoveToBack(e.lruEl)
}

// removeLocked 从缓存与 LRU 链表同时摘除一页。调用方必须持有 bp.mu。
func (bp *BufferPool) removeLocked(pageID uint32, e *pageEntry) {
	if e.lruEl != nil {
		bp.lru.Remove(e.lruEl)
		e.lruEl = nil
	}
	delete(bp.pages, pageID)
}

// 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.touchLocked(e, p.PageID)
		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++
		}
		bp.removeLocked(pid, entry)
	}
	bp.dirtyCount -= removedDirty
	if bp.dirtyCount < 0 {
		bp.dirtyCount = 0
	}
}

// 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
		}
	}
	bp.removeLocked(pageID, entry)
	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)
	if old, ok := bp.pages[pageID]; ok {
		// 复用条目与 lru 节点, 只换内容 —— 语义与"换成新条目"相同 (干净页回填),
		// 但不产生孤儿链表节点。
		old.data = cached
		old.dirty = false
		old.pageLen = len(data)
		bp.stampLocked(old) // 内容换过了, 版本号必须往前走, 否则旧快照会误认为这页没变
		bp.touchLocked(old, pageID)
	} else {
		e := &pageEntry{data: cached, dirty: false, pageLen: len(data)}
		bp.stampLocked(e)
		bp.pages[pageID] = e
		bp.touchLocked(e, pageID)
	}
	// 回填的页是干净页, 不受 dirtyCount 约束; 不加这一步, 大量回收页会把缓存
	// 无限撑大 (回收 N 页 → 缓存增加 N 页), 并让 WritePage 的容量判断形同虚设。
	bp.evictCleanForRoomLocked()
}

// evictCleanForRoomLocked 容量已满时按 LRU 顺序淘汰干净页, 直到降到容量以下。
//
// 从链表头 (最久未用) 开始扫描: 干净页直接淘汰 (内容与磁盘一致, 丢弃后下次读取
// 重新载入, 语义不受影响); 脏页跳过 —— 淘汰脏页会静默丢掉尚未落盘的修改。
// 若缓存里全是脏页, 这里不做任何事 —— 允许缓存临时超出容量, 由脏页阈值触发的
// Commit 尽快回落。
//
// 旧实现按 map 遍历顺序随机挑干净页淘汰 (并非 LRU), 热页与冷页被同等概率丢掉。
//
// 调用方必须持有 bp.mu。
func (bp *BufferPool) evictCleanForRoomLocked() {
	if bp.capacity <= 0 || len(bp.pages) < bp.capacity {
		return
	}
	for el := bp.lru.Front(); el != nil && len(bp.pages) >= bp.capacity; {
		next := el.Next() // 先记住下一个节点: removeLocked 会把当前节点从链里摘掉
		pid := el.Value.(uint32)
		entry, ok := bp.pages[pid]
		if !ok || entry.lruEl != el {
			// 孤儿链表节点 (正常不应出现): 清掉继续, 避免无限循环
			bp.lru.Remove(el)
		} else if !entry.dirty {
			bp.removeLocked(pid, entry)
		}
		// 脏页: 跳过, 也不移动到尾部 (移动会让全脏扫描永远走不到旧页之外)
		el = next
	}
}

// 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 读页 — 先查缓存, 未命中则读磁盘并作为【干净页载入缓存】。
//
// 读缓存让重复扫描/点查不必每次都走系统调用; 载入的页与磁盘一致, 可被 LRU 淘汰。
// 读盘期间另一线程可能已把该页写进缓存 (写缓存优先): 重新检查后以池内内容为准,
// 绝不拿这份可能更旧的读盘结果覆盖脏页。直写盘路径 (WritePageDirect) 自带缓存失效,
// 读缓存不会滞留陈旧内容。
func (bp *BufferPool) ReadPage(pageID uint32, buf []byte) error {
	bp.mu.Lock()
	if entry, ok := bp.pages[pageID]; ok {
		copy(buf, entry.data)
		bp.touchLocked(entry, pageID)
		bp.mu.Unlock()
		return nil
	}
	bp.mu.Unlock()

	// 缓存未命中 → 读磁盘
	if err := bp.diskRead(pageID, buf); err != nil {
		return err
	}

	// 作为干净页入池。入池前重新检查: 读盘这段时间里别的线程可能已经写入了这一页
	// (缓存里是更新/未落盘的内容), 此时以池内为准, 丢弃读盘结果。
	bp.mu.Lock()
	if entry, ok := bp.pages[pageID]; ok {
		copy(buf, entry.data)
		bp.touchLocked(entry, pageID)
		bp.mu.Unlock()
		return nil
	}
	bp.evictCleanForRoomLocked()
	entry := &pageEntry{data: make([]byte, len(buf)), dirty: false, pageLen: len(buf)}
	copy(entry.data, buf)
	bp.pages[pageID] = entry
	bp.touchLocked(entry, pageID)
	bp.mu.Unlock()
	return nil
}

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

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

	var triggerThreshold bool
	bp.mu.Lock()

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

	existing, exists := bp.pages[pageID]
	switch {
	case exists && existing.dirty:
		// 已是脏页, 直接覆盖, 不增加脏页计数。
		//
		// 【原地覆写】旧实现在这里每次都新分配一块 len(buf) 的副本再整个换掉
		// existing.data —— 对逐行写入路径 (B+ 树叶 + 数据页 + 表信息页, 每行约 6 次
		// 16KB 写) 意味着每行 ~96KB 的纯垃圾, 5 万行灌数实测 ~4.7GB 分配, GC 压力
		// 吃掉约 1/3 的插入吞吐。这里改为: 页长一致时直接 copy 进旧缓冲 (旧缓冲没有
		// 任何外部别名 —— CollectDirty 的快照是独立副本, UpdatePage/RestoreDirty 各自
		// 持有自己的数据), 只留一次 16KB memcpy, 分配量从"每次写"降到"每张页"。
		if len(existing.data) == len(buf) {
			copy(existing.data, buf)
		} else {
			cached := make([]byte, len(buf))
			copy(cached, buf)
			existing.data = cached
		}
		bp.stampLocked(existing) // 内容变了: 让提交窗口里的旧快照认不出这一页
		bp.touchLocked(existing, pageID)
	case exists:
		// 干净页 (读缓存/回收回填的页) 被改写: 复用条目与 lru 节点, 标脏
		if len(existing.data) == len(buf) {
			copy(existing.data, buf)
		} else {
			cached := make([]byte, len(buf))
			copy(cached, buf)
			existing.data = cached
		}
		existing.dirty = true
		existing.pageLen = len(buf)
		bp.stampLocked(existing)
		bp.dirtyCount++
		bp.touchLocked(existing, pageID)
	default:
		// 新页, 增加脏页计数
		cached := make([]byte, len(buf))
		copy(cached, buf)
		e := &pageEntry{data: cached, dirty: true, pageLen: len(buf)}
		bp.stampLocked(e)
		bp.pages[pageID] = e
		bp.touchLocked(e, pageID)
		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(uint32)) {
	bp.mu.Lock()
	bp.onPageWrite = callback
	bp.mu.Unlock()
}