从零构建：用 Rust 实现一个简易键值存储引擎

从零构建：用 Rust 实现一个简易键值存储引擎

一、存储引擎的入门难题：从理论到可运行代码的鸿沟

数据库内核是后端开发中最"黑盒"的领域之一。MySQL 的 InnoDB、RocksDB 的 LSM-Tree，源码动辄数十万行，初学者很难从中学到核心原理。但一个最小化的键值存储引擎只需要三个组件：日志写入器（WAL）、内存索引（MemTable）、磁盘持久化（SSTable）。用 Rust 从零实现一个简易 KV Store，不仅能理解存储引擎的核心机制，还能实践 Rust 的文件 I/O、序列化和并发安全。

graph TB
    A[写入请求] --> B[写入 WAL 日志<br/>顺序写入, 保证持久性]
    B --> C[写入 MemTable<br/>内存中的排序跳表]
    C --> D{MemTable 是否满?}
    D -->|否| E[返回成功]
    D -->|是| F[冻结 MemTable → Immutable MemTable]
    F --> G[后台刷写为 SSTable<br/>有序磁盘文件]
    G --> H[创建新 MemTable<br/>继续接收写入]
    H --> E

    I[读取请求] --> J{MemTable 中查找}
    J -->|命中| K[返回值]
    J -->|未命中| L{Immutable MemTable}
    L -->|命中| K
    L -->|未命中| M[SSTable 查找<br/>从新到旧二分搜索]
    M -->|命中| K
    M -->|未命中| N[返回 Key 不存在]

二、LSM-Tree 存储引擎的核心机制

2.1 WAL 的写入与恢复

WAL（Write-Ahead Log）保证崩溃恢复：每次写入操作先追加到 WAL 文件，再更新内存索引。重启时重放 WAL 中的所有操作，重建 MemTable。WAL 是顺序追加写入，性能极高。

2.2 MemTable 的有序性保证

MemTable 使用跳表（SkipList）或 B-Tree 实现，保证键有序。有序性是后续刷写 SSTable 的前提——SSTable 要求键有序存储，这样查找时可以用二分搜索。

graph LR
    subgraph MemTable 跳表
        A[Head] --> B[key: a]
        B --> C[key: c]
        C --> D[key: f]
        D --> E[key: h]

        A -->|skip| D
        B -->|skip| E
    end

    subgraph SSTable 磁盘文件
        F[Data Block: a,c,f,h] --> G[Index Block: a→0, c→16, f→32, h→48]
        G --> H[Footer: index_offset]
    end

2.3 SSTable 的结构与查找

SSTable（Sorted String Table）是磁盘上的有序键值文件，由数据块、索引块和页脚组成。索引块记录每个数据块的起始键和偏移量，查找时先在索引块二分定位数据块，再在数据块内二分查找。

三、生产级代码实现与最佳实践

3.1 WAL 实现

use std::fs::{File, OpenOptions};
use std::io::{Write, BufWriter, BufRead, BufReader};
use std::path::Path;
use serde::{Serialize, Deserialize};

/// WAL 记录
#[derive(Serialize, Deserialize, Debug)]
pub enum WalEntry {
    Set { key: String, value: String },
    Delete { key: String },
}

/// WAL 写入器
pub struct WalWriter {
    writer: BufWriter<File>,
}

impl WalWriter {
    pub fn open(path: &Path) -> Result<Self, std::io::Error> {
        let file = OpenOptions::new()
            .create(true)
            .append(true)
            .open(path)?;
        Ok(Self {
            writer: BufWriter::new(file),
        })
    }

    /// 追加一条记录
    pub fn append(&mut self, entry: &WalEntry) -> Result<(), std::io::Error> {
        let json = serde_json::to_string(entry)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
        self.writer.write_all(json.as_bytes())?;
        self.writer.write_all(b"\n")?;
        self.writer.flush()?; // 每次写入后立即刷盘，保证持久性
        Ok(())
    }
}

/// WAL 恢复器
pub struct WalRecovery;

impl WalRecovery {
    /// 重放 WAL 重建内存状态
    pub fn replay(path: &Path) -> Result<Vec<WalEntry>, std::io::Error> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let mut entries = Vec::new();

        for line in reader.lines() {
            let line = line?;
            if line.trim().is_empty() {
                continue;
            }
            match serde_json::from_str::<WalEntry>(&line) {
                Ok(entry) => entries.push(entry),
                Err(_) => continue, // 跳过损坏的行（可能是崩溃时未写完）
            }
        }

        Ok(entries)
    }
}

3.2 MemTable 实现（基于 BTreeMap）

use std::collections::BTreeMap;

/// 内存表：键值有序存储
pub struct MemTable {
    data: BTreeMap<String, Option<String>>, // None 表示已删除
    approximate_size: usize,
}

impl MemTable {
    pub fn new() -> Self {
        Self {
            data: BTreeMap::new(),
            approximate_size: 0,
        }
    }

    /// 插入键值对
    pub fn set(&mut self, key: String, value: String) {
        self.approximate_size += key.len() + value.len();
        self.data.insert(key, Some(value));
    }

    /// 删除键（插入墓碑标记）
    pub fn delete(&mut self, key: String) {
        self.data.insert(key, None);
    }

    /// 查找键
    pub fn get(&self, key: &str) -> Option<Option<&String>> {
        self.data.get(key).map(|v| v.as_ref())
    }

    /// 获取近似大小（用于判断是否需要刷盘）
    pub fn approximate_size(&self) -> usize {
        self.approximate_size
    }

    /// 获取有序迭代器（用于刷写 SSTable）
    pub fn iter(&self) -> impl Iterator<Item = (&String, &Option<String>)> {
        self.data.iter()
    }

    /// 清空 MemTable
    pub fn clear(&mut self) {
        self.data.clear();
        self.approximate_size = 0;
    }
}

3.3 SSTable 实现与查找

use std::fs::{File, OpenOptions};
use std::io::{Write, Read, Seek, SeekFrom, BufWriter};
use std::path::{Path, PathBuf};

/// SSTable 索引条目
#[derive(Serialize, Deserialize)]
pub struct IndexEntry {
    pub key: String,
    pub offset: u64,
    pub length: u64,
}

/// SSTable 文件
pub struct SSTable {
    path: PathBuf,
    index: Vec<IndexEntry>,
}

impl SSTable {
    /// 从 MemTable 刷写新的 SSTable
    pub fn flush_from_memtable(
        memtable: &MemTable,
        dir: &Path,
        sequence: u64,
    ) -> Result<Self, Box<dyn std::error::Error>> {
        let filename = format!("sstable_{:06}.dat", sequence);
        let path = dir.join(&filename);
        let file = File::create(&path)?;
        let mut writer = BufWriter::new(file);

        let mut index = Vec::new();
        let mut offset: u64 = 0;

        for (key, value) in memtable.iter() {
            let entry_start = offset;

            // 写入数据行: [key_len:u16][key][value_len:u16][value_or_tombstone]
            let key_bytes = key.as_bytes();
            writer.write_all(&(key_bytes.len() as u16).to_le_bytes())?;
            writer.write_all(key_bytes)?;

            match value {
                Some(val) => {
                    let val_bytes = val.as_bytes();
                    writer.write_all(&(val_bytes.len() as u16).to_le_bytes())?;
                    writer.write_all(val_bytes)?;
                    offset += 2 + key_bytes.len() as u64 + 2 + val_bytes.len() as u64;
                }
                None => {
                    // 墓碑标记: value_len = 0xFFFF
                    writer.write_all(&0xFFFFu16.to_le_bytes())?;
                    offset += 2 + key_bytes.len() as u64 + 2;
                }
            }

            index.push(IndexEntry {
                key: key.clone(),
                offset: entry_start,
                length: offset - entry_start,
            });
        }

        writer.flush()?;

        Ok(SSTable { path, index })
    }

    /// 在 SSTable 中查找键
    pub fn get(&self, key: &str) -> Result<Option<Option<String>>, Box<dyn std::error::Error>> {
        // 在索引中二分查找
        let pos = self.index.binary_search_by(|entry| entry.key.as_str().cmp(key));

        match pos {
            Ok(idx) => {
                let entry = &self.index[idx];
                let mut file = File::open(&self.path)?;
                file.seek(SeekFrom::Start(entry.offset))?;

                let mut buf = vec![0u8; entry.length as usize];
                file.read_exact(&mut buf)?;

                // 解析数据行
                let key_len = u16::from_le_bytes([buf[0], buf[1]]) as usize;
                let _stored_key = &buf[2..2 + key_len];
                let val_len_offset = 2 + key_len;
                let val_len = u16::from_le_bytes([
                    buf[val_len_offset],
                    buf[val_len_offset + 1],
                ]);

                if val_len == 0xFFFF {
                    // 墓碑标记：键已删除
                    Ok(Some(None))
                } else {
                    let val_start = val_len_offset + 2;
                    let val_end = val_start + val_len as usize;
                    let value = String::from_utf8(buf[val_start..val_end].to_vec())?;
                    Ok(Some(Some(value)))
                }
            }
            Err(_) => Ok(None), // 键不在此 SSTable 中
        }
    }
}

3.4 KV Store 主结构

/// 简易键值存储引擎
pub struct KvStore {
    dir: PathBuf,
    memtable: MemTable,
    sstables: Vec<SSTable>,
    wal: WalWriter,
    flush_threshold: usize, // MemTable 刷盘阈值（字节）
}

impl KvStore {
    pub fn open(dir: &Path) -> Result<Self, Box<dyn std::error::Error>> {
        std::fs::create_dir_all(dir)?;

        // 1. 从 WAL 恢复 MemTable
        let wal_path = dir.join("wal.log");
        let mut memtable = MemTable::new();
        if wal_path.exists() {
            let entries = WalRecovery::replay(&wal_path)?;
            for entry in entries {
                match entry {
                    WalEntry::Set { key, value } => memtable.set(key, value),
                    WalEntry::Delete { key } => memtable.delete(key),
                }
            }
        }

        // 2. 加载已有 SSTable
        let mut sstables = Vec::new();
        // 实际实现需要扫描目录、加载索引

        let wal = WalWriter::open(&wal_path)?;

        Ok(Self {
            dir: dir.to_path_buf(),
            memtable,
            sstables,
            wal,
            flush_threshold: 4 * 1024 * 1024, // 4MB
        })
    }

    /// 设置键值对
    pub fn set(&mut self, key: String, value: String) -> Result<(), Box<dyn std::error::Error>> {
        // 1. 写入 WAL
        self.wal.append(&WalEntry::Set {
            key: key.clone(),
            value: value.clone(),
        })?;

        // 2. 写入 MemTable
        self.memtable.set(key, value);

        // 3. 检查是否需要刷盘
        if self.memtable.approximate_size() >= self.flush_threshold {
            self.flush_memtable()?;
        }

        Ok(())
    }

    /// 查找键
    pub fn get(&self, key: &str) -> Result<Option<String>, Box<dyn std::error::Error>> {
        // 1. 在 MemTable 中查找
        if let Some(value) = self.memtable.get(key) {
            return Ok(value.cloned());
        }

        // 2. 在 SSTable 中查找（从新到旧）
        for sstable in self.sstables.iter().rev() {
            if let Some(result) = sstable.get(key)? {
                return Ok(result);
            }
        }

        Ok(None)
    }

    /// 刷写 MemTable 到 SSTable
    fn flush_memtable(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        let sequence = self.sstables.len() as u64;
        let sstable = SSTable::flush_from_memtable(
            &self.memtable, &self.dir, sequence
        )?;

        self.sstables.push(sstable);
        self.memtable.clear();

        // 清空 WAL（新 MemTable 为空，旧 WAL 不再需要）
        let wal_path = self.dir.join("wal.log");
        std::fs::write(&wal_path, "")?;

        Ok(())
    }
}

四、简易 KV Store 的架构权衡

4.1 写入性能 vs 读取性能

4.2 Compaction 策略缺失的影响

本实现没有 Compaction（合并 SSTable），随着写入增多，SSTable 文件数量线性增长，读取需要遍历所有文件。生产级引擎（RocksDB）使用 Level Compaction 或 Size-Tiered Compaction 定期合并 SSTable，消除重复和墓碑记录。

4.3 适用边界与禁用场景

适用场景：

• 学习存储引擎核心原理
• 写多读少、数据量小的嵌入式场景
• 不需要事务和复杂查询的简单缓存

禁用场景：

• 生产环境（缺少 Compaction、并发控制、事务）
• 大数据量（无 Compaction 导致读取性能退化）
• 需要范围查询（本实现只支持点查）

五、总结

从零实现一个 KV Store 是理解存储引擎最直接的方式。核心只有三个组件：WAL 保证持久性、MemTable 提供快速写入、SSTable 实现磁盘持久化。LSM-Tree 的设计哲学是"写入优化"：所有写入都是顺序追加，避免随机写。代价是读取需要多级查找——MemTable → Immutable MemTable → SSTable（从新到旧）。这个简易实现缺少 Compaction、布隆过滤器和并发控制，但已经包含了存储引擎的核心骨架。理解了这个骨架，再读 RocksDB 的源码就有了地图。