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 // 底层磁盘写

	// 锁表机制
	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 流程统一写回
// 同时清除干净页面缓存，释放内存并确保后续脏页计数准确
func (bp *BufferPool) CollectDirty() []DirtyPage {
	// 1. 锁定写入，防止新的写操作
	bp.LockWrites()
	defer bp.UnlockWrites()

	bp.mu.Lock()
	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 // 重置脏页计数

	// 2. 通知等待容量的 goroutine
	if bp.capacityWait != nil {
		close(bp.capacityWait)
		bp.capacityWait = nil
	}

	bp.mu.Unlock()
	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 {
	// 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. 检查 BufferPool 容量，如果满了则等待
	for bp.capacity > 0 && len(bp.pages) >= bp.capacity {
		// 创建等待 channel
		if bp.capacityWait == nil {
			bp.capacityWait = make(chan struct{})
		}
		waitChan := bp.capacityWait
		bp.mu.Unlock()

		// 阻塞等待空间释放
		<-waitChan

		// 重新获取锁，继续检查
		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++
	}

	// 3. 检查是否超过阈值 (脏页或总页面数)
	totalPages := len(bp.pages)
	if bp.dirtyThreshold > 0 && bp.onThreshold != nil {
		if bp.dirtyCount >= bp.dirtyThreshold || totalPages >= 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()
}