TsDB2/DB/core/pool/buffer_pool.go

package pool

import (
	"sync"
)

// pageEntry 缓冲池中的页条目
type pageEntry struct {
	data    []byte // 页数据副本
	dirty   bool   // 是否被修改（需要写回磁盘）
	pageLen int    // 页字节长度（用于计算文件偏移）
}

// DirtyPage 脏页快照
type DirtyPage struct {
	PageID  uint32
	Data    []byte
	PageLen int
}

// 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                                   // 当前脏页数 (实时追踪)
	diskRead       func(pageID uint32, buf []byte) error // 底层磁盘读
	diskWrite      func(pageID uint32, buf []byte) error // 底层磁盘写
}

// 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
}

// CollectDirty 取出所有脏页快照并清除脏标记
// 由 DB.Commit() 调用，配合 WAL 流程统一写回
// 同时清除干净页面缓存，释放内存并确保后续脏页计数准确
func (bp *BufferPool) CollectDirty() []DirtyPage {
	bp.mu.Lock()
	defer bp.mu.Unlock()
	var dirty []DirtyPage
	for pid, entry := range bp.pages {
		if entry.dirty {
			data := make([]byte, len(entry.data))
			copy(data, entry.data)
			dirty = append(dirty, DirtyPage{pid, data, entry.pageLen})
		}
		// 无论脏净，都从缓存中移除
		// 脏页已被复制，干净页不需要缓存
		// 下次访问时会从磁盘重新加载
		delete(bp.pages, pid)
	}
	bp.dirtyCount = 0 // 重置脏页计数
	return dirty
}

// Flush 将所有脏页写入磁盘 (仅在 Close 时使用)
func (bp *BufferPool) Flush() {
	dirty := bp.CollectDirty()
	for _, item := range dirty {
		_ = bp.diskWrite(item.PageID, item.Data)
	}
}

// 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)
	bp.pages[pageID] = &pageEntry{data: cached, dirty: false, pageLen: len(data)}
}

// ClearAll 清除所有脏标记 (WAL 提交后调用)
// 保留干净页在缓存中, 避免下次 WritePage 重新缓存导致 BufferPool 反复膨胀
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 {
	cached := make([]byte, len(buf))
	copy(cached, buf)

	var triggerThreshold bool
	bp.mu.Lock()
	existing, exists := bp.pages[pageID]
	if exists && existing.dirty {
		// 已是脏页, 直接覆盖, 不增加脏页计数
		existing.data = cached
	} else {
		// 新页或干净页, 增加脏页计数
		bp.pages[pageID] = &pageEntry{data: cached, dirty: true, pageLen: len(buf)}
		bp.dirtyCount++
	}

	// 检查是否超过阈值 (脏页或总页面数)
	// 使用总页面数作为第二道安全网，防止脏页计数增长缓慢时内存无限膨胀
	totalPages := len(bp.pages)
	if bp.dirtyThreshold > 0 && bp.onThreshold != nil {
		// 脏页阈值或总页面数阈值 (两者相同，确保及时刷盘)
		if bp.dirtyCount >= bp.dirtyThreshold || totalPages >= bp.dirtyThreshold {
			triggerThreshold = true
		}
	}
	bp.mu.Unlock()

	// 在锁外触发回调, 避免死锁
	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
}