package srdb import ( "encoding/binary" "fmt" "os" "slices" "sort" "github.com/edsrzf/mmap-go" ) /* B+Tree 存储格式 B+Tree 用于索引 SSTable 和 Index 文件，提供 O(log n) 查询性能。 节点结构 (4096 bytes): ┌─────────────────────────────────────────────────────────────┐ │ Node Header (32 bytes) │ │ ├─ NodeType (1 byte): 0=Internal, 1=Leaf │ │ ├─ KeyCount (2 bytes): 节点中的 key 数量 │ │ ├─ Level (1 byte): 层级 (0=叶子层) │ │ └─ Reserved (28 bytes): 预留空间 │ ├─────────────────────────────────────────────────────────────┤ │ Keys Array (variable) │ │ └─ Key[0..KeyCount-1]: int64 (8 bytes each) │ ├─────────────────────────────────────────────────────────────┤ │ Values (variable, 取决于节点类型) │ │ │ │ 内部节点 (Internal Node): │ │ └─ Children[0..KeyCount]: int64 (8 bytes each) │ │ - 子节点的文件偏移量 │ │ - Children[i] 包含 < Key[i] 的所有 key │ │ - Children[KeyCount] 包含 >= Key[KeyCount-1] 的 key │ │ │ │ 叶子节点 (Leaf Node): │ │ └─ Data Pairs[0..KeyCount-1]: 交错存储 (12 bytes each) │ │ ├─ DataOffset: int64 (8 bytes) - 数据块的文件偏移量 │ │ └─ DataSize: int32 (4 bytes) - 数据块的大小 │ └─────────────────────────────────────────────────────────────┘ 节点头格式 (32 bytes): Offset | Size | Field | Description -------|------|------------|---------------------------------- 0 | 1 | NodeType | 0=Internal, 1=Leaf 1 | 2 | KeyCount | key 数量 (0 ~ BTreeOrder) 3 | 1 | Level | 层级 (0=叶子层, 1+=内部层) 4 | 28 | Reserved | 预留空间 内部节点布局 (示例: KeyCount=3): [Header: 32B] [Keys: Key0(8B), Key1(8B), Key2(8B)] [Children: Child0(8B), Child1(8B), Child2(8B), Child3(8B)] 查询规则: - key < Key0 → Child0 - Key0 ≤ key < Key1 → Child1 - Key1 ≤ key < Key2 → Child2 - key ≥ Key2 → Child3 叶子节点布局 (示例: KeyCount=3): [Header: 32B] [Keys: Key0(8B), Key1(8B), Key2(8B)] [Data: (Offset0, Size0), (Offset1, Size1), (Offset2, Size2)] - 交错存储: Offset0(8B), Size0(4B), Offset1(8B), Size1(4B), Offset2(8B), Size2(4B) 查询规则: - 找到 key == Key[i] - 返回 (DataOffsets[i], DataSizes[i]) B+Tree 特性: - 阶数 (Order): 200 (每个节点最多 200 个 key) - 节点大小: 4096 bytes (4 KB，对齐页大小) - 高度: log₂₀₀(N) (100万条数据约 3 层) - 查询复杂度: O(log n) - 范围查询: 支持（叶子节点有序） 文件布局示例: SSTable/Index 文件: [Header: 256B] [B+Tree Nodes: 4KB each] ├─ Root Node (Internal) ├─ Level 1 Nodes (Internal) └─ Leaf Nodes [Data Blocks: variable] 性能优化: - mmap 零拷贝: 直接从内存映射读取节点 - 节点对齐: 4KB 对齐，利用操作系统页缓存 - 有序存储: 叶子节点有序，支持范围查询 - 紧凑编码: 最小化节点大小，提高缓存命中率 */ const ( BTreeNodeSize = 4096 // 节点大小 (4 KB) BTreeOrder = 200 // B+Tree 阶数 (保守估计，叶子节点每个entry 20 bytes) BTreeHeaderSize = 32 // 节点头大小 BTreeNodeTypeInternal = 0 // 内部节点 BTreeNodeTypeLeaf = 1 // 叶子节点 ) // BTreeNode 表示一个 B+Tree 节点 (4 KB) type BTreeNode struct { // Header (32 bytes) NodeType byte // 0=Internal, 1=Leaf KeyCount uint16 // key 数量 Level byte // 层级 (0=叶子层) Reserved [28]byte // 预留字段 // Keys (variable, 最多 256 个) Keys []int64 // key 数组 // Values (variable) // Internal Node: 子节点指针 Children []int64 // 子节点的文件 offset // Leaf Node: 数据位置 DataOffsets []int64 // 数据块的文件 offset DataSizes []int32 // 数据块大小 } // NewInternalNode 创建内部节点 func NewInternalNode(level byte) *BTreeNode { return &BTreeNode{ NodeType: BTreeNodeTypeInternal, Level: level, Keys: make([]int64, 0, BTreeOrder), Children: make([]int64, 0, BTreeOrder+1), } } // NewLeafNode 创建叶子节点 func NewLeafNode() *BTreeNode { return &BTreeNode{ NodeType: BTreeNodeTypeLeaf, Level: 0, Keys: make([]int64, 0, BTreeOrder), DataOffsets: make([]int64, 0, BTreeOrder), DataSizes: make([]int32, 0, BTreeOrder), } } // Marshal 序列化节点到 4 KB // // 布局： // [Header: 32B] // [Keys: KeyCount * 8B] // [Values: 取决于节点类型] // - Internal: Children (KeyCount+1) * 8B // - Leaf: 交错存储 (Offset, Size) 对，每对 12B，共 KeyCount * 12B // // 示例（叶子节点，KeyCount=3）： // Offset | Size | Content // -------|------|---------------------------------- // 0 | 1 | NodeType = 1 (Leaf) // 1 | 2 | KeyCount = 3 // 3 | 1 | Level = 0 // 4 | 28 | Reserved // 32 | 24 | Keys [100, 200, 300] // 56 | 8 | DataOffset0 = 1000 // 64 | 4 | DataSize0 = 50 // 68 | 8 | DataOffset1 = 2000 // 76 | 4 | DataSize1 = 60 // 80 | 8 | DataOffset2 = 3000 // 88 | 4 | DataSize2 = 70 // 92 | 4004 | Padding (unused) func (n *BTreeNode) Marshal() []byte { buf := make([]byte, BTreeNodeSize) // 写入 Header (32 bytes) buf[0] = n.NodeType binary.LittleEndian.PutUint16(buf[1:3], n.KeyCount) buf[3] = n.Level copy(buf[4:32], n.Reserved[:]) // 写入 Keys offset := BTreeHeaderSize for _, key := range n.Keys { if offset+8 > BTreeNodeSize { break } binary.LittleEndian.PutUint64(buf[offset:offset+8], uint64(key)) offset += 8 } // 写入 Values if n.NodeType == BTreeNodeTypeInternal { // Internal Node: 写入子节点指针 for _, child := range n.Children { if offset+8 > BTreeNodeSize { break } binary.LittleEndian.PutUint64(buf[offset:offset+8], uint64(child)) offset += 8 } } else { // Leaf Node: 写入数据位置 for i := 0; i < len(n.Keys); i++ { if offset+12 > BTreeNodeSize { break } binary.LittleEndian.PutUint64(buf[offset:offset+8], uint64(n.DataOffsets[i])) offset += 8 binary.LittleEndian.PutUint32(buf[offset:offset+4], uint32(n.DataSizes[i])) offset += 4 } } return buf } // UnmarshalBTree 从字节数组反序列化节点 // // 参数： // data: 4KB 节点数据（通常来自 mmap） // // 返回： // *BTreeNode: 反序列化后的节点 // // 零拷贝优化： // - 直接从 mmap 数据读取，不复制整个节点 // - 只复制必要的字段（Keys, Children, DataOffsets, DataSizes） func UnmarshalBTree(data []byte) *BTreeNode { if len(data) < BTreeNodeSize { return nil } node := &BTreeNode{} // 读取 Header node.NodeType = data[0] node.KeyCount = binary.LittleEndian.Uint16(data[1:3]) node.Level = data[3] copy(node.Reserved[:], data[4:32]) // 读取 Keys offset := BTreeHeaderSize node.Keys = make([]int64, node.KeyCount) for i := 0; i < int(node.KeyCount); i++ { if offset+8 > len(data) { break } node.Keys[i] = int64(binary.LittleEndian.Uint64(data[offset : offset+8])) offset += 8 } // 读取 Values if node.NodeType == BTreeNodeTypeInternal { // Internal Node: 读取子节点指针 childCount := int(node.KeyCount) + 1 node.Children = make([]int64, childCount) for i := range childCount { if offset+8 > len(data) { break } node.Children[i] = int64(binary.LittleEndian.Uint64(data[offset : offset+8])) offset += 8 } } else { // Leaf Node: 读取数据位置 node.DataOffsets = make([]int64, node.KeyCount) node.DataSizes = make([]int32, node.KeyCount) for i := 0; i < int(node.KeyCount); i++ { if offset+12 > len(data) { break } node.DataOffsets[i] = int64(binary.LittleEndian.Uint64(data[offset : offset+8])) offset += 8 node.DataSizes[i] = int32(binary.LittleEndian.Uint32(data[offset : offset+4])) offset += 4 } } return node } // IsFull 检查节点是否已满 func (n *BTreeNode) IsFull() bool { return len(n.Keys) >= BTreeOrder } // AddKey 添加 key (仅用于构建) func (n *BTreeNode) AddKey(key int64) { n.Keys = append(n.Keys, key) n.KeyCount = uint16(len(n.Keys)) } // AddChild 添加子节点 (仅用于内部节点) func (n *BTreeNode) AddChild(offset int64) error { if n.NodeType != BTreeNodeTypeInternal { return fmt.Errorf("AddChild called on leaf node") } n.Children = append(n.Children, offset) return nil } // AddData 添加数据位置 (仅用于叶子节点) func (n *BTreeNode) AddData(key int64, offset int64, size int32) error { if n.NodeType != BTreeNodeTypeLeaf { return fmt.Errorf("AddData called on internal node") } n.Keys = append(n.Keys, key) n.DataOffsets = append(n.DataOffsets, offset) n.DataSizes = append(n.DataSizes, size) n.KeyCount = uint16(len(n.Keys)) return nil } // BTreeBuilder 从下往上构建 B+Tree // // 构建流程： // 1. Add(): 添加所有 (key, offset, size) 到叶子节点 // - 当叶子节点满时，创建新的叶子节点 // - 所有叶子节点按 key 有序 // // 2. Build(): 从叶子层向上构建 // - Level 0: 叶子节点（已创建） // - Level 1: 为叶子节点创建父节点（内部节点） // - Level 2+: 递归创建更高层级 // - 最终返回根节点偏移量 // // 示例（100 个 key，Order=200）： // - 叶子层: 1 个叶子节点（100 个 key） // - 根节点: 叶子节点本身 // // 示例（500 个 key，Order=200）： // - 叶子层: 3 个叶子节点（200, 200, 100 个 key） // - Level 1: 1 个内部节点（3 个子节点） // - 根节点: Level 1 的内部节点 type BTreeBuilder struct { order int // B+Tree 阶数 file *os.File // 输出文件 offset int64 // 当前写入位置 leafNodes []*BTreeNode // 叶子节点列表 } // NewBTreeBuilder 创建构建器 func NewBTreeBuilder(file *os.File, startOffset int64) *BTreeBuilder { return &BTreeBuilder{ order: BTreeOrder, file: file, offset: startOffset, leafNodes: make([]*BTreeNode, 0), } } // Add 添加一个 key-value 对 (数据必须已排序) func (b *BTreeBuilder) Add(key int64, dataOffset int64, dataSize int32) error { // 获取或创建当前叶子节点 var leaf *BTreeNode if len(b.leafNodes) == 0 || b.leafNodes[len(b.leafNodes)-1].IsFull() { // 创建新的叶子节点 leaf = NewLeafNode() b.leafNodes = append(b.leafNodes, leaf) } else { leaf = b.leafNodes[len(b.leafNodes)-1] } // 添加到叶子节点 if err := leaf.AddData(key, dataOffset, dataSize); err != nil { return err } return nil } // Build 构建完整的 B+Tree，返回根节点的 offset func (b *BTreeBuilder) Build() (rootOffset int64, err error) { if len(b.leafNodes) == 0 { return 0, nil } // 1. 写入所有叶子节点，记录它们的 offset leafOffsets := make([]int64, len(b.leafNodes)) for i, leaf := range b.leafNodes { leafOffsets[i] = b.offset data := leaf.Marshal() _, err := b.file.WriteAt(data, b.offset) if err != nil { return 0, err } b.offset += BTreeNodeSize } // 2. 如果只有一个叶子节点，它就是根 if len(b.leafNodes) == 1 { return leafOffsets[0], nil } // 3. 从下往上构建内部节点 currentLevel := b.leafNodes currentOffsets := leafOffsets level := 1 for len(currentLevel) > 1 { nextLevel, nextOffsets, err := b.buildLevel(currentLevel, currentOffsets, level) if err != nil { return 0, err } currentLevel = nextLevel currentOffsets = nextOffsets level++ } // 4. 返回根节点的 offset return currentOffsets[0], nil } // buildLevel 构建一层内部节点 func (b *BTreeBuilder) buildLevel(children []*BTreeNode, childOffsets []int64, level int) ([]*BTreeNode, []int64, error) { var parents []*BTreeNode var parentOffsets []int64 // 每 order 个子节点创建一个父节点 for i := 0; i < len(children); i += b.order { end := min(i+b.order, len(children)) // 创建父节点 parent := NewInternalNode(byte(level)) // 添加第一个子节点 (没有对应的 key) if err := parent.AddChild(childOffsets[i]); err != nil { return nil, nil, err } // 添加剩余的子节点和分隔 key for j := i + 1; j < end; j++ { // 分隔 key 是子节点的第一个 key separatorKey := children[j].Keys[0] parent.AddKey(separatorKey) if err := parent.AddChild(childOffsets[j]); err != nil { return nil, nil, err } } // 写入父节点 parentOffset := b.offset data := parent.Marshal() _, err := b.file.WriteAt(data, b.offset) if err != nil { return nil, nil, err } b.offset += BTreeNodeSize parents = append(parents, parent) parentOffsets = append(parentOffsets, parentOffset) } return parents, parentOffsets, nil } // BTreeReader 用于查询 B+Tree (mmap) // // 查询流程： // 1. 从根节点开始 // 2. 如果是内部节点： // - 二分查找确定子节点 // - 跳转到子节点继续查找 // 3. 如果是叶子节点： // - 二分查找 key // - 返回 (dataOffset, dataSize) // // 性能优化： // - mmap 零拷贝：直接从内存映射读取节点 // - 二分查找：O(log KeyCount) 在节点内查找 // - 总复杂度：O(log n) = O(height * log Order) // // 示例（100万条数据，Order=200）： // - 高度: log₂₀₀(1000000) ≈ 3 // - 查询次数: 3 次节点读取 + 3 次二分查找 type BTreeReader struct { mmap mmap.MMap rootOffset int64 } // NewBTreeReader 创建查询器 func NewBTreeReader(mmap mmap.MMap, rootOffset int64) *BTreeReader { return &BTreeReader{ mmap: mmap, rootOffset: rootOffset, } } // Get 查询 key，返回数据位置 // // 参数： // key: 要查询的 key // // 返回： // dataOffset: 数据块的文件偏移量 // dataSize: 数据块的大小 // found: 是否找到 // // 查询流程： // 1. 从根节点开始遍历 // 2. 内部节点：二分查找确定子节点，跳转 // 3. 叶子节点：二分查找 key，返回数据位置 func (r *BTreeReader) Get(key int64) (dataOffset int64, dataSize int32, found bool) { if r.rootOffset == 0 { return 0, 0, false } nodeOffset := r.rootOffset for { // 读取节点 (零拷贝) if nodeOffset+BTreeNodeSize > int64(len(r.mmap)) { return 0, 0, false } nodeData := r.mmap[nodeOffset : nodeOffset+BTreeNodeSize] node := UnmarshalBTree(nodeData) if node == nil { return 0, 0, false } // 叶子节点 if node.NodeType == BTreeNodeTypeLeaf { // 二分查找 idx := sort.Search(len(node.Keys), func(i int) bool { return node.Keys[i] >= key }) if idx < len(node.Keys) && node.Keys[idx] == key { return node.DataOffsets[idx], node.DataSizes[idx], true } return 0, 0, false } // 内部节点，继续向下 // keys[i] 是分隔符，children[i] 包含 < keys[i] 的数据 // children[i+1] 包含 >= keys[i] 的数据 idx := sort.Search(len(node.Keys), func(i int) bool { return node.Keys[i] > key }) // idx 现在指向第一个 > key 的位置 // 我们应该走 children[idx] if idx >= len(node.Children) { idx = len(node.Children) - 1 } nodeOffset = node.Children[idx] } } // GetAllKeys 获取 B+Tree 中所有的 key（按顺序） func (r *BTreeReader) GetAllKeys() []int64 { if r.rootOffset == 0 { return nil } var keys []int64 r.traverseLeafNodes(r.rootOffset, func(node *BTreeNode) { keys = append(keys, node.Keys...) }) // 显式排序以确保返回的 keys 严格有序 // 虽然 B+Tree 构建时应该已经是有序的，但这是一个安全保障 // 特别是在 compaction 后，确保查询结果正确排序 slices.Sort(keys) return keys } // traverseLeafNodes 遍历所有叶子节点 func (r *BTreeReader) traverseLeafNodes(nodeOffset int64, callback func(*BTreeNode)) { if nodeOffset+BTreeNodeSize > int64(len(r.mmap)) { return } nodeData := r.mmap[nodeOffset : nodeOffset+BTreeNodeSize] node := UnmarshalBTree(nodeData) if node == nil { return } if node.NodeType == BTreeNodeTypeLeaf { // 叶子节点，执行回调 callback(node) } else { // 内部节点，递归遍历所有子节点 for _, childOffset := range node.Children { r.traverseLeafNodes(childOffset, callback) } } }