TsDB2/DB/core/index/primary_handler.go

package index

import (
	"TsDB2/DB/core/index/btree"
	"TsDB2/DB/core/index/datapage"
	"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

	// 尝试复用 DataPage: 只在路由目标叶子内查找
	if rootPID != 0 {
		if dpID, found := btree.FindLastLeafDP(handle, rootPID, rowSize, indexKey); found {
			dp, err := datapage.Load(handle, dpID)
			if err == nil {
				// 用 InsertRowSorted 而不是 AppendRow: 主键由调用方显式给出时
				// (自定义主键 / BaseModel.ID) 插入顺序与 pkID 顺序无关, 而
				// FindRowByPK 依赖"页内行按 pkID 升序"提前终止 —— 追加小号会让
				// 行在页里却查不到 (详见 datapage.InsertRowSorted 的说明)。
				if dp.InsertRowSorted(storeRowData) {
					dp.Save()
					newRoot, _, dup, err := btree.BTreePKInsert(rootPID, indexKey, dpID, handle)
					if err != nil {
						return 0, err
					}
					if !dup {
						*entryPage = newRoot
					}
					return newRoot, nil
				}
			}
		}
	}

	// 无可用空间或树为空: 创建新 DataPage
	dp := datapage.New(handle)
	dp.AppendRow(storeRowData)
	dp.Save()
	newRoot, _, _, err := btree.BTreePKInsert(rootPID, indexKey, dp.PageID(), 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 读取整行。
// 主键 B+ 树的 DataPage 按 pkID 有序, 用 cache (各页 pkID 区间) 二分即可定位;
// cache 未命中时退回 B+ 树路由 (内含全叶扫描兜底), 保证不会漏行。
func (h *PrimaryHandler) ResolveRow(handle page.HandleDb, entryPage uint32, pkID uint64, cache *[]DPRange) ([]byte, bool) {
	pkKeyBE := make([]byte, 8)
	binary.BigEndian.PutUint64(pkKeyBE, pkID)
	if cache != nil {
		if *cache == nil {
			var ranges []DPRange
			for _, pid := range btree.GetDataPageIDs(handle, entryPage) {
				dp, err := datapage.Load(handle, pid)
				if err != nil || dp.Count() == 0 {
					continue
				}
				ranges = append(ranges, DPRange{PageID: pid, MinPK: dp.MinPK(), MaxPK: dp.MaxPK()})
			}
			*cache = ranges
		}
		ranges := *cache
		// 二分: 找到最后一个 MinPK <= pkID 的页 (各页区间有序且不重叠)
		pos := -1
		lo, hi := 0, len(ranges)-1
		for lo <= hi {
			mid := (lo + hi) / 2
			if ranges[mid].MinPK <= pkID {
				pos = mid
				lo = mid + 1
			} else {
				hi = mid - 1
			}
		}
		if pos >= 0 && pkID <= ranges[pos].MaxPK {
			if dp, err := datapage.Load(handle, ranges[pos].PageID); err == nil {
				if row, ok := dp.FindRowByPK(pkKeyBE); ok {
					return row, true
				}
			}
		}
	}
	return h.SearchByPK(handle, entryPage, pkID)
}

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)
	dpID, found := btree.BTreePKSearch(rootPID, pkKeyBE, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			if rowData, ok := dp.FindRowByPK(pkKeyBE); ok {
				return rowData, true
			}
		}
	}
	// 路由失败或未找到, 扫描所有叶子
	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, found := btree.BTreePKSearch(rootPID, indexKey, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			deleted := false
			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() {
					if handle.FreePage != nil {
						handle.FreePage(dpID)
					}
				} else {
					dp.Save()
				}
				return nil
			}
		}
	}
	// 索引键路由失败, 扫描所有叶子
	return btree.ScanAndDelete(rootPID, handle, indexKey, indexKey)
}

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, found := btree.BTreePKSearch(rootPID, pkKeyBE, handle)
	if found {
		dp, err := datapage.Load(handle, dpID)
		if err == nil {
			if dp.DeleteByPK(pkKeyBE) {
				// 同 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() {
					if handle.FreePage != nil {
						handle.FreePage(dpID)
					}
				} else {
					dp.Save()
				}
				return nil
			}
		}
	}
	// 路由失败, 扫描所有叶子
	return btree.ScanAndDeleteByPK(rootPID, handle, pkKeyBE)
}

// 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
	dpID, found := btree.BTreePKSearch(rootPID, pkKeyBE, handle)
	if found {
		if dp, err := datapage.Load(handle, dpID); err == nil {
			if err := dp.UpdateByPK(pkKeyBE, storeRowData); err == nil {
				return nil
			}
			if err := dp.RebuildWith(pkKeyBE, 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)
	newDP.AppendRow(storeRowData)
	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, 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() {
				if handle.FreePage != nil {
					handle.FreePage(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()
}

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