walrus_rust/wal/runtime/
index.rs

1use crate::wal::paths::WalPathManager;
2use rkyv::{Archive, Deserialize, Serialize};
3use std::collections::HashMap;
4use std::fs;
5
6#[derive(Archive, Deserialize, Serialize, Debug, Clone)]
7pub struct BlockPos {
8    pub cur_block_idx: u64,
9    pub cur_block_offset: u64,
10}
11
12pub struct WalIndex {
13    store: HashMap<String, BlockPos>,
14    path: String,
15}
16
17impl WalIndex {
18    pub fn new(file_name: &str) -> std::io::Result<Self> {
19        let paths = WalPathManager::default();
20        Self::new_in(&paths, file_name)
21    }
22
23    pub(super) fn new_in(paths: &WalPathManager, file_name: &str) -> std::io::Result<Self> {
24        paths.ensure_root()?;
25        let path = paths.index_path(file_name);
26        let store = path
27            .exists()
28            .then(|| fs::read(&path).ok())
29            .flatten()
30            .and_then(|bytes| {
31                if bytes.is_empty() {
32                    return None;
33                }
34                // SAFETY: `bytes` comes from our persisted index file which we control;
35                // we only proceed when the file is non-empty and rkyv can interpret it.
36                let archived = unsafe { rkyv::archived_root::<HashMap<String, BlockPos>>(&bytes) };
37                archived.deserialize(&mut rkyv::Infallible).ok()
38            })
39            .unwrap_or_default();
40
41        Ok(Self {
42            store,
43            path: path.to_string_lossy().into_owned(),
44        })
45    }
46
47    pub fn set(&mut self, key: String, idx: u64, offset: u64) -> std::io::Result<()> {
48        self.store.insert(
49            key,
50            BlockPos {
51                cur_block_idx: idx,
52                cur_block_offset: offset,
53            },
54        );
55        self.persist()
56    }
57
58    pub fn get(&self, key: &str) -> Option<&BlockPos> {
59        self.store.get(key)
60    }
61
62    pub fn remove(&mut self, key: &str) -> std::io::Result<Option<BlockPos>> {
63        let result = self.store.remove(key);
64        if result.is_some() {
65            self.persist()?;
66        }
67        Ok(result)
68    }
69
70    fn persist(&self) -> std::io::Result<()> {
71        let tmp_path = format!("{}.tmp", self.path);
72        let bytes = rkyv::to_bytes::<_, 256>(&self.store).map_err(|e| {
73            std::io::Error::new(
74                std::io::ErrorKind::Other,
75                format!("index serialize failed: {:?}", e),
76            )
77        })?;
78
79        fs::write(&tmp_path, &bytes)?;
80        fs::File::open(&tmp_path)?.sync_all()?;
81        fs::rename(&tmp_path, &self.path)?;
82        Ok(())
83    }
84}
walrus_rust/wal/runtime/index.rs

walrus_rust/wal/runtime/
index.rs
