0289-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 writer = PosBufWriter::new(path)?;
        let mut index = BTreeMap::default();
        let mut reader = PosBufReader::new(path)?;
        reader.load(&mut index)?;
        Ok(Self {
            reader,
            writer,
            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::ffi::OsStr;
use std::fs::{File, OpenOptions};
use std::io::{BufReader, Read, Seek, SeekFrom};
use std::path::{Path, PathBuf};

use serde_json::Deserializer;

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

use super::KvResult;

pub struct PosBufReader {
    reader: BufReader<File>,
}

impl PosBufReader {
    pub fn new(path: &Path) -> anyhow::Result<PosBufReader> {
        let number = *sorted_gen_list(path)?.first().unwrap();
        let path = path.join(format!("{}.log", number));
        let file = OpenOptions::new().read(true).open(&path)?;
        let reader = BufReader::new(file);
        Ok(Self { reader })
    }

    pub fn load(&mut self, map: &mut Index) -> anyhow::Result<()> {
        let reader = &mut self.reader;
        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, .. } => {
                    map.insert(key, CommandPosition::new(old, new - old));
                }
                Command::Remove { .. } => (),
            };
            old = new;
        }
        Ok(())
    }

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

        self.reader.seek(SeekFrom::Start(position.pos as u64))?;
        let command = (&mut self.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) -> anyhow::Result<Vec<u64>> {
    let mut list: Vec<u64> = 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<u64> {
    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,
}

impl PosBufWriter {
    pub fn new(path: &Path) -> 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 })
    }

    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 += result.len;
        Ok(result)
    }
}

总结

根据提供的目录，到目录中读取文件，然后将文件中的内容加入到索引中。

附录