TsDB2/DB/core/pool/buffer_pool_test.go

package pool

import (
	"sync"
	"testing"
)

// newTestPool 构造一个带内存"磁盘"的缓冲池，便于观察缓存与落盘行为。
func newTestPool(t *testing.T, capacity int) (*BufferPool, *diskStub) {
	t.Helper()
	d := &diskStub{pages: map[uint32][]byte{}}
	return NewBufferPool(capacity, d.read, d.write), d
}

type diskStub struct {
	mu    sync.Mutex
	pages map[uint32][]byte
}

func (d *diskStub) read(pageID uint32, buf []byte) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	if data, ok := d.pages[pageID]; ok {
		copy(buf, data)
	}
	return nil
}

func (d *diskStub) write(pageID uint32, buf []byte) error {
	d.mu.Lock()
	defer d.mu.Unlock()
	data := make([]byte, len(buf))
	copy(data, buf)
	d.pages[pageID] = data
	return nil
}

// TestCollectDirtyKeepsDirtyPagesReadable 收集脏页快照之后, 脏页必须【仍然留在缓存里】
// 且内容可读 —— 这是"并发写 + 提交"崩溃的直接回归点。
//
// 旧实现把脏页移出缓存, 于是从收集完到 WAL 逐页回写完成(中间隔着两次 fsync)这段窗口
// 里, 这些页既不在缓存、也还没落盘: 并发的读一旦未命中就回源, 拿到的是旧内容, 对
// 一页本会话刚分配、盘上从没写过的页则直接是全零。btree 把全零页当成 EntryLen==0 的
// 节点, 接着就在 btPID 里切片越界 panic; 文件还没长到那个偏移时则读出错, 表现为
// "insert: load failed"。
//
// 注意旧版本的这条用例恰好把 bug 当成预期钉住了: 它断言"脏页回落到磁盘, 期望 0"。
func TestCollectDirtyKeepsDirtyPagesReadable(t *testing.T) {
	bp, _ := newTestPool(t, 64)

	_ = bp.WritePage(1, []byte{1}) // 脏页 (盘上还没有)
	bp.UpdatePage(2, []byte{9, 9}) // 干净页 (WAL 回填/回收页回填)
	_ = bp.WritePage(3, []byte{7}) // 脏页 (盘上还没有)

	if bp.Len() != 3 || bp.DirtyCount() != 2 {
		t.Fatalf("收集前期望池内 3 页 / 脏页 2, 实际 %d / %d", bp.Len(), bp.DirtyCount())
	}

	dirty := bp.CollectDirty()
	if len(dirty) != 2 {
		t.Fatalf("期望取出 2 个脏页, 实际 %d", len(dirty))
	}

	// 快照不等于"页离开了池子": 页数与脏页数都不该变。
	if bp.Len() != 3 {
		t.Fatalf("收集后页仍应全部留在池中, 实际 %d", bp.Len())
	}
	if bp.DirtyCount() != 2 {
		t.Fatalf("收集后脏页计数不该清零 (它们还没落盘), 实际 %d", bp.DirtyCount())
	}

	// 干净页 2 命中缓存
	out2 := make([]byte, 16)
	if err := bp.ReadPage(2, out2); err != nil {
		t.Fatalf("ReadPage(2): %v", err)
	}
	if out2[0] != 9 || out2[1] != 9 {
		t.Fatalf("干净页应命中缓存并返回回填值, 实际 %v", out2[:2])
	}

	// 脏页 1 / 3 也必须命中缓存并读到真实内容 —— 拿到 0 就说明又出现"两头都不在"。
	for _, c := range []struct {
		pageID uint32
		want   byte
	}{{1, 1}, {3, 7}} {
		out := make([]byte, 16)
		if err := bp.ReadPage(c.pageID, out); err != nil {
			t.Fatalf("ReadPage(%d): %v", c.pageID, err)
		}
		if out[0] != c.want {
			t.Fatalf("收集快照后读取页 %d 应当是 %d, 实际 %d (页被移出缓存就会出现这种全零读)",
				c.pageID, c.want, out[0])
		}
	}
}

// TestMarkCleanOnlyClearsUnchangedPages 提交成功收脏标记时必须按版本号逐页判定:
// 提交窗口里被写入线程重新改过的页要继续留脏。
//
// 反例就是 ClearAll 式的"提交成功 → 整个缓存标干净": 那些页并不在本笔 WAL 里, 标干净
// 等于宣布"已持久化" —— 新内容再也不会被写盘, 页被淘汰时静默消失。
func TestMarkCleanOnlyClearsUnchangedPages(t *testing.T) {
	bp, _ := newTestPool(t, 64)

	_ = bp.WritePage(1, []byte{1}) // 会被重新改脏
	_ = bp.WritePage(2, []byte{2}) // 保持原样
	dirty := bp.CollectDirty()
	if len(dirty) != 2 {
		t.Fatalf("期望 2 个脏页, 实际 %d", len(dirty))
	}

	// 模拟"提交进行中, 写入线程又改了页 1"
	_ = bp.WritePage(1, []byte{0xFF})

	// 提交成功 → WAL 回填步骤 (WritePageHome) 已在盘上写下快照内容
	bp.MarkClean(dirty)

	if bp.DirtyCount() != 1 {
		t.Fatalf("被重新改脏的页必须继续留脏, 期望脏页 1, 实际 %d", bp.DirtyCount())
	}
	out := make([]byte, 16)
	if err := bp.ReadPage(1, out); err != nil {
		t.Fatalf("ReadPage(1): %v", err)
	}
	if out[0] != 0xFF {
		t.Fatalf("页 1 应保留提交窗口里写入的新内容 0xFF, 实际 %d", out[0])
	}
	// 页 2 没被改过 → 转干净页 (内容不受影响)
	if entry, ok := bp.pages[2]; !ok || entry.dirty {
		t.Fatalf("未被改动的页应转成干净页, 实际 %+v (存在=%v)", entry, ok)
	}
}

// TestFlushWritesAndClearsDirty Flush 必须真的落盘, 并且把写出去的页转成干净页 ——
// CollectDirty 不再移出缓存, 少了收尾这一步, 同一批内容会被下一次提交再写一遍。
func TestFlushWritesAndClearsDirty(t *testing.T) {
	bp, disk := newTestPool(t, 64)

	_ = bp.WritePage(1, []byte{0x11})
	_ = bp.WritePage(2, []byte{0x22})

	bp.Flush()

	if bp.DirtyCount() != 0 {
		t.Fatalf("Flush 后不该还有脏页, 实际 %d", bp.DirtyCount())
	}
	for _, c := range []struct {
		pageID uint32
		want   byte
	}{{1, 0x11}, {2, 0x22}} {
		got, ok := disk.pages[c.pageID]
		if !ok {
			t.Fatalf("页 %d 没有被写盘", c.pageID)
		}
		if got[0] != c.want {
			t.Fatalf("页 %d 落盘内容应当是 %#x, 实际 %#x", c.pageID, c.want, got[0])
		}
	}
}

// TestWritePageThresholdCountsOnlyDirty 验证自动提交阈值只由脏页数决定，
// 常驻的干净页不会把阈值反复顶爆、导致提交频率失控。
func TestWritePageThresholdCountsOnlyDirty(t *testing.T) {
	bp, _ := newTestPool(t, 1000)

	triggers := 0
	bp.SetDirtyThreshold(10, func() { triggers++ })

	// 灌入大量干净页 (模拟大量回收页经 UpdatePage 回填缓存)
	for i := uint32(0); i < 200; i++ {
		bp.UpdatePage(1000+i, []byte{1})
	}

	// 9 张脏页 —— 未达阈值，不应触发
	for i := uint32(0); i < 9; i++ {
		_ = bp.WritePage(i, []byte{2})
	}
	if triggers != 0 {
		t.Fatalf("干净页不应触发阈值, 实际触发 %d 次", triggers)
	}

	// 第 10 张脏页 → 命中阈值，触发一次
	_ = bp.WritePage(9, []byte{2})
	if triggers != 1 {
		t.Fatalf("脏页达阈值应触发 1 次, 实际 %d 次", triggers)
	}
	if bp.DirtyCount() != 10 {
		t.Fatalf("期望脏页计数 10, 实际 %d", bp.DirtyCount())
	}
	if bp.Len() < 200 {
		t.Fatalf("干净页应保留在池中, 实际池内仅 %d 页", bp.Len())
	}
}


// countingRead 包装 diskStub.read 计数 (用于断言"命中缓存后不再走磁盘")。
type countingRead struct {
	d *diskStub
	n int
}

func (c *countingRead) read(pageID uint32, buf []byte) error {
	c.n++
	return c.d.read(pageID, buf)
}

// TestReadPageCachesCleanPages 读未命中必须把磁盘内容作为干净页载入缓存:
// 第二次读同一页必须命中缓存 (磁盘读次数不增加), 这是"LRU BufferPool"读路径
// 收益的主体 (B+ 树下降 / Hash 目录链 / 堆页 Load 都是重复读)。
func TestReadPageCachesCleanPages(t *testing.T) {
	d := &diskStub{pages: map[uint32][]byte{42: {0x2A}}}
	cr := &countingRead{d: d}
	bp := NewBufferPool(8, cr.read, d.write)

	buf := make([]byte, 1)
	if err := bp.ReadPage(42, buf); err != nil {
		t.Fatalf("第一次读: %v", err)
	}
	if buf[0] != 0x2A {
		t.Fatalf("读到 %d, 想要 0x2A", buf[0])
	}
	if cr.n != 1 {
		t.Fatalf("第一次读应走 1 次磁盘, 实际 %d", cr.n)
	}

	// 第二次读: 必须命中缓存
	if err := bp.ReadPage(42, buf); err != nil {
		t.Fatalf("第二次读: %v", err)
	}
	if cr.n != 1 {
		t.Fatalf("第二次读应命中缓存 (磁盘读仍为 1), 实际 %d", cr.n)
	}

	// 载入的是干净页: 不占脏页计数, 可被 LRU 淘汰
	if bp.DirtyCount() != 0 {
		t.Fatalf("读缓存页不应计脏, 实际脏页 %d", bp.DirtyCount())
	}
}

// TestLRUEvictsColdestPage LRU 淘汰必须按"最久未用"顺序: 最近读过的热页留下,
// 最冷的页被淘汰。旧实现按 map 随机序淘汰, 热页与冷页同等概率被丢。
func TestLRUEvictsColdestPage(t *testing.T) {
	d := &diskStub{pages: map[uint32][]byte{}}
	cr := &countingRead{d: d}
	bp := NewBufferPool(3, cr.read, d.write)

	// 三张干净页入池 (模拟读缓存), lru 序: [1, 2, 3]
	bp.UpdatePage(1, []byte{1})
	bp.UpdatePage(2, []byte{2})
	bp.UpdatePage(3, []byte{3})

	// 读页 1 → 它变热: lru 序 [2, 3, 1]
	if err := bp.ReadPage(1, make([]byte, 1)); err != nil {
		t.Fatalf("ReadPage(1): %v", err)
	}

	// 写页 4 (触发容量淘汰): 最冷的页 2 必须先被淘汰
	if err := bp.WritePage(4, []byte{4}); err != nil {
		t.Fatalf("WritePage(4): %v", err)
	}

	if _, ok := bp.pages[2]; ok {
		t.Fatalf("最久未用的页 2 应被淘汰, 却仍在池中")
	}
	for _, pid := range []uint32{1, 3, 4} {
		if _, ok := bp.pages[pid]; !ok {
			t.Fatalf("页 %d 应在池中 (淘汰顺序错误)", pid)
		}
	}

	// 淘汰后链表与 map 必须一致 (无孤儿节点)
	if bp.lru.Len() != len(bp.pages) {
		t.Fatalf("lru 链表 (%d) 与页表 (%d) 长度不一致", bp.lru.Len(), len(bp.pages))
	}

	// 当前 lru 序 [3, 1, 4]: 读页 2 (miss, 重新入池) 应淘汰最冷的页 3
	cr.n = 0
	if err := bp.ReadPage(2, make([]byte, 1)); err != nil {
		t.Fatalf("ReadPage(2): %v", err)
	}
	if cr.n != 1 {
		t.Fatalf("被淘汰的页 2 应重新走磁盘 1 次, 实际 %d 次", cr.n)
	}
	if _, ok := bp.pages[3]; ok {
		t.Fatalf("页 3 现在是最冷的, 应被页 2 的入池挤掉")
	}
	if err := bp.ReadPage(1, make([]byte, 1)); err != nil {
		t.Fatalf("ReadPage(1): %v", err)
	}
	if cr.n != 1 {
		t.Fatalf("热页 1 应命中缓存, 实际磁盘读 %d 次", cr.n)
	}
}

// TestReadPageNeverOverwritesConcurrentDirty 读盘期间另一线程写入了同一页时,
// ReadPage 必须以池内 (更新) 内容为准 —— 绝不能用更旧的读盘结果覆盖脏页。
func TestReadPageNeverOverwritesConcurrentDirty(t *testing.T) {
	d := &diskStub{pages: map[uint32][]byte{7: {0x07}}}
	started := make(chan struct{}, 1)
	release := make(chan struct{})
	bp := NewBufferPool(8, func(pageID uint32, buf []byte) error {
		started <- struct{}{}
		<-release // 卡住读盘, 制造并发窗口
		return d.read(pageID, buf)
	}, d.write)

	out := make([]byte, 1)
	done := make(chan error, 1)
	go func() { done <- bp.ReadPage(7, out) }()

	<-started // 读盘进行中 (页 7 尚未入池)
	// 并发写入同页 (单写入者模型下不该发生, 但缓冲池是并发的最后防线)
	if err := bp.WritePage(7, []byte{0x99}); err != nil {
		t.Fatalf("WritePage(7): %v", err)
	}
	close(release)
	if err := <-done; err != nil {
		t.Fatalf("ReadPage(7): %v", err)
	}

	if out[0] != 0x99 {
		t.Fatalf("读到的是旧读盘内容 %#x, 池内新内容 0x99 被覆盖了", out[0])
	}
	if bp.DirtyCount() != 1 {
		t.Fatalf("页 7 必须保持脏, 实际脏页 %d", bp.DirtyCount())
	}
	// 再读一次: 必须命中且仍是新内容
	if err := bp.ReadPage(7, out); err != nil {
		t.Fatalf("重读: %v", err)
	}
	if out[0] != 0x99 {
		t.Fatalf("重读拿到 %#x", out[0])
	}
}

// TestWriteOverReadCachedPageKeepsLRUConsistent 读缓存的干净页被 WritePage 改脏后,
// 条目原地复用: lru 链表与页表保持一致, 不泄漏孤儿节点。
func TestWriteOverReadCachedPageKeepsLRUConsistent(t *testing.T) {
	d := &diskStub{pages: map[uint32][]byte{5: {0x05}}}
	bp := NewBufferPool(8, d.read, d.write)

	// 读入池 (干净页)
	if err := bp.ReadPage(5, make([]byte, 1)); err != nil {
		t.Fatalf("ReadPage(5): %v", err)
	}
	// 改脏
	if err := bp.WritePage(5, []byte{0x55}); err != nil {
		t.Fatalf("WritePage(5): %v", err)
	}

	if bp.DirtyCount() != 1 {
		t.Fatalf("改脏后脏页应计 1, 实际 %d", bp.DirtyCount())
	}
	if bp.lru.Len() != len(bp.pages) {
		t.Fatalf("lru 链表 (%d) 与页表 (%d) 长度不一致", bp.lru.Len(), len(bp.pages))
	}
	// 淘汰路径上它现在是脏页: 容量满时不得被淘汰
	bp.UpdatePage(100, []byte{1})
	if err := bp.WritePage(6, []byte{6}); err != nil {
		t.Fatalf("WritePage(6): %v", err)
	}
	if _, ok := bp.pages[5]; !ok {
		t.Fatalf("脏页 5 绝不能被容量淘汰")
	}
}