用 Rust 从 0 到 1 实现最简化的 KV 存储引擎

Knight丶

2024-04-09 帮助1人

本文将从下层的数据编码到上层的 kv 数据读写接口实现完整介绍如何实现一个最简化的 kv 存储引擎，适合 Bitcask 存储模型和 Rust 语言的入门。本文的完整代码已开源在：GitHub - Morgan279/miniDB: A mini kv database demo that using simplified bitcask storage model with rust implementation.

功能和架构设计

学新通

如上图所示，本文的存储引擎支持 Get, Set, Remove 这三种最基础的命令。架构上可以分为 KV Store 和 Storage 两个部分，其中 KV Store 负责与客户端交互，接收外部命令和返回操作结果；Storage 负责持久化经过编码之后的数据。图里面的 Entry 是指数据的二进制编码，即每一份数据都会先封装成 Entry 然后编码成二进制之后再存入 Storage（实际的物理存储中，Entry 之间不会有缝隙）。

存储流程概述

本文的存储流程将采用经过简化后的 Bistcask 存储模型。Bitcask 是一种 LSM 类别下的简单的存储模型，其主要思想是通过顺序写和非原地更新（通过 append 新数据来更新或删除旧数据）节省磁盘随机寻道的开销，从而提高写吞吐量，但是这个特性也会造成读放大，适合写多读少的场景。

其基本的读写流程是：

从磁盘中读取数据文件在内存中生成索引（该步骤只在初始化时执行一次）
接收到读写命令之后通过在内存中的索引获取数据在磁盘中的位置
根据在第 2 步中拿到的具体位置信息在磁盘中进行数据读写，成功后更新内存中的索引

如下图所示：

学新通

由于数据是非原地更新，所以当接收到更新或删除命令时，只需要在文件末尾追加新的数据而不用管旧的数据，因为每次做数据读写都要先经过内存中的索引且索引指向的都是最新且有效的数据，所以该模型可以保证读写的正确性。

通过观察上图的存储模型，在实现层面上有两个问题需要解决：

存储的 key-value 数据的长度可以是任意的，这就导致每个 Entry 的长度不是固定的（如图中的浅蓝色 Entry 和深蓝色 Entry）且相邻 Entry 之间没有缝隙（即 Entry 之间无分隔符，视图上看起来都是无差别的二进制数据），而我们通过索引只能知道数据在磁盘中的起始位置，如何在 Entry 长度不确定的情况下正确解码出数据（数据编码与解码问题）。
资源是有限的，不可能无止境地追加新数据而不管旧数据，需要定期地对无效的旧数据做清理，回收对应的资源（无效数据回收问题）。

下面将分别对这两个问题的解决方案进行介绍。

数据编码与解码

数据正确解码的前提是明确知道 Entry 的在磁盘中的起始位置和终止位置，然后根据设计好的编码模型进行反序列解码即可。Entry 的起始位置我们可以从索引中很容易地获取到，重点在于如何知道 Entry 的终止位置。最简单的实现方案是在索引中额外存储 Entry 序列化之后的长度字段，起始位置加上长度即可得到终止位置，但是相对于磁盘，内存的存储资源是更为宝贵的，所以本文将基于磁盘通过 定长变长 的编码方案解决这个问题。

虽然我们并不能限制用户输入的 key、value 数据的长度，但是我们可以用固定长度的字节来描述输入的 key 和 value 的长度信息，比如我们可以用固定的 4 字节（即一个无符号整型占的空间）来描述长度在 [0,2^32] 的 key 的长度，已经可以满足常规的需求，而 value 也是同样的道理。

这里的 4 字节就是“定长”，key 或者 value 序列化之后的二进制数据就是 “变长”。我们把“定长”部分放在 Entry 的头部，后面跟“变长”部分，这样当我们在索引中拿到 Entry 的起始位置时，由于定长部分的长度是一个常量，所以我们可以直接相加得到 Entry 头部的终止位置，这样我们就可以解码出 Entry 头部的信息，即 key 和 value 的长度的信息，最后我们根据这两个长度信息计算得到整个 Entry 的终止位置，从而正确地解码出数据。

值得注意的是：数据解码的主要开销在于 I/O 操作，所以相对于直接在内存中存储 Entry 长度信息的方案，本方案在节约了内存空间的同时，并不会造成性能上的损耗。

本方案具体使用的对 Entry 的编码格式如下图所示：

学新通

这里编码中的 kind 字段是为了标识追加的新数据是 “新增、更新” 还是 “删除”，因为在服务重启时需要读取数据文件在内存中建立索引，假如只有新增和更新操作，因为我们只是从头开始顺序读取数据文件，而最新的数据总是是后面，即对于相同的 key 来说，我们顺序读取时，新数据的位置会不断覆盖索引中旧数据的位置，所以此时没有标识也不会影响正确性；

但是引入删除操作之后，如果被删除的这个 key 是唯一的，即这个 key 不会被新的有效的数据覆盖，如果不做特殊标识的话，那么就会导致用户读取到已经被删除的数据，从而造成正确性问题。所以，这里需要一个字段来标识操作类型，这样就可以在初始化索引时通过该标识跳过带有删除标识的 Entry（同时这个标识也会在无效数据回收过程中用到）。

无效数据回收

无效数据的回收在 Bitcask 模型中被称为 merge 操作或者 compact 操作，该操作将会清除磁盘中无效的数据，释放对应的资源。这里介绍最简单的一种实现方式：顺序扫描并解码磁盘中数据文件中的所有 Entry，判断这个 Entry 是不是有效的，然后把所有有效的 Entry 写入到一个新文件，然后用这个新文件替换掉原来的旧数据文件即可，这样就自然而然地把无效数据清理掉了。而判断 Entry 是否有效的方法也很简单，首先如果这个 Entry 带有删除标识，那么一定就是无效的，其次将这个 Entry 的起始位置与索引中的对应 key 的起始位置对比，如果相等就是有效的（因为索引中存储的总是最新数据的起始位置），否则就是无效的。

上面的问题都解决后，一个最简化的 KV 存储引擎就完成了。

下面将从编码实现层面进行进一步介绍。

Rust 实现

首先我们先对 Entry 进行定义：

pub struct Entry {
key_len: usize,
value_len: usize,
key: String,
value: String,
kind: CmdKind,
}
pub enum CmdKind {
PUT = 1,
DEL = 2,
}

这里的定义与上文图中描述的一致。

然后是 Entry 的编码与解码，首先我们定义出“定长部分”：

const USIZE_LEN: usize = std::mem::size_of::<usize>();
// key_len 和 val_len 都占 USIZE_LEN 的长度，操作类型标识字段（PUT、DEL）占 1 字节的长度
const ENTRY_HEAD_LEN: usize = USIZE_LEN * 2 1;

“定长”部分由 4 字节或 8 字节（usize 在 32 位机器上为 4 字节， 64位机器上为 8 字节）的 key 长度字段，value 长度字段，以及 1 字节的标识字段组成，这样我们可以通过计算得到整个 Entry 的长度信息：

pub fn size(&self) -> usize {
ENTRY_HEAD_LEN self.key_len self.value_len
}

在编码时，我们对各个字段逐一进行序列化即可：

pub fn encode(&self) -> Vec<u8> {
let mut buf = vec![0; self.size()];
// encode key len
buf[0..USIZE_LEN].copy_from_slice(&self.key_len.to_be_bytes());
// encode value length
buf[USIZE_LEN..USIZE_LEN * 2].copy_from_slice(&self.value_len.to_be_bytes());
// encode kind
buf[USIZE_LEN * 2..ENTRY_HEAD_LEN]
.copy_from_slice(bincode::serialize(&self.kind).unwrap().as_slice());
// encode key
buf[ENTRY_HEAD_LEN..ENTRY_HEAD_LEN self.key_len].copy_from_slice(self.key.as_bytes());
// encode value
buf[ENTRY_HEAD_LEN self.key_len..].copy_from_slice(self.value.as_bytes());
buf
}

在解码时我们先解码出 Entry 的头部信息：

pub fn decode(b: &[u8; ENTRY_HEAD_LEN]) -> Result<Entry> {
let key_len = usize::from_be_bytes(b[0..USIZE_LEN].try_into()?);
let value_len = usize::from_be_bytes(b[USIZE_LEN..USIZE_LEN * 2].try_into()?);
let kind: CmdKind = bincode::deserialize(&b[USIZE_LEN * 2..ENTRY_HEAD_LEN])?;
Ok(Entry {
key_len,
value_len,
kind,
key: String::new(),
value: String::new(),
})
}

然后解码出整个 Entry 的信息：

let mut e = Entry::decode(&buf)?;
let mut key_buf = vec![0; e.key_len];
self.reader.read_exact(key_buf.as_mut_slice())?;
e.key = String::from_utf8(key_buf)?;
let mut val_buf = vec![0; e.value_len];
self.reader.read_exact(val_buf.as_mut_slice())?;
e.value = String::from_utf8(val_buf)?;

然后是无效数据回收的实现，首先筛选出有效的 Entry：

let mut offset = 0;
let mut valid_entry = Vec::new();
loop {
match self.read_at(offset) {
Ok(e) => {
let size = e.size() as u64;
if let Some(valid_pos) = self.index.get(&e.key) {
// 滤出带有删除标记的数据且只保留最新数据
if e.kind == CmdKind::PUT && *valid_pos == offset {
valid_entry.push(e);
}
}
offset = size;
}
Err(KvsError::EOF) => {
// 全部解码完毕
break;
}
Err(e) => {
return Err(e);
}
}
}

然后将有效的 Entry 写入新数据文件然后直接替换掉旧数据文件：

if !valid_entry.is_empty() {
let mut data_path_ancestors = self.data_path_buf.ancestors();
data_path_ancestors.next();
let merge_path_buf = data_path_ancestors
.next()
.ok_or(KvsError::InvalidDataPath)?
.join(STORAGE_FILE_PREFIX.to_string() ".merge");
let merge_file = File::create(merge_path_buf.as_path())?;
// 在数据文件的目录下新建回收文件
let mut write_buf = BufWriterWithPos::new(merge_file)?;
// 向回收文件中写入所有有效的数据
for e in &valid_entry {
let key = e.key.clone();
self.index.insert(key, write_buf.pos);
write_buf.write(&e.encode())?;
}
// 用回收文件替换原来的数据文件
self.writer = write_buf;
self.reader = BufReaderWithPos::new(File::open(merge_path_buf.as_path())?)?;
std::fs::remove_file(self.data_path_buf.as_path())?;
std::fs::rename(merge_path_buf.as_path(), self.data_path_buf.as_path())?;
}

值得注意的是，本文基于“最简化”的考虑，无效数据回收的触发方式是执行 set 操作之后如果无效的数据大小超过一定的阈值，则同步阻塞地进行回收操作：

fn put(&mut self, key: String, val: String) -> Result<()> {
let e = Entry::new(key, val, CmdKind::PUT);
self.write(e)?;
// 无效的数据超过一定的阈值时则执行回收操作
if self.pending_compact >= COMPACTION_THRESHOLD {
self.merge()?;
}
Ok(())
}

这里是一个非常大的可以进行性能优化的地方，比如创建另一个线程非阻塞地定期执行无效数据回收操作，但是需要额外处理多线程操作同一个数据源而带来的“数据竞争”问题。此外还可以将整个单线程模型优化为多线程模型，为了减小临界区的竞争还可以将数据存储为多个小文件，从而减小锁的粒度等等很多性能优化手段，这篇博文旨在于从功能性方面介绍如何设计和实现一种最简单的 KV 存储模型，性能优化方面不做过多讨论。

以上关键逻辑实现之后，可以利用 trait 抽象出 Storage 的共有行为然后组合上面的核心逻辑对其进行实现即可：

pub trait Storage {
fn get(&mut self, key: String) -> Result<Option<String>>;
fn put(&mut self, key: String, val: String) -> Result<()>;
fn remove(&mut self, key: String) -> Result<()>;
}
impl Storage for SimplifiedBitcask {
fn get(&mut self, key: String) -> Result<Option<String>> {
match self.read(&key) {
Ok(e) => Ok(Some(e.value)),
Err(KvsError::KeyNotFound) => Ok(None),
Err(e) => Err(e),
}
}
fn put(&mut self, key: String, val: String) -> Result<()> {
let e = Entry::new(key, val, CmdKind::PUT);
self.write(e)?;
if self.pending_compact >= COMPACTION_THRESHOLD {
self.merge()?;
}
Ok(())
}
fn remove(&mut self, key: String) -> Result<()> {
if self.index.contains_key(&key) {
let e = Entry::new(key.clone(), String::new(), CmdKind::DEL);
self.write(e)?;
self.index.remove(&key);
return Ok(());
}
Err(KvsError::KeyNotFound)
}
}

最后再简单地实现 KV Store 作为对外交互的接口：

pub struct KvStore {
storage: Box<dyn Storage>,
}
impl KvStore {
pub fn open(path: &Path) -> Result<KvStore> {
let storage = SimplifiedBitcask::open(path.to_path_buf())?;
Ok(KvStore {
storage: Box::new(storage),
})
}
pub fn get(&mut self, key: String) -> Result<Option<String>> {
self.storage.get(key)
}
pub fn set(&mut self, key: String, val: String) -> Result<()> {
self.storage.put(key, val)
}
pub fn remove(&mut self, key: String) -> Result<()> {
self.storage.remove(key)
}
}

本文在理论层面，介绍的是最简单的一种存储模型；在实现层面，关键在于 I/O 操作方面的编码实现，其次在于一些异常情况的处理。本文的完整代码已开源在 GitHub - Morgan279/miniDB: A mini kv database demo that using simplified bitcask storage model with rust implementation.

欢迎大家提出修改建议或者在此基础上进行进一步完善。

这篇好文章是转载于：学新通技术网

用 Rust 从 0 到 1 实现最简化的 KV 存储引擎

功能和架构设计

存储流程概述

数据编码与解码

无效数据回收

Rust 实现

photoshop保存的图片太大微信发不了怎么办

Android 11 保存文件到外部存储，并分享文件

word里面弄一个表格后上面的标题会跑到下面怎么办

《学习通》视频自动暂停处理方法

photoshop扩展功能面板显示灰色怎么办

微信公众号没有声音提示怎么办

怎样阻止微信小程序自动打开

excel下划线不显示怎么办

excel打印预览压线压字怎么办

photoshop蒙版画笔没反应怎么办