0290-KVS-读取多个文件

环境

• Time 2022-12-16
• 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"

log.rs

use std::collections::BTreeMap;
use std::path::Path;

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

use self::{reader::PosBufReader, writer::PosBufWriter};

mod reader;
mod writer;

pub struct CommandLog {
    reader: PosBufReader,
    writer: PosBufWriter,
    index: Index,
}

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

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

    pub fn get(&mut self, key: &str) -> KvResult {
        self.reader.read(self.index.get(key))
    }

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

    pub fn remove(&mut self, key: String) -> KvResult {
        let result = self.reader.read(self.index.get(&key));
        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);
        result
    }
}

reader.rs

use std::collections::HashMap;
use std::ffi::OsStr;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};

use anyhow::Result;
use serde_json::Deserializer;

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

use super::KvResult;

pub struct PosBufReader {
    readers: HashMap<usize, BufReader<File>>,
    pub gen: usize,
}

impl PosBufReader {
    pub fn new(path: &Path) -> Result<PosBufReader> {
        let list = sorted_gen_list(path)?;
        let mut reader = Self {
            readers: HashMap::default(),
            gen: list.last().unwrap_or(&0) + 1,
        };
        for gen in list {
            reader.add_reader(path, gen)?;
        }
        File::create(path.join(format!("{}.log", reader.gen)))?;
        reader.add_reader(path, reader.gen)?;
        Ok(reader)
    }

    fn add_reader(&mut self, path: &Path, gen: usize) -> Result<()> {
        let path = path.join(format!("{}.log", gen));
        let file = OpenOptions::new().read(true).open(&path)?;
        self.readers.insert(gen, BufReader::new(file));
        Ok(())
    }

    pub fn load(&mut self, map: &mut Index) -> Result<()> {
        for (gen, reader) in &mut self.readers {
            let mut old = reader.seek(SeekFrom::Start(0))? as usize;
            let mut stream = Deserializer::from_reader(reader).into_iter();
            while let Some(cmd) = stream.next() {
                let new = stream.byte_offset();
                match cmd? {
                    Command::Set { key, .. } => {
                        let value = CommandPosition::new(*gen, old, new - old);
                        map.insert(key, value);
                    }
                    Command::Remove { key } => {
                        map.remove(&key);
                    }
                };
                old = new;
            }
        }
        Ok(())
    }

    pub fn read(&mut self, position: Option<&CommandPosition>) -> KvResult {
        let position = match position {
            Some(position) => position,
            None => return Ok(None),
        };

        let reader = &mut self.readers.get_mut(&position.gen).unwrap();
        reader.seek(SeekFrom::Start(position.pos as u64))?;
        let command = (reader).take(position.len as u64);

        match serde_json::from_reader(command)? {
            Command::Set { value, .. } => Ok(Some(value)),
            _ => unreachable!(),
        }
    }
}

fn sorted_gen_list(path: &Path) -> Result<Vec<usize>> {
    let mut list: Vec<usize> = path
        .read_dir()?
        .flat_map(|res| res.map(|entry| entry.path()))
        .filter(|path| path.is_file())
        .flat_map(file_name_to_number)
        .collect();
    list.sort_unstable();
    Ok(list)
}

fn file_name_to_number(path: PathBuf) -> Option<usize> {
    let name = path.file_name().and_then(OsStr::to_str)?;

    match name.ends_with(".log") {
        false => None,
        true => name.trim_end_matches(".log").parse().ok(),
    }
}

writer.rs

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

use crate::cmd::CommandPosition;

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

impl PosBufWriter {
    pub fn new(path: &Path, gen: usize) -> anyhow::Result<Self> {
        let path = path.join("1.log");
        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, gen })
    }

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

总结

根据提供的目录，到目录中读取所有的文件，然后将偏移量加入到索引中。

附录