TsDB2/DB/core/index/btree/release_test.go

package btree

import (
	"encoding/binary"
	"sort"
	"testing"

	"TsDB2/DB/core/page"
)

// ============================================================================
// 页回收回归测试 —— 钉死「二级 B+ 树整棵树必须被完整归还」。
//
// 背景 (bench_10m 审计发现的真实缺陷):
//   CollectAllPageIDs 下探内部节点时一律用 btPID(e) 读「偏移 8..12」的子页号。
//   这只对主键条目成立 ([key(8B)][childPID(4B)])。二级 B+ 树的内部条目是变长的
//   ([key(keyLen)][childPID(4B)], entryLen = keyLen+4 或 keyLen+8), 子页号在**末 4 字节**。
//   于是 BFS 在二级树上读到的永远是 0 (或键值字节, 一个从未分配过的垃圾页号):
//   整棵子树只回收了根页, 其余全部成为永久孤儿页 —— 文件只涨不缩。
//
//   实测: 千万行库里 DROP INDEX / 批量删除重建 / DROP TABLE 之后残留 32444 个孤儿
//   索引页 (约 3.2 万页 × 16KB ≈ 500MB), 全部来自二级树子树。
//
// 这些用例用「按键分配的内存页句柄」直接观测「分配了多少页 / 归还了多少页」,
// 不依赖 CollectAllPageIDs 自身去数页 —— 避免循环论证 (若遍历与回收都漏收同一批页,
// 互相印证反而会「通过」)。
// ============================================================================

// memHandle 一个只在内存里保存 16KB 页的 HandleDb:
// 记录每个分配出去的页号与每次归还的页号, 供断言直接比对。
type memHandle struct {
	pages map[uint32][]byte
	next  uint32 // 最近一次分配出去的页号 (0/1 预留, 首个可分配页是 2)
	freed []uint32
}

func newMemHandle() *memHandle {
	return &memHandle{pages: map[uint32][]byte{}, next: 1}
}

func (m *memHandle) handle() page.HandleDb {
	return page.HandleDb{
		NewPage: func() uint32 {
			m.next++
			m.pages[m.next] = make([]byte, page.IndexPageSize)
			return m.next
		},
		SavePage: func(pid uint32, buf []byte) error {
			dst := make([]byte, page.IndexPageSize)
			copy(dst, buf)
			m.pages[pid] = dst
			return nil
		},
		ReadPage: func(pid uint32, buf []byte) error {
			src, ok := m.pages[pid]
			if !ok {
				return page.ErrNotFound
			}
			copy(buf, src)
			return nil
		},
		// 直接归还与延迟归还都记账, 用例只关心「有没有被登记回收」。
		FreePage:      func(pid uint32) { m.freed = append(m.freed, pid) },
		DeferFreePage: func(pid uint32) { m.freed = append(m.freed, pid) },
		// 与引擎一致: 0 是头页, 1 是建库预留表页, >= next 的页从未分配过。
		ValidPage: func(pid uint32) bool { return pid >= 2 && pid < m.next },
	}
}

// allocated 返回全部已分配页号 (升序)。
func (m *memHandle) allocated() []uint32 {
	out := make([]uint32, 0, len(m.pages))
	for pid := range m.pages {
		out = append(out, pid)
	}
	sort.Slice(out, func(i, j int) bool { return out[i] < out[j] })
	return out
}

func sortedSet(ids []uint32) []uint32 {
	seen := map[uint32]bool{}
	out := make([]uint32, 0, len(ids))
	for _, id := range ids {
		if !seen[id] {
			seen[id] = true
			out = append(out, id)
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i] < out[j] })
	return out
}

func sameSet(a, b []uint32) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// assertReclaimsWholeTree 构建一棵树后断言: CollectAllPageIDs 枚举到的页
// == 实际分配出去的全部页 (即整棵树), 且 btReleaseSubtree 把它们全部登记回收。
func assertReclaimsWholeTree(t *testing.T, m *memHandle, root uint32) {
	t.Helper()

	allocated := m.allocated()
	if len(allocated) <= 1 {
		t.Fatalf("样本太小: 只分配了 %d 页, 覆盖不到「多级树」; 请增大条目数", len(allocated))
	}

	got := sortedSet(CollectAllPageIDs(m.handle(), root))
	if !sameSet(got, allocated) {
		t.Fatalf("CollectAllPageIDs 漏收页: 分配到 %d 页 %v, 只枚举出 %d 页 %v\n"+
			"内部节点下探用了错误的子页号偏移 (二级树子页号在末 4 字节, 不在偏移 8..12)",
			len(allocated), allocated, len(got), got)
	}

	m.freed = nil
	btReleaseSubtree(m.handle(), root)
	if freed := sortedSet(m.freed); !sameSet(freed, allocated) {
		t.Fatalf("btReleaseSubtree 漏归还: 应归还 %d 页 %v, 实际归还 %d 页 %v",
			len(allocated), allocated, len(freed), freed)
	}
}

// TestCollectAllPageIDsSecondaryTree 二级 B+ 树 (IndexTypeBTree, 内部条目子页号在末 4 字节)。
//
// 键取 16 字节且第 2 个 8 字节非零, 正是为了把旧实现的读法暴露出来:
// 偏移 8..12 落在键值上 —— 旧的 btPID 会截出一个巨大的、从未分配过的「页号」。
// 这类垃圾页号若不加 ValidPage 兜底就会直接写进空闲链, 那比泄漏更严重 (静默数据损坏)。
func TestCollectAllPageIDsSecondaryTree(t *testing.T) {
	const (
		keyLen = 16
		rows   = 4000 // 二级叶容量 ≈ 681 条 (entryLen=24) → 6 叶 + 1 根
	)
	m := newMemHandle()
	entries := make([][]byte, rows)
	for i := 0; i < rows; i++ {
		e := make([]byte, keyLen+8)
		binary.BigEndian.PutUint64(e[0:8], uint64(i))         // 键分量 1
		binary.BigEndian.PutUint64(e[8:16], uint64(i)*7919+1) // 键分量 2 (非零, 旧实现会误读)
		binary.BigEndian.PutUint64(e[keyLen:], uint64(i)+1)   // PKID
		entries[i] = e
	}
	root := SBTBuildFromSorted(entries, m.handle())
	if root == 0 {
		t.Fatal("SBTBuildFromSorted 返回根页 0")
	}
	assertReclaimsWholeTree(t, m, root)
}

// TestCollectAllPageIDsPrimaryTree 主键 B+ 树 (IndexTypePrimary, 子页号在偏移 8..12)。
// 这条路径本来就对, 一并钉住 —— 修复分流后不能把主键树改坏。
func TestCollectAllPageIDsPrimaryTree(t *testing.T) {
	const rows = 4000 // 主键叶容量 ≈ 1168 条 (entryLen=14) → 4 叶 + 1 根
	m := newMemHandle()
	entries := make([][]byte, rows)
	for i := 0; i < rows; i++ {
		e := make([]byte, btLeafEntryLen)
		binary.BigEndian.PutUint64(e[0:8], uint64(i))
		btSetPID(e, 12)         // dpID (回收路径不看叶条目的 dpID, 只要合法即可)
		btSetSlot(e, uint16(1)) // slot
		entries[i] = e
	}
	root := BuildFromSorted(entries, m.handle())
	if root == 0 {
		t.Fatal("BuildFromSorted 返回根页 0")
	}
	assertReclaimsWholeTree(t, m, root)
}