TsDB2/DB/core/index/primary_handler.go

package index

import (
	"git.ycl.pub/yl32940/TsDB2/DB/core/index/btree"
	"git.ycl.pub/yl32940/TsDB2/DB/core/index/datapage"
	"git.ycl.pub/yl32940/TsDB2/DB/core/page"
	"encoding/binary"
)

// PrimaryHandler — 主键 B+ 树索引 (IndexTypePrimary)
// 存储: storeRow = [pkID(8B LE) + rowData]
type PrimaryHandler struct{}

func init() {
	Compose.RegisterHandler(page.IndexTypePrimary, &PrimaryHandler{})

	// 主键 B+ 树分裂参数 —— 写死在本文件, 不随索引元数据传递。
	// 主键多为自增/递增: 顺序追加时右叶只保留 1 条 (SequentialRight), 避免每次
	// 分裂都把整页对半搬运; 非顺序插入仍按 SplitPercent 的比例分裂。
	btree.PKSplitCfg = btree.SplitConfig{SplitPercent: 50, SequentialRight: true}
}

func (h *PrimaryHandler) Insert(handle page.HandleDb, entryPage *uint32, indexKey []byte, pkID uint64, rowData []byte, keyLen int) (uint32, error) {
	// 一次分配即得 [pkID(8B BE) + rowData]: 前 8 字节同时作为 B+ 树 key (只读)
	storeRowData := make([]byte, 8+len(rowData))
	binary.BigEndian.PutUint64(storeRowData[:8], pkID)
	copy(storeRowData[8:], rowData)
	// 主键 B+ 树始终以 pkID(8B BE) 为 key, 与 Search/Delete/Update 保持一致
	return h.insertPrebuilt(handle, entryPage, storeRowData[:8], storeRowData)
}

// insertPrebuilt 写入已构建好的 storeRowData (storeRowData = [pkID(8B BE) + rowData])
func (h *PrimaryHandler) insertPrebuilt(handle page.HandleDb, entryPage *uint32, indexKey []byte, storeRowData []byte) (uint32, error) {
	rootPID := *entryPage
	rowSize := uint32(len(storeRowData)) + 4 // 4B len prefix + row data

	// appendInto 把行写进 dpID 这一页, 并把 (key -> dpID, slot) 挂进 B+ 树叶。
	// 返回 (B+ 树新根, 是否写入成功, 错误) —— "未写入"只表示这一页装不下, 由调用方决定换页。
	//
	// 用 InsertRowSorted 而不是 AppendRow: 主键由调用方显式给出时
	// (自定义主键 / BaseModel.ID) 插入顺序与 pkID 顺序无关, 而
	// FindRowByPK 依赖"页内行按 pkID 升序"提前终止 —— 追加小号会让
	// 行在页里却查不到 (详见 datapage.InsertRowSorted 的说明)。
	//
	// 返回的 slot 会一并写进叶条目, 之后查找/范围扫描就能 O(1) 取行。
	appendInto := func(dpID uint32) (uint32, bool, error) {
		dp, err := datapage.Load(handle, dpID)
		if err != nil {
			return 0, false, nil
		}
		slot, ok := dp.InsertRowSortedSlot(storeRowData)
		if !ok {
			dp.Release()
			return 0, false, nil
		}
		dp.Save()
		dp.Release() // Save 已把整页交给缓冲池, 此后不再访问 dp
		newRoot, _, _, err := btree.BTreePKInsert(rootPID, indexKey, dpID, uint16(slot), handle)
		if err != nil {
			return 0, false, err
		}
		return newRoot, true, nil
	}
	// commit 记录新的树根 (B+ 树未改根时 newRoot == rootPID, 无需回写)
	commit := func(newRoot uint32) uint32 {
		if newRoot != rootPID {
			*entryPage = newRoot
		}
		return newRoot
	}

	// 1) 尝试复用 DataPage: 只在路由目标叶子内查找
	if rootPID != 0 {
		if dpID, found := btree.FindLastLeafDP(handle, rootPID, rowSize, indexKey); found {
			if newRoot, ok, err := appendInto(dpID); err != nil {
				return 0, err
			} else if ok {
				return commit(newRoot), nil
			}
		}

		// 2) 末页判满: 开新页之前先按配置就地整理一次 (墓碑占比达到 CompactDeadRatio 才动手,
		//    配置 <=0 即关闭)。整理会把墓碑占的空间还给活行, 常常就够放下这一行;
		//    CompactForInsert 会先判断"整理后确实装得下", 判断不成立就放弃, 不白搬一趟。
		//
		//    为什么值得做: 被删掉大半的页若一直占着位置, 写入就只能不断开新页 (每页只放
		//    少量行), 全表扫描要读的页数随之膨胀 —— 这正是"删多了以后表越写越散"的来路。
		if dpID, ok := btree.LastLeafDPID(handle, rootPID, indexKey); ok {
			if dp, lerr := datapage.Load(handle, dpID); lerr == nil && dp.CompactForInsert(rowSize) {
				if newRoot, ok, err := appendInto(dpID); err != nil {
					return 0, err
				} else if ok {
					return commit(newRoot), nil
				}
			}
		}
	}

	// 3) 整理后仍装不下 (或本页没有值得回收的墓碑, 或树为空): 创建新 DataPage
	dp := datapage.New(handle)
	slot := dp.AppendRowSlot(storeRowData)
	if slot == 0 {
		// 空页都装不下 → 行本身超长 (调用方应先经 CheckRowSize / 溢出页处理)
		return 0, page.ErrRowTooLarge
	}
	dp.Save()
	newRoot, _, _, err := btree.BTreePKInsert(rootPID, indexKey, dp.PageID(), uint16(slot), handle)
	if err != nil {
		return 0, err
	}
	*entryPage = newRoot
	return newRoot, nil
}

func (h *PrimaryHandler) Search(handle page.HandleDb, entryPage uint32, indexKey []byte, keyLen int, resolve RowResolver) ([]byte, bool) {
	// 主键索引自己就是行数据所在, 不需要回表 (resolve 仅对二级索引有意义)
	if entryPage == 0 || len(indexKey) < 8 {
		return nil, false
	}
	// 主键索引的键就是 pkID (见 Insert: 叶键 = storeRow[:8]), 所以这里必须按 pkID 定位。
	// 旧实现是 dp.FindRowByIndexKey(indexKey, 9) —— 即"拿 key 去和第 1 个字段的值比",
	// 只有当第 1 列的值恒等于 pkID (ORM 的 _id/_ts) 时才凑巧成立; 主键索引声明在别的
	// 列上时, SearchByIndex(pk_.., pkKey(id)) 会查不到行 (数据在, 但 pkID 不可达),
	// 而且它按偏移 9 匹配的是"页内第一行满足者", 大表上还会路由到别的页而漏行。
	// 统一按 pkID 走, 与 SearchByPK/DeleteByPK/Update/ResolveRow 同一套语义。
	return h.SearchByPK(handle, entryPage, binary.BigEndian.Uint64(indexKey[:8]))
}

// IsRowStore 主键索引直接承载行数据
func (h *PrimaryHandler) IsRowStore() bool { return true }

// ResolveRow 按 pkID 读取整行。
//
// 叶条目现在直接给出 (dpID, slot), 一次树下降 + ReadRowBySlotPK 就是 O(1); 而旧的
// 「页区间二分 cache + 页内全扫」在 slot 之后反而更慢 (它仍要逐行扫页), 所以这里统一
// 交给 SearchByPK。cache 参数仅为兼容接口保留 —— 不再需要它来避免重复扫页。
func (h *PrimaryHandler) ResolveRow(handle page.HandleDb, entryPage uint32, pkID uint64, cache *[]DPRange) ([]byte, bool) {
	_ = cache
	return h.SearchByPK(handle, entryPage, pkID)
}

// resolveChunk 批量回表的单批行数: 每批"游标走查 + 物化"完成后才触发上层回调,
// 回调里对行存储索引的 Update/Delete 不会踩坏正在推进的游标。
const resolveChunk = 512

// ResolveMany 批量回表: 见 index.BatchRowResolver 接口契约。
//
// 分批执行: 每批"游标走查 + 物化"完成后才调 sink —— 回调里对行存储索引的
// Update/Delete 不会踩坏正在推进的叶链游标 (踩坏场景: 删除触发叶合并, NextPage
// 被改写, 游标读到回收页; demo/stress_test 的 Scan+Delete 实测踩过)。
// 行快照跨 sink 调用有效 (实现内部缓冲)。
func (h *PrimaryHandler) ResolveMany(handle page.HandleDb, entryPage uint32,
	sortedPKIDs []uint64, sink func(pkID uint64, storeRow []byte) bool) {
	if entryPage == 0 || len(sortedPKIDs) == 0 {
		return
	}
	for start := 0; start < len(sortedPKIDs); start += resolveChunk {
		end := start + resolveChunk
		if end > len(sortedPKIDs) {
			end = len(sortedPKIDs)
		}
		chunk := sortedPKIDs[start:end]
		keys := make([][]byte, len(chunk))
		for i, pkID := range chunk {
			k := make([]byte, 8)
			binary.BigEndian.PutUint64(k, pkID)
			keys[i] = k
		}
		// 本批物化: 游标走查期间【没有】用户回调在跑
		type hit struct {
			pkID uint64
			row  []byte
		}
		var batch []hit
		var lastDP *datapage.DataPage
		lastDPID := uint32(0)
		btree.BTreePKBatchSearch(entryPage, keys, handle,
			func(keyIdx int, dpID uint32, slot uint16) {
				if dpID != lastDPID {
					dp, err := datapage.Load(handle, dpID)
					if err != nil {
						lastDPID = dpID
						lastDP = nil
						return
					}
					lastDP, lastDPID = dp, dpID
				}
				if lastDP == nil {
					return
				}
				if row, ok := lastDP.ReadRowBySlotPK(uint32(slot), keys[keyIdx]); ok {
					batch = append(batch, hit{pkID: chunk[keyIdx], row: row})
				}
			})
		// 走查完成, 现在才触发上层回调
		for _, h := range batch {
			if !sink(h.pkID, h.row) {
				return
			}
		}
	}
}

func (h *PrimaryHandler) SearchByPK(handle page.HandleDb, entryPage uint32, pkID uint64) ([]byte, bool) {
	rootPID := entryPage
	if rootPID == 0 {
		return nil, false
	}
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	// 叶条目带 slot: 命中叶即 O(1) 取行, 无需在页内按 pkID 扫一遍
	dpID, slot, st := btree.BTreePKSearchEx(rootPID, pkKeyBE, handle)
	if st == btree.PKMiss {
		// 正常下降到达叶子且键不在: 该主键不存在, 直接返回。
		// 旧实现在这里无条件全叶扫描 —— 查不存在的键是 O(叶数), 插入查重
		// (rowStoreHasPK) 逐行触发它, 显式主键的批量导入退化成 O(N²)。
		return nil, false
	}
	if st == btree.PKHit {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			if rowData, ok := dp.ReadRowBySlotPK(uint32(slot), pkKeyBE); ok {
				return rowData, true
			}
		}
		// 叶条目命中但槽位取不到行 (slot 失配): 数据页可能被换页/重建过, 仍走兜底确认
	}
	// 路由失败 (树结构异常): 扫描所有叶子
	return btree.ScanAllLeavesForPK(rootPID, handle, pkKeyBE)
}

func (h *PrimaryHandler) Delete(handle page.HandleDb, entryPage *uint32, indexKey []byte, pkID uint64, keyLen int) error {
	rootPID := *entryPage
	if rootPID == 0 {
		return page.ErrNotFound
	}
	dpID, slot, st := btree.BTreePKSearchEx(rootPID, indexKey, handle)
	if st == btree.PKMiss {
		return page.ErrNotFound // 键不存在: 树里没有条目可删, 也无行可删
	}
	if st == btree.PKHit {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			deleted := false
			if _, ok := dp.ReadRowBySlotPK(uint32(slot), indexKey); ok {
				// 槽位命中 → O(1) 删行
				deleted = dp.DeleteBySlot(uint32(slot))
			} else if dp.DeleteByIndexKey(indexKey, 8) {
				deleted = true
			} else if dp.DeleteByPK(indexKey) {
				deleted = true
			}
			if deleted {
				// 摘除本行在 B+ 树叶里的条目 (pk -> dpID)。必须「每次删除都做」,
				// 不能只在 DataPage 变空时做: 旧实现只在 AllDeleted() 时删除
				// 「触发变空的那一个键」, 同一页里先被删掉的行会留下陈旧叶条目。
				// 该 DataPage 随后被回收复用 (例如改作 WAL 页) 时, 这些条目就指向了
				// 非数据页, 空闲链上也挂着仍被索引引用的页
				// (analyzer 报 not-data-page / freeLive, 属数据损坏前兆)。
				newRoot, _, derr := btree.BTreePKDelete(rootPID, indexKey, handle)
				if derr != nil && derr != page.ErrNotFound {
					return derr
				}
				if newRoot != rootPID {
					*entryPage = newRoot
				}
				if dp.AllDeleted() {
					// 延迟还链: 页仍被 B+ 树叶条目引用 (回滚/崩溃重放后树会
					// 复活并继续指向它), 立即还链会让它被复用后读到别人的页
					// (数据损坏)。等 Commit 成功 (或 Close) 后由 pendingFree 统一还。
					handle.DeferFreeDataPage(dpID)
				} else {
					dp.Save()
				}
				return nil
			}
		}
	}
	// 路由失败 (树结构异常), 扫描所有叶子
	newRoot, err := btree.ScanAndDelete(rootPID, handle, indexKey, indexKey)
	if err == nil && newRoot != rootPID {
		*entryPage = newRoot // 兜底删除也可能换根 (树删空/降高), 必须回写
	}
	return err
}

func (h *PrimaryHandler) DeleteByPK(handle page.HandleDb, entryPage *uint32, pkID uint64) error {
	rootPID := *entryPage
	if rootPID == 0 {
		return page.ErrNotFound
	}
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	dpID, slot, st := btree.BTreePKSearchEx(rootPID, pkKeyBE, handle)
	if st == btree.PKMiss {
		return page.ErrNotFound // 键不存在
	}
	if st == btree.PKHit {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			deleted := false
			if _, ok := dp.ReadRowBySlotPK(uint32(slot), pkKeyBE); ok {
				deleted = dp.DeleteBySlot(uint32(slot))
			} else {
				deleted = dp.DeleteByPK(pkKeyBE)
			}
			if deleted {
				// 同 Delete: 每次成功删除都要摘掉本行的叶条目, 否则 DataPage 回收复用后
				// 会留下指向非数据页的陈旧条目 (详见 Delete 处的说明)。
				newRoot, _, derr := btree.BTreePKDelete(rootPID, pkKeyBE, handle)
				if derr != nil && derr != page.ErrNotFound {
					return derr
				}
				if newRoot != rootPID {
					*entryPage = newRoot
				}
				if dp.AllDeleted() {
					// 延迟还链: 页仍被 B+ 树叶条目引用 (回滚/崩溃重放后树会
					// 复活并继续指向它), 立即还链会让它被复用后读到别人的页
					// (数据损坏)。等 Commit 成功 (或 Close) 后由 pendingFree 统一还。
					handle.DeferFreeDataPage(dpID)
				} else {
					dp.Save()
				}
				return nil
			}
		}
	}
	// 路由失败 (树结构异常), 扫描所有叶子
	newRoot, err := btree.ScanAndDeleteByPK(rootPID, handle, pkKeyBE)
	if err == nil && newRoot != rootPID {
		*entryPage = newRoot // 兜底删除也可能换根, 必须回写
	}
	return err
}

// Update 更新主键索引中 pkID 对应的行。
//
// schema 演进 (AddField) 后, 旧行的 schemaVer 与行长都会变长, "原地覆盖"不再可行,
// 因此采用三级策略, 保证更新要么真正生效, 要么返回明确错误:
//
//	1. 原地更新   — 新行长 <= 旧行长 (最常见, 零成本)
//	2. 页内重建   — 页内腾挪并顺带回收墓碑, 支持行变长而无需换页
//	3. 换页       — 页内确实装不下时, 把行搬到新 DataPage 并改写 B+ 树叶条目
func (h *PrimaryHandler) Update(handle page.HandleDb, entryPage *uint32, oldKey, newKey, rowData []byte, pkID uint64, keyLen int) error {
	rootPID := *entryPage
	if rootPID == 0 {
		return page.ErrNotFound
	}
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	// 一次分配构造 [pkID(8B BE) + rowData], 取代 pkBuf+append 的二次分配
	storeRowData := make([]byte, 8+len(rowData))
	copy(storeRowData, pkKeyBE)
	copy(storeRowData[8:], rowData)

	// 1/2 级: B+ 树路由定位 DataPage (slot 直接给出该行的槽位号)
	dpID, slot, st := btree.BTreePKSearchEx(rootPID, pkKeyBE, handle)
	if st == btree.PKMiss {
		return page.ErrNotFound // 键不存在, 无行可更新
	}
	if st == btree.PKHit {
		if dp, err := datapage.Load(handle, dpID); err == nil {
			if err := dp.UpdateBySlot(uint32(slot), storeRowData); err == nil {
				return nil
			}
			if err := dp.RebuildWithSlot(uint32(slot), storeRowData); err == nil {
				return nil
			}
			// 页内装不下 → 直接换页, 不再做无意义的全树扫描
			return h.relocateRow(handle, entryPage, pkKeyBE, storeRowData, dpID)
		}
	}

	// 路由失败: 回退扫描所有叶子 (内部同样含 1/2 级兜底)
	if err := btree.ScanAndUpdate(rootPID, handle, pkKeyBE, storeRowData); err == nil {
		return nil
	}

	// 3 级: 目标行位于"已满且需变长"的页上, 定位后换页
	if oldDPID, ok := btree.FindDataPageIDByPK(rootPID, handle, pkKeyBE); ok {
		return h.relocateRow(handle, entryPage, pkKeyBE, storeRowData, oldDPID)
	}
	return page.ErrNotFound
}

// relocateRow 把 pkID 对应的行搬到新的 DataPage 承载。
//
// 步骤 (无事务, 顺序上优先保证"索引始终可达该行"):
//  1. 新建 DataPage 并写入新行 (此时尚未被索引引用)
//  2. 摘除该键的旧叶条目, 再插入指向新页的条目
//  3. 旧页软删该行; 若旧页已无活行则直接回收
func (h *PrimaryHandler) relocateRow(handle page.HandleDb, entryPage *uint32, pkKeyBE, storeRowData []byte, oldDPID uint32) error {
	newDP := datapage.New(handle)
	newSlot := newDP.AppendRowSlot(storeRowData)
	if newSlot == 0 {
		return page.ErrRowTooLarge
	}
	newDP.Save()
	newDPID := newDP.PageID()

	// BTreePKInsert 遇到已存在的键会判为重复而不更新指向, 故必须先摘除旧条目
	newRoot, _, err := btree.BTreePKDelete(*entryPage, pkKeyBE, handle)
	if err != nil {
		if handle.FreePage != nil {
			handle.FreePage(newDPID)
		}
		return err
	}
	newRoot2, _, _, err := btree.BTreePKInsert(newRoot, pkKeyBE, newDPID, uint16(newSlot), handle)
	if err != nil {
		*entryPage = newRoot
		if handle.FreePage != nil {
			handle.FreePage(newDPID)
		}
		return err
	}
	*entryPage = newRoot2

	// 旧页收尾: 只软删这一行, 其它行仍由各自的叶条目引用
	if oldDP, lerr := datapage.Load(handle, oldDPID); lerr == nil {
		if oldDP.DeleteByPK(pkKeyBE) {
			if oldDP.AllDeleted() {
				// 延迟还链, 理由同 Delete 路径
				handle.DeferFreeDataPage(oldDPID)
			} else {
				oldDP.Save()
			}
		}
	}
	return nil
}

func (h *PrimaryHandler) GetDataPageIDs(handle page.HandleDb, entryPage uint32) []uint32 {
	return btree.GetDataPageIDs(handle, entryPage)
}

func (h *PrimaryHandler) ReadDataPageRows(handle page.HandleDb, dpID uint32) [][]byte {
	dp, err := datapage.Load(handle, dpID)
	if err != nil {
		return nil
	}
	return dp.GetAllRows()
}

// ReadDataPageLiveCount 返回数据页中未删除的行数 (不物化行)。
// 无条件 COUNT(*) 的快速路径; 通过接口断言暴露, 不强加给所有 IndexHandler。
func (h *PrimaryHandler) ReadDataPageLiveCount(handle page.HandleDb, dpID uint32) int {
	dp, err := datapage.Load(handle, dpID)
	if err != nil {
		return 0
	}
	return dp.CountLiveRows()
}

func (h *PrimaryHandler) RangeScan(handle page.HandleDb, entryPage uint32,
	startKey, endKey []byte, startInclusive, endInclusive bool, reverse bool,
	callback func(row []byte) bool) error {
	return btree.BTreePKRangeScan(entryPage, handle, startKey, endKey, startInclusive, endInclusive, reverse, callback)
}

// InsertBatch 批量写入行存储 (主键 B+ 树 + 数据页)。
//
// 两段式:
//   1) 数据页驻留灌行 —— 路由到首行归属的叶末数据页, 加载一次后连续写入,
//      装满才落盘换新页 (逐条路径每行 Load+Save 整页, 16KB×2 memcpy/行;
//      批量路径摊到 ~每页一次)。行内仍走 InsertRowSortedSlot 保持"页内按
//      pkID 升序"不变式。
//   2) 主键树批量挂条目 —— 收集 (pkID → 数据页, 槽位) 后交给 BTreePKBulkInsert。
//
// rows 必须按 PKID 升序 (Table.InsertRowsWithIDs 保证; 自增批量天然满足)。
// 页内整理 (CompactForInsert) 不做: 批量灌行不依赖旧页腾挪, 装不下直接开新页。
func (h *PrimaryHandler) InsertBatch(handle page.HandleDb, entryPage *uint32, rows []BatchRow) error {
	if len(rows) == 0 {
		return nil
	}

	// 数据页: 路由到首行归属叶的末数据页, 驻留复用
	rootPID := *entryPage
	var dp *datapage.DataPage
	if rootPID != 0 {
		firstKey := make([]byte, 8)
		binary.BigEndian.PutUint64(firstKey, rows[0].PKID)
		if dpID, ok := btree.LastLeafDPID(handle, rootPID, firstKey); ok {
			if d, err := datapage.Load(handle, dpID); err == nil {
				dp = d
			}
		}
	}

	items := make([]btree.PKBulkItem, len(rows))
	for i := range rows {
		r := &rows[i]
		// [pkID(8B BE) + rowData]: 前 8 字节同时作为 B+ 树 key (与 Insert 一致)
		storeRow := make([]byte, 8+len(r.RowData))
		binary.BigEndian.PutUint64(storeRow[:8], r.PKID)
		copy(storeRow[8:], r.RowData)

		if dp != nil {
			if slot, ok := dp.InsertRowSortedSlot(storeRow); ok {
				items[i] = btree.PKBulkItem{DPID: dp.PageID(), Slot: uint16(slot)}
				binary.BigEndian.PutUint64(items[i].Key[:], r.PKID)
				continue
			}
			dp.Save()
			dp.Release()
			dp = nil
		}
		dp = datapage.New(handle)
		slot := dp.AppendRowSlot(storeRow)
		if slot == 0 {
			dp.Save()
			dp.Release()
			return page.ErrRowTooLarge
		}
		items[i] = btree.PKBulkItem{DPID: dp.PageID(), Slot: uint16(slot)}
		binary.BigEndian.PutUint64(items[i].Key[:], r.PKID)
	}
	if dp != nil {
		dp.Save()
		dp.Release()
	}

	newRoot, err := btree.BTreePKBulkInsert(rootPID, items, handle)
	if err != nil {
		return err
	}
	if newRoot != *entryPage {
		*entryPage = newRoot
	}
	return nil
}