lora_database/snapshot/
store.rs

1use std::fs::{self, File, OpenOptions};
2use std::io::{BufReader, BufWriter, Write};
3use std::path::{Path, PathBuf};
4use std::sync::atomic::{AtomicU64, Ordering};
5
6use anyhow::{anyhow, Context, Result};
7use lora_snapshot::{read_snapshot, write_snapshot, SnapshotOptions};
8use lora_store::{InMemoryGraph, SnapshotMeta};
9use lora_wal::{Lsn, WalRecorder};
10
11const CURRENT_FILE: &str = "CURRENT";
12const SNAPSHOT_PREFIX: &str = "snapshot-";
13const SNAPSHOT_SUFFIX: &str = ".lsnap";
14
15#[derive(Debug, Clone)]
16pub struct SnapshotConfig {
17    pub dir: PathBuf,
18    /// When set, create a managed checkpoint after this many committed WAL
19    /// transactions. `None` keeps checkpointing manual via `sync()` /
20    /// `checkpoint_managed()`.
21    pub checkpoint_every_commits: Option<u64>,
22    /// Number of older checkpoint files to retain in addition to `CURRENT`.
23    pub keep_old: usize,
24    /// Columnar snapshot codec options. Defaults to fast gzip compression and
25    /// no encryption.
26    pub codec: SnapshotOptions,
27}
28
29impl SnapshotConfig {
30    pub fn enabled(dir: impl Into<PathBuf>) -> Self {
31        Self {
32            dir: dir.into(),
33            checkpoint_every_commits: None,
34            keep_old: 1,
35            codec: SnapshotOptions::default(),
36        }
37    }
38
39    pub fn every_commits(mut self, commits: u64) -> Self {
40        self.checkpoint_every_commits = Some(commits.max(1));
41        self
42    }
43
44    pub fn keep_old(mut self, keep_old: usize) -> Self {
45        self.keep_old = keep_old;
46        self
47    }
48
49    pub fn codec(mut self, codec: SnapshotOptions) -> Self {
50        self.codec = codec;
51        self
52    }
53}
54
55pub(crate) struct ManagedSnapshotStore {
56    config: SnapshotConfig,
57    commits_since_checkpoint: AtomicU64,
58}
59
60impl ManagedSnapshotStore {
61    pub(crate) fn open(config: SnapshotConfig) -> Result<Self> {
62        fs::create_dir_all(&config.dir)
63            .with_context(|| format!("create snapshot dir {}", config.dir.display()))?;
64        Ok(Self {
65            config,
66            commits_since_checkpoint: AtomicU64::new(0),
67        })
68    }
69
70    pub(crate) fn load_latest(&self, graph: &mut InMemoryGraph) -> Result<Lsn> {
71        let Some(path) = self.latest_snapshot_path()? else {
72            return Ok(Lsn::ZERO);
73        };
74        let file =
75            File::open(&path).with_context(|| format!("open snapshot {}", path.display()))?;
76        let (payload, info) =
77            read_snapshot(BufReader::new(file), self.config.codec.encryption.as_ref())
78                .with_context(|| format!("load snapshot {}", path.display()))?;
79        graph.load_snapshot_payload(payload)?;
80        Ok(info.wal_lsn.map(Lsn::new).unwrap_or(Lsn::ZERO))
81    }
82
83    pub(crate) fn checkpoint(
84        &self,
85        graph: &InMemoryGraph,
86        recorder: &WalRecorder,
87    ) -> Result<SnapshotMeta> {
88        recorder
89            .force_fsync()
90            .map_err(|e| anyhow!("WAL fsync before managed snapshot failed: {e}"))?;
91        let snapshot_lsn = recorder.wal().durable_lsn();
92        let meta = self.write_snapshot(graph, snapshot_lsn)?;
93
94        recorder
95            .checkpoint_marker(snapshot_lsn)
96            .map_err(|e| anyhow!("WAL checkpoint marker failed: {e}"))?;
97        recorder
98            .force_fsync()
99            .map_err(|e| anyhow!("WAL fsync after checkpoint marker failed: {e}"))?;
100        if let Err(err) = recorder.truncate_up_to(snapshot_lsn) {
101            tracing::warn!(
102                lsn = snapshot_lsn.raw(),
103                error = %err,
104                "WAL truncation after managed checkpoint failed; will retry later"
105            );
106        }
107
108        self.commits_since_checkpoint.store(0, Ordering::Relaxed);
109        self.prune_old_snapshots(snapshot_lsn)?;
110        Ok(meta)
111    }
112
113    pub(crate) fn observe_commit(
114        &self,
115        graph: &InMemoryGraph,
116        recorder: &WalRecorder,
117    ) -> Result<()> {
118        let Some(every) = self.config.checkpoint_every_commits else {
119            return Ok(());
120        };
121        let commits = self
122            .commits_since_checkpoint
123            .fetch_add(1, Ordering::Relaxed)
124            + 1;
125        if commits >= every {
126            self.checkpoint(graph, recorder)?;
127        }
128        Ok(())
129    }
130
131    fn write_snapshot(&self, graph: &InMemoryGraph, snapshot_lsn: Lsn) -> Result<SnapshotMeta> {
132        let target = snapshot_path(&self.config.dir, snapshot_lsn);
133        let tmp = tmp_path(&target);
134        let file = OpenOptions::new()
135            .write(true)
136            .create(true)
137            .truncate(true)
138            .open(&tmp)
139            .with_context(|| format!("open temp snapshot {}", tmp.display()))?;
140        let mut writer = BufWriter::new(file);
141        let payload = graph.snapshot_payload();
142        let info = write_snapshot(
143            &mut writer,
144            &payload,
145            Some(snapshot_lsn.raw()),
146            &self.config.codec,
147        )
148        .map_err(|e| anyhow!("encode managed snapshot failed: {e}"))?;
149        let meta = SnapshotMeta {
150            format_version: info.format_version,
151            node_count: info.node_count,
152            relationship_count: info.relationship_count,
153            wal_lsn: info.wal_lsn,
154        };
155        writer.flush()?;
156        let file = writer.into_inner().map_err(|e| e.into_error())?;
157        file.sync_all()?;
158        drop(file);
159
160        fs::rename(&tmp, &target)
161            .with_context(|| format!("rename {} to {}", tmp.display(), target.display()))?;
162        sync_dir(&self.config.dir)
163            .with_context(|| format!("sync snapshot dir {}", self.config.dir.display()))?;
164        write_current(&self.config.dir, &target)?;
165        Ok(meta)
166    }
167
168    fn latest_snapshot_path(&self) -> Result<Option<PathBuf>> {
169        let current = self.config.dir.join(CURRENT_FILE);
170        match fs::read_to_string(&current) {
171            Ok(name) => {
172                let name = name.trim();
173                if name.is_empty() {
174                    return Ok(None);
175                }
176                let path = self.config.dir.join(name);
177                if path.exists() {
178                    return Ok(Some(path));
179                }
180            }
181            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
182            Err(err) => return Err(err).with_context(|| format!("read {}", current.display())),
183        }
184
185        let latest = snapshot_files(&self.config.dir)?
186            .into_iter()
187            .max_by_key(|(lsn, _)| *lsn)
188            .map(|(_, path)| path);
189        Ok(latest)
190    }
191
192    fn prune_old_snapshots(&self, current_lsn: Lsn) -> Result<()> {
193        let mut snapshots = snapshot_files(&self.config.dir)?;
194        snapshots.retain(|(lsn, _)| *lsn != current_lsn);
195        snapshots.sort_by_key(|(lsn, _)| *lsn);
196        let retain = self.config.keep_old;
197        let remove_count = snapshots.len().saturating_sub(retain);
198        for (_, path) in snapshots.into_iter().take(remove_count) {
199            fs::remove_file(&path)
200                .with_context(|| format!("remove old snapshot {}", path.display()))?;
201        }
202        sync_dir(&self.config.dir)
203            .with_context(|| format!("sync snapshot dir {}", self.config.dir.display()))?;
204        Ok(())
205    }
206}
207
208fn snapshot_path(dir: &Path, lsn: Lsn) -> PathBuf {
209    dir.join(format!(
210        "{SNAPSHOT_PREFIX}{:020}{SNAPSHOT_SUFFIX}",
211        lsn.raw()
212    ))
213}
214
215fn snapshot_files(dir: &Path) -> Result<Vec<(Lsn, PathBuf)>> {
216    let mut out = Vec::new();
217    for entry in
218        fs::read_dir(dir).with_context(|| format!("read snapshot dir {}", dir.display()))?
219    {
220        let entry = entry?;
221        let path = entry.path();
222        let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
223            continue;
224        };
225        let Some(raw) = name
226            .strip_prefix(SNAPSHOT_PREFIX)
227            .and_then(|name| name.strip_suffix(SNAPSHOT_SUFFIX))
228        else {
229            continue;
230        };
231        if let Ok(lsn) = raw.parse::<u64>() {
232            out.push((Lsn::new(lsn), path));
233        }
234    }
235    Ok(out)
236}
237
238fn write_current(dir: &Path, target: &Path) -> Result<()> {
239    let name = target
240        .file_name()
241        .and_then(|name| name.to_str())
242        .ok_or_else(|| {
243            anyhow!(
244                "snapshot path has no portable filename: {}",
245                target.display()
246            )
247        })?;
248    let current = dir.join(CURRENT_FILE);
249    let tmp = tmp_path(&current);
250    let mut file = OpenOptions::new()
251        .write(true)
252        .create(true)
253        .truncate(true)
254        .open(&tmp)
255        .with_context(|| format!("open temp CURRENT {}", tmp.display()))?;
256    writeln!(file, "{name}")?;
257    file.sync_all()?;
258    drop(file);
259    fs::rename(&tmp, &current)
260        .with_context(|| format!("rename {} to {}", tmp.display(), current.display()))?;
261    sync_dir(dir).with_context(|| format!("sync snapshot dir {}", dir.display()))?;
262    Ok(())
263}
264
265fn tmp_path(path: &Path) -> PathBuf {
266    let mut tmp = path.as_os_str().to_owned();
267    tmp.push(".tmp");
268    PathBuf::from(tmp)
269}
270
271#[cfg(unix)]
272fn sync_dir(path: &Path) -> Result<()> {
273    let dir = File::open(path).with_context(|| format!("open dir {}", path.display()))?;
274    dir.sync_all()
275        .with_context(|| format!("sync dir {}", path.display()))
276}
277
278#[cfg(not(unix))]
279fn sync_dir(_path: &Path) -> Result<()> {
280    Ok(())
281}
lora_database/snapshot/store.rs

lora_database/snapshot/
store.rs
