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

lora_database/snapshot/
store.rs
