0287-KVS-记录写入时的索引

环境

• Time 2022-12-13
• WSL-Ubuntu 22.04
• Rust 1.65.0

前言

说明

参考：https://github.com/pingcap/talent-plan

目标

在上一节的基础上，在写入时，记录写入文件时的索引信息。

Cargo.toml

[package]
edition = "2021"
name = "kvs"
version = "1.0.0"

[dependencies]
clap = {version = "4", features = ["derive"]}
serde = {version = "1", features = ["derive"]}
serde_json = "1"

cmd.rs

#[derive(serde::Serialize, serde::Deserialize)]
pub enum Command {
    Set { key: String, value: String },
    Remove { key: String },
}

pub struct CommandPosition {
    position: usize,
    length: usize,
}

impl CommandPosition {
    pub fn new(position: usize, length: usize) -> Self {
        CommandPosition { position, length }
    }
}

log.rs

use std::collections::{BTreeMap, HashMap};

use crate::cmd::{Command, CommandPosition};

use self::writer::PosBufWriter;

mod reader;
mod writer;

pub(crate) struct CommandLog {
    writer: PosBufWriter,
    map: HashMap<String, String>,
    index: BTreeMap<String, CommandPosition>,
}

pub type KvResult = anyhow::Result<Option<String>>;

impl CommandLog {
    pub(crate) fn new() -> anyhow::Result<Self> {
        let path = "/root/log/1.log";
        let writer = PosBufWriter::new(path)?;
        let mut map = HashMap::default();
        reader::load(path, &mut map)?;
        Ok(Self {
            writer,
            map,
            index: BTreeMap::default(),
        })
    }

    pub(crate) fn get(&mut self, key: &str) -> KvResult {
        Ok(self.map.get(key).map(String::from))
    }

    pub fn set(&mut self, key: String, value: String) -> KvResult {
        let command = Command::Set {
            key: key.to_string(),
            value: value.to_string(),
        };
        let json = serde_json::to_string(&command)?;
        let position = self.writer.write(json.as_bytes())?;
        self.index.insert(key.to_string(), position);
        Ok(self.map.insert(key, value))
    }

    pub(crate) fn remove(&mut self, key: String) -> KvResult {
        let command = Command::Remove {
            key: key.to_string(),
        };
        let json = serde_json::to_string(&command)?;
        self.writer.write(json.as_bytes())?;
        self.index.remove(&key);
        Ok(self.map.remove(&key))
    }
}

writer.rs

use std::fs::{File, OpenOptions};
use std::io::{BufWriter, Seek, SeekFrom, Write};

use crate::cmd::CommandPosition;

pub struct PosBufWriter {
    writer: BufWriter<File>,
    pos: usize,
}

impl PosBufWriter {
    pub fn new(path: &str) -> anyhow::Result<Self> {
        let file = OpenOptions::new().append(true).create(true).open(path)?;
        let mut writer = BufWriter::new(file);
        let pos = writer.seek(SeekFrom::End(0))? as usize;
        Ok(Self { writer, pos })
    }

    pub fn write(&mut self, buf: &[u8]) -> anyhow::Result<CommandPosition> {
        let result = CommandPosition::new(self.pos, buf.len());
        self.writer.write_all(buf)?;
        self.writer.flush()?;
        self.pos += buf.len();
        Ok(result)
    }
}

总结

在写入键值对到文件时，将写入时的索引记录下来。

附录