用Rust实现MMKV-开篇

前言 Link to heading

在Android上，基于mmap的key-value存储库MMKV因为其高效安全非常受欢迎。由于最近在学习Rust，我决定尝试用Rust来实现一遍MMKV。

对应的仓库在github：https://github.com/yangkx1024/MMKV

使用mmap来映射文件存储，使用protobuf来序列化数据，内存中使用map来保存kev-value对（很直观，就是内存缓存），每次添加新key，都是往map中加数据，然后往文件存储的末尾添加一条记录做持久化（文件中有重复的key也无所谓，重新初始化时顺序读取文件，解码数据依次添加进map，末尾的新记录会覆盖旧记录，值依然是正确的）。当写持久化记录时如果遇到空间不足，则考虑去重：丢弃掉全部的文件内容，把内存map中的值重新写入。如果此后空间依然不够，才考虑扩容。

在Rust中mmap可以使用libc的Rust binding，protobuf也有现成的crate：https://crates.io/crates/protobuf。

那么万事具备，直接开始做。

先设计protobuf数据结构：

message KV {
  string key = 1;
  bytes value = 2;
}

为了简单，先只保留key和数据的字节序列，不关心类型。

proto类型要先经过protoc编译才能生成Rust可以使用的代码，为此需要自定义编译过程。

需要在build.rs之中引入protobuf相关的依赖，分别是protobuf_codegen和protoc_bin_vendored。

前者提供Rust语言API来操作proto文件，后者会自动识别当前平台下载对应的proto编译器。

然后就可以添加build.rs来实现预编译：

fn main() {
    println!("process proto files");
    protobuf_codegen::Codegen::new()
        // Use `protoc` parser, optional.
        .protoc()
        .protoc_path(&protoc_bin_vendored::protoc_bin_path().unwrap())
        // All inputs and imports from the inputs must reside in `includes` directories.
        .includes(&["src/protos"])
        // Inputs must reside in some of include paths.
        .input("src/protos/kv.proto")
        // Specify output directory relative to Cargo output directory.
        .cargo_out_dir("protos")
        .run_from_script();
}

在此之后，编译器会在target目录下生成proto翻译之后的Rust代码，之后就可以使用include!宏来在代码中引入对应的mod：

include!(concat!(env!("OUT_DIR"), "/protos/mod.rs"));

同时也能看到proto生成的Rust代码：

#[derive(PartialEq,Clone,Default,Debug)]
// @@protoc_insertion_point(message:KV)
pub struct KV {
    // message fields
    // @@protoc_insertion_point(field:KV.key)
    pub key: ::std::string::String,
    // @@protoc_insertion_point(field:KV.value)
    pub value: ::std::vec::Vec<u8>,
    // special fields
    // @@protoc_insertion_point(special_field:KV.special_fields)
    pub special_fields: ::protobuf::SpecialFields,
}

然后着手设计数据编码和存储结构。

参考原版的MMKV的思路，每次存储Key都是往文件末尾添加新的记录，当文件空间不足时才考虑Key去重，这样做的好处是不用seek查找，非常高效。所以这里需要设计一种结构能够顺序添加，同时保证重新加载数据时能准确顺序读取。最终采取的结构如下：

4 bytes	data
存储数据长度	存储完整的字节序列

对于每个key-value对，先将其编码成proto结构，也就是上述的KV，然后将其转换成字节序列，先在头4个字节之中写入KV字节序列的长度，然后再写入完整的KV字节序列。这样在重新加载数据时，我们就能顺序的读取每一个key-value的字节序列然后将其解码到proto结构KV中。

然后考虑到这么一个问题，假如先写入一个String类型的数据，然后尝试以int32类型去读取，由于存储的都是字节序列，当将一个String类型的字节序列转换为int32时，是会出错的。所以需要加入类型验证。当尝试解码数据时，先验证类型，如果类型不匹配，则表示当前类型对应的key没有值。所以最后protobuf数据结构如下：

enum Types {
  I32 = 0;
  STR = 1;
  BYTE = 2;
}

message KV {
  string key = 1;
  Types type = 2;
  bytes value = 3;
}

添加了一个枚举来标识支持的类型。

然后设计一个结构体来承载上述的设计思路，这个结构需要有序列化输出bytes的能力，需要能够按照类型解码数据，对应的代码都在buffer这个mod之中：

pub struct Buffer(KV);

impl Buffer {
    // 将i32转换为字节序列后构造对应的KV
    pub fn from_i32(key: &str, value: i32) -> Self {
        Buffer::from_kv(key, Types::I32, value.to_be_bytes().as_ref())
    }
  
    ......

    // 根据数据类型构造KV
    fn from_kv(key: &str, t: Types, value: &[u8]) -> Self {
        let mut kv = KV::new();
        kv.key = key.to_string();
        kv.type_ = EnumOrUnknown::new(t);
        kv.value = value.to_vec();
        Buffer(kv)
    }
  
    // 将数据序列化输出为bytes以便存储
    pub fn to_bytes(&self) -> Vec<u8> {
        // 先用protobuf的API将raw_data(即proto数据KV)序列化为bytes
        let bytes_to_write = self.raw_data.write_to_bytes().unwrap();
        // 拿到protobuf数据的长度
        let len = bytes_to_write.len() as u32;
        // 先将长度写入头4个字节(这里u32类型转换为bytes默认占4个字节)，再将protobuf接在后面
        let mut data = len.to_be_bytes().to_vec();
        data.extend_from_slice(bytes_to_write.as_slice());
        data
    }
  
    // 从序列化之后的bytes数据中解码
    pub fn from_encoded_bytes(data: &[u8]) -> (Self, u32) {
        // 先从前4个字节读取长度
        let item_len = u32::from_be_bytes(
            data[0..4].try_into().unwrap()
        );
        // 然后使用protobuf的API从剩下的数据之中解码KV，同时返回已经读取的数据的长度
        (Buffer(
            KV::parse_from_bytes(&data[4..(4 + item_len as usize)]).unwrap()
        ), 4 + item_len)
    }
  
    // 从原始数据中解码str
    fn decode_str(&self) -> Option<&str> {
        match self.raw_data.type_.enum_value() {
            // 类型为STR则将原始数据转码为str
            Ok(Types::STR) => {
                str::from_utf8(self.raw_data.value.as_slice()).ok()
            }
            // 否则认为数据不存在
            _ => None
        }
    }
    ......
}

以上就是完整的数据结构设计以及序列化和反序列化的关键代码。

接下来我们设计一些test case来验证想法：

#[test]
fn test_buffer() {
    // 构造一个str类型的buffer
    let buffer = Buffer::from_str("first_key", "first_value");
    let bytes = buffer.to_bytes();
    let data_len = bytes.len();
    // 从字节数据解码一份copy
    let (copy, len) = Buffer::from_encoded_bytes(bytes.as_slice());
    // 验证数据正确性
    assert_eq!(copy, buffer);
    assert_eq!(len , data_len as u32);
    assert_eq!(copy.decode_str().unwrap(), "first_value");
    assert_eq!(copy.decode_i32(), None);
    assert_eq!(copy.decode_bool(), None);
    ......
}

当所有的test case通过，那么buffer这一个mod就完成了，至此已经有了基本的数据编解码能力。

下一篇讲memory map的使用。

前言 Link to heading

设计 Link to heading