Skip to main content

weavatrix_memory/snapshot/
file.rs

1use super::SnapshotStore;
2use crate::{Codec, Durability, MemoryError, ProjectionSnapshot, Result, store::frame::crc32c};
3use std::{
4    fs::{self, File, OpenOptions},
5    io::{Read, Write},
6    marker::PhantomData,
7    path::{Path, PathBuf},
8    sync::atomic::{AtomicU64, Ordering},
9};
10
11const SNAPSHOT_HEADER: &[u8; 8] = b"WMEMSN01";
12const FRAME_HEADER_LEN: usize = 20;
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub struct SnapshotOptions {
16    pub durability: Durability,
17    pub max_snapshot_bytes: usize,
18}
19
20impl Default for SnapshotOptions {
21    fn default() -> Self {
22        Self {
23            durability: Durability::SyncData,
24            max_snapshot_bytes: 512 * 1024 * 1024,
25        }
26    }
27}
28
29pub struct FileSnapshotStore<P, C> {
30    directory: PathBuf,
31    prefix: String,
32    codec: C,
33    options: SnapshotOptions,
34    marker: PhantomData<fn() -> P>,
35}
36
37impl<P, C> FileSnapshotStore<P, C>
38where
39    C: Codec<ProjectionSnapshot<P>>,
40{
41    /// Creates an immutable, generation-named snapshot store.
42    ///
43    /// # Errors
44    ///
45    /// Rejects an invalid prefix or inaccessible directory.
46    pub fn open(
47        directory: impl AsRef<Path>,
48        prefix: impl Into<String>,
49        codec: C,
50        options: SnapshotOptions,
51    ) -> Result<Self> {
52        let prefix = prefix.into();
53        if prefix.is_empty() || prefix.trim() != prefix || prefix.contains(['/', '\\']) {
54            return Err(MemoryError::InvalidValue {
55                field: "snapshot.prefix",
56                reason: "must be a simple non-empty file prefix",
57            });
58        }
59        if options.max_snapshot_bytes == 0 {
60            return Err(MemoryError::InvalidValue {
61                field: "max_snapshot_bytes",
62                reason: "must be greater than zero",
63            });
64        }
65        let directory = directory.as_ref().to_path_buf();
66        fs::create_dir_all(&directory).map_err(|error| io("create snapshot directory", error))?;
67        Ok(Self {
68            directory,
69            prefix,
70            codec,
71            options,
72            marker: PhantomData,
73        })
74    }
75
76    fn final_path(&self, position: u64) -> PathBuf {
77        self.directory
78            .join(format!("{}-{position:020}.wmsnap", self.prefix))
79    }
80
81    fn latest_path(&self) -> Result<Option<(u64, PathBuf)>> {
82        let start = format!("{}-", self.prefix);
83        let mut latest = None;
84        for entry in
85            fs::read_dir(&self.directory).map_err(|error| io("read snapshot directory", error))?
86        {
87            let entry = entry.map_err(|error| io("read snapshot entry", error))?;
88            let name = entry.file_name();
89            let Some(name) = name.to_str() else {
90                continue;
91            };
92            let Some(raw) = name
93                .strip_prefix(&start)
94                .and_then(|value| value.strip_suffix(".wmsnap"))
95            else {
96                continue;
97            };
98            let Ok(position) = raw.parse::<u64>() else {
99                continue;
100            };
101            if latest
102                .as_ref()
103                .is_none_or(|(current, _)| position > *current)
104            {
105                latest = Some((position, entry.path()));
106            }
107        }
108        Ok(latest)
109    }
110
111    fn read_path(&self, path: &Path, expected_position: u64) -> Result<ProjectionSnapshot<P>> {
112        let mut file = File::open(path).map_err(|error| io("open snapshot", error))?;
113        let mut header = [0_u8; FRAME_HEADER_LEN];
114        file.read_exact(&mut header)
115            .map_err(|error| io("read snapshot header", error))?;
116        if &header[..8] != SNAPSHOT_HEADER {
117            return Err(corrupt("unsupported snapshot header"));
118        }
119        let length = usize::try_from(u64::from_le_bytes(header[8..16].try_into().unwrap()))
120            .map_err(|_| corrupt("snapshot length exceeds platform capacity"))?;
121        if length > self.options.max_snapshot_bytes {
122            return Err(corrupt("snapshot exceeds configured size limit"));
123        }
124        let expected_crc = u32::from_le_bytes(header[16..20].try_into().unwrap());
125        let mut bytes = vec![0; length];
126        file.read_exact(&mut bytes)
127            .map_err(|error| io("read snapshot payload", error))?;
128        let actual_len = file
129            .metadata()
130            .map_err(|error| io("read snapshot metadata", error))?
131            .len();
132        let framed_len =
133            u64::try_from(FRAME_HEADER_LEN + length).map_err(|_| MemoryError::CapacityOverflow)?;
134        if actual_len != framed_len || crc32c(&bytes) != expected_crc {
135            return Err(corrupt("snapshot length or checksum mismatch"));
136        }
137        let snapshot = self.codec.decode(&bytes)?;
138        if snapshot.cursor.global_position != Some(expected_position) {
139            return Err(corrupt("snapshot filename and cursor disagree"));
140        }
141        Ok(snapshot)
142    }
143
144    fn encode_frame(&self, snapshot: &ProjectionSnapshot<P>) -> Result<Vec<u8>> {
145        let bytes = self.codec.encode(snapshot)?;
146        if bytes.len() > self.options.max_snapshot_bytes {
147            return Err(MemoryError::InvalidValue {
148                field: "snapshot",
149                reason: "encoded snapshot exceeds max_snapshot_bytes",
150            });
151        }
152        let length = u64::try_from(bytes.len()).map_err(|_| MemoryError::CapacityOverflow)?;
153        let mut frame = Vec::with_capacity(FRAME_HEADER_LEN + bytes.len());
154        frame.extend_from_slice(SNAPSHOT_HEADER);
155        frame.extend_from_slice(&length.to_le_bytes());
156        frame.extend_from_slice(&crc32c(&bytes).to_le_bytes());
157        frame.extend_from_slice(&bytes);
158        Ok(frame)
159    }
160}
161
162impl<P, C> SnapshotStore<P> for FileSnapshotStore<P, C>
163where
164    C: Codec<ProjectionSnapshot<P>>,
165{
166    fn save(&mut self, snapshot: &ProjectionSnapshot<P>) -> Result<()> {
167        let position = snapshot
168            .cursor
169            .global_position
170            .ok_or(MemoryError::InvalidValue {
171                field: "snapshot.cursor",
172                reason: "cannot persist an empty replay cursor",
173            })?;
174        let final_path = self.final_path(position);
175        let frame = self.encode_frame(snapshot)?;
176        if final_path.exists() {
177            let existing =
178                fs::read(&final_path).map_err(|error| io("read existing snapshot", error))?;
179            if existing == frame {
180                return Ok(());
181            }
182            return Err(MemoryError::InvalidValue {
183                field: "snapshot",
184                reason: "different snapshot already exists at this position",
185            });
186        }
187        let temporary = temporary_path(&final_path);
188        let mut guard = TempGuard::new(temporary.clone());
189        let mut file = OpenOptions::new()
190            .write(true)
191            .create_new(true)
192            .open(&temporary)
193            .map_err(|error| io("create temporary snapshot", error))?;
194        file.write_all(&frame)
195            .map_err(|error| io("write snapshot", error))?;
196        match self.options.durability {
197            Durability::Flush => file.flush().map_err(|error| io("flush snapshot", error))?,
198            Durability::SyncData => file
199                .sync_data()
200                .map_err(|error| io("sync snapshot", error))?,
201        }
202        drop(file);
203        fs::rename(&temporary, &final_path).map_err(|error| io("commit snapshot", error))?;
204        guard.committed = true;
205        Ok(())
206    }
207
208    fn load_latest(&self) -> Result<Option<ProjectionSnapshot<P>>> {
209        self.latest_path()?
210            .map(|(position, path)| self.read_path(&path, position))
211            .transpose()
212    }
213}
214
215fn temporary_path(final_path: &Path) -> PathBuf {
216    static NEXT: AtomicU64 = AtomicU64::new(0);
217    let id = NEXT.fetch_add(1, Ordering::Relaxed);
218    let name = final_path
219        .file_name()
220        .and_then(|name| name.to_str())
221        .unwrap_or("snapshot");
222    final_path.with_file_name(format!(".{name}.tmp-{}-{id}", std::process::id()))
223}
224
225struct TempGuard {
226    path: PathBuf,
227    committed: bool,
228}
229
230impl TempGuard {
231    fn new(path: PathBuf) -> Self {
232        Self {
233            path,
234            committed: false,
235        }
236    }
237}
238
239impl Drop for TempGuard {
240    fn drop(&mut self) {
241        if !self.committed {
242            let _ = fs::remove_file(&self.path);
243        }
244    }
245}
246
247fn corrupt(reason: &str) -> MemoryError {
248    MemoryError::CorruptLog {
249        offset: 0,
250        reason: reason.to_owned(),
251    }
252}
253
254#[allow(clippy::needless_pass_by_value)]
255fn io(operation: &'static str, error: std::io::Error) -> MemoryError {
256    MemoryError::Io {
257        operation,
258        message: error.to_string(),
259    }
260}