Skip to main content

weavatrix_memory/snapshot/file/
mod.rs

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