3. 基于B+树的磁盘存储引擎

基于B+树的磁盘存储引擎 学习卡片

任务说明

实现一个基于B+树的磁盘存储引擎核心模块,支持键值对的插入、查找、范围扫描与持久化。数据以定长页(4KB)为单位存储于单个磁盘文件中,B+树索引与数据页共用同一文件空间管理。本案例聚焦存储引擎的页管理、B+树索引结构与基本读写路径,不涉及事务、并发控制与崩溃恢复。

核心设计

文件按4KB页划分,页号从0开始。页类型分为:文件头页(页0)、B+树内部节点页、叶子节点页。每个节点页可容纳ORDER-1个键(ORDER=8),叶子节点额外存储值(固定8字节)。所有节点在内存中修改后通过write_page写回磁盘,读取时通过read_page加载。

页0 文件头 页1 根(内部) 页2 内部节点 页3 叶子 页4 叶子 页结构 (4KB) ┌──────────┬──────────────┬──────────────────────┐ │ 页头 32B │ 键数组 │ 子指针/值数组 │ │ type,count│ int64 键 │ 页号/固定8B值 │ └──────────┴──────────────┴──────────────────────┘ 内部节点: 键[i] 与 child[i] 对应,child[count] 为最右指针 叶子节点: 键与值一一对应,next_leaf 指向右兄弟页号

代码实现


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#define PAGE_SIZE 4096
#define ORDER 8
#define MAX_KEYS (ORDER - 1)
#define LEAF 1
#define INTERNAL 2

typedef struct {
    uint32_t type;        // LEAF 或 INTERNAL
    uint32_t count;       // 键数量
    uint32_t next_leaf;   // 叶子节点的右兄弟页号
    uint32_t parent;      // 父节点页号
} PageHeader;

typedef struct {
    int64_t keys[MAX_KEYS];
    union {
        uint32_t children[ORDER];  // 内部节点:子页号
        int64_t values[MAX_KEYS];  // 叶子节点:值
    } u;
    PageHeader hdr;
    uint8_t _pad[PAGE_SIZE - sizeof(PageHeader) - sizeof(int64_t)*MAX_KEYS - sizeof(uint32_t)*ORDER];
} Page;

typedef struct {
    FILE *fp;
    uint32_t root_pg;
    uint32_t next_free_pg;
} BTree;

static void write_page(BTree *t, uint32_t pgno, Page *pg) {
    fseek(t->fp, (long)pgno * PAGE_SIZE, SEEK_SET);
    fwrite(pg, PAGE_SIZE, 1, t->fp);
}

static void read_page(BTree *t, uint32_t pgno, Page *pg) {
    fseek(t->fp, (long)pgno * PAGE_SIZE, SEEK_SET);
    fread(pg, PAGE_SIZE, 1, t->fp);
}

static uint32_t alloc_page(BTree *t) {
    return t->next_free_pg++;
}

static void init_leaf(Page *pg) {
    memset(pg, 0, PAGE_SIZE);
    pg->hdr.type = LEAF;
    pg->hdr.count = 0;
    pg->hdr.next_leaf = 0;
}

static void init_internal(Page *pg) {
    memset(pg, 0, PAGE_SIZE);
    pg->hdr.type = INTERNAL;
    pg->hdr.count = 0;
}

// 在叶子节点中查找键,返回所在槽位
static int leaf_find(Page *leaf, int64_t key) {
    for (int i = 0; i < leaf->hdr.count; i++) {
        if (leaf->keys[i] == key) return i;
        if (leaf->keys[i] > key) return -1;
    }
    return -1;
}

// 分裂叶子节点
static void split_leaf(BTree *t, uint32_t old_pgno, Page *old_leaf,
                       int64_t new_key, int64_t new_val,
                       uint32_t *new_pgno, int64_t *up_key) {
    Page new_leaf;
    init_leaf(&new_leaf);
    *new_pgno = alloc_page(t);
    int total = old_leaf->hdr.count + 1;
    int64_t all_keys[ORDER];
    int64_t all_vals[ORDER];
    int i = 0, j = 0;
    while (i < old_leaf->hdr.count && old_leaf->keys[i] < new_key) {
        all_keys[j] = old_leaf->keys[i];
        all_vals[j] = old_leaf->u.values[i];
        i++; j++;
    }
    all_keys[j] = new_key; all_vals[j] = new_val; j++;
    while (i < old_leaf->hdr.count) {
        all_keys[j] = old_leaf->keys[i];
        all_vals[j] = old_leaf->u.values[i];
        i++; j++;
    }
    int mid = ORDER / 2;
    old_leaf->hdr.count = mid;
    for (int k = 0; k < mid; k++) {
        old_leaf->keys[k] = all_keys[k];
        old_leaf->u.values[k] = all_vals[k];
    }
    new_leaf.hdr.count = total - mid;
    for (int k = 0; k < new_leaf.hdr.count; k++) {
        new_leaf.keys[k] = all_keys[mid + k];
        new_leaf.u.values[k] = all_vals[mid + k];
    }
    new_leaf.hdr.next_leaf = old_leaf->hdr.next_leaf;
    old_leaf->hdr.next_leaf = *new_pgno;
    *up_key = new_leaf.keys[0];
    write_page(t, old_pgno, old_leaf);
    write_page(t, *new_pgno, &new_leaf);
}

// 递归插入,返回是否需要向上分裂
static int insert_recursive(BTree *t, uint32_t pgno, int64_t key, int64_t val,
                            uint32_t *new_child_pgno, int64_t *up_key) {
    Page pg;
    read_page(t, pgno, &pg);
    if (pg.hdr.type == LEAF) {
        int pos = leaf_find(&pg, key);
        if (pos >= 0) {
            pg.u.values[pos] = val;
            write_page(t, pgno, &pg);
            return 0;
        }
        if (pg.hdr.count < MAX_KEYS) {
            int ins = pg.hdr.count;
            while (ins > 0 && pg.keys[ins-1] > key) {
                pg.keys[ins] = pg.keys[ins-1];
                pg.u.values[ins] = pg.u.values[ins-1];
                ins--;
            }
            pg.keys[ins] = key;
            pg.u.values[ins] = val;
            pg.hdr.count++;
            write_page(t, pgno, &pg);
            return 0;
        } else {
            split_leaf(t, pgno, &pg, key, val, new_child_pgno, up_key);
            return 1;
        }
    } else {
        // 内部节点:找到子节点
        int pos = pg.hdr.count - 1;
        while (pos >= 0 && pg.keys[pos] > key) pos--;
        uint32_t child_pg = pg.u.children[pos + 1];
        uint32_t child_new_pg;
        int64_t child_up_key;
        int need_split = insert_recursive(t, child_pg, key, val, &child_new_pg, &child_up_key);
        if (!need_split) return 0;
        // 子节点分裂,需要在本节点插入 up_key 和 child_new_pg
        if (pg.hdr.count < MAX_KEYS) {
            int ins = pg.hdr.count;
            while (ins > 0 && pg.keys[ins-1] > child_up_key) {
                pg.keys[ins] = pg.keys[ins-1];
                pg.u.children[ins+1] = pg.u.children[ins];
                ins--;
            }
            pg.keys[ins] = child_up_key;
            pg.u.children[ins+1] = child_new_pg;
            pg.hdr.count++;
            write_page(t, pgno, &pg);
            return 0;
        } else {
            // 内部节点分裂
            Page new_internal;
            init_internal(&new_internal);
            *new_child_pgno = alloc_page(t);
            int total = pg.hdr.count + 1;
            int64_t all_keys[ORDER];
            uint32_t all_children[ORDER+1];
            int i = 0, j = 0;
            while (i < pg.hdr.count && pg.keys[i] < child_up_key) {
                all_keys[j] = pg.keys[i];
                all_children[j] = pg.u.children[i];
                i++; j++;
            }
            all_children[j] = pg.u.children[i]; // 左子指针
            all_keys[j] = child_up_key;
            all_children[j+1] = child_new_pg;
            j++;
            i++;
            while (i < pg.hdr.count) {
                all_keys[j] = pg.keys[i];
                all_children[j+1] = pg.u.children[i+1];
                i++; j++;
            }
            int mid = ORDER / 2;
            *up_key = all_keys[mid];
            pg.hdr.count = mid;
            for (int k = 0; k < mid; k++) {
                pg.keys[k] = all_keys[k];
                pg.u.children[k] = all_children[k];
            }
            pg.u.children[mid] = all_children[mid];
            new_internal.hdr.count = total - mid - 1;
            for (int k = 0; k < new_internal.hdr.count; k++) {
                new_internal.keys[k] = all_keys[mid+1+k];
                new_internal.u.children[k] = all_children[mid+1+k];
            }
            new_internal.u.children[new_internal.hdr.count] = all_children[total];
            write_page(t, pgno, &pg);
            write_page(t, *new_child_pgno, &new_internal);
            return 1;
        }
    }
}

void bt_insert(BTree *t, int64_t key, int64_t val) {
    uint32_t new_pg;
    int64_t up_key;
    int split = insert_recursive(t, t->root_pg, key, val, &new_pg, &up_key);
    if (split) {
        Page new_root;
        init_internal(&new_root);
        uint32_t old_root = t->root_pg;
        uint32_t new_root_pg = alloc_page(t);
        new_root.hdr.count = 1;
        new_root.keys[0] = up_key;
        new_root.u.children[0] = old_root;
        new_root.u.children[1] = new_pg;
        write_page(t, new_root_pg, &new_root);
        t->root_pg = new_root_pg;
    }
}

int bt_get(BTree *t, int64_t key, int64_t *val) {
    uint32_t pgno = t->root_pg;
    Page pg;
    while (1) {
        read_page(t, pgno, &pg);
        if (pg.hdr.type == LEAF) {
            for (int i = 0; i < pg.hdr.count; i++) {
                if (pg.keys[i] == key) {
                    *val = pg.u.values[i];
                    return 1;
                }
            }
            return 0;
        } else {
            int pos = pg.hdr.count - 1;
            while (pos >= 0 && pg.keys[pos] > key) pos--;
            pgno = pg.u.children[pos + 1];
        }
    }
}

void bt_range_scan(BTree *t, int64_t start, int64_t end,
                   void (*callback)(int64_t, int64_t)) {
    uint32_t pgno = t->root_pg;
    Page pg;
    // 先找到起始叶子
    while (1) {
        read_page(t, pgno, &pg);
        if (pg.hdr.type == LEAF) break;
        int pos = pg.hdr.count - 1;
        while (pos >= 0 && pg.keys[pos] > start) pos--;
        pgno = pg.u.children[pos + 1];
    }
    while (pgno != 0) {
        read_page(t, pgno, &pg);
        for (int i = 0; i < pg.hdr.count; i++) {
            if (pg.keys[i] >= start && pg.keys[i] <= end)
                callback(pg.keys[i], pg.u.values[i]);
        }
        pgno = pg.hdr.next_leaf;
    }
}

void bt_init(BTree *t, const char *path) {
    t->fp = fopen(path, "w+b");
    t->next_free_pg = 1;
    Page root;
    init_leaf(&root);
    write_page(t, 0, &root);
    t->root_pg = 0;
}

void bt_close(BTree *t) {
    fclose(t->fp);
}

// 回调函数:打印扫描结果
void print_kv(int64_t k, int64_t v) {
    printf("%ld:%ld ", k, v);
}

int main() {
    BTree tree;
    bt_init(&tree, "test.db");
    // 插入 20 个键值对
    for (int64_t i = 1; i <= 20; i++) {
        bt_insert(&tree, i * 10, i * 100);
    }
    // 查找
    int64_t val;
    if (bt_get(&tree, 100, &val))
        printf("get(100)=%ld\n", val);
    else
        printf("get(100) not found\n");
    // 范围扫描 [50, 150]
    printf("range[50,150]: ");
    bt_range_scan(&tree, 50, 150, print_kv);
    printf("\n");
    bt_close(&tree);
    return 0;
}

编译与运行


gcc -o btree_disk btree_disk.c
./btree_disk

验证方法

程序执行后,观察输出是否包含以下内容:get(100)=1000,以及range[50,150]: 50:500 60:600 70:700 80:800 90:900 100:1000 110:1100 120:1200 130:1300 140:1400 150:1500。同时检查当前目录下生成test.db文件,其大小应为4KB的整数倍(至少包含文件头页与若干数据页)。若输出与上述一致,且文件大小合理,则判定存储引擎的插入、查找与范围扫描功能通过验证。