1mod json;
16pub(crate) mod store;
17
18pub use json::{snapshot_credentials_from_json, snapshot_options_from_json};
19pub(crate) use store::ManagedSnapshotStore;
20pub use store::SnapshotConfig;
21
22use std::fs::OpenOptions;
23use std::io::{BufWriter, Read, Write};
24use std::path::{Path, PathBuf};
25
26use anyhow::Result;
27
28use lora_snapshot::{
29 decode_snapshot as decode_database_snapshot, read_snapshot as read_database_snapshot,
30 write_snapshot as write_database_snapshot, Compression, SnapshotCodecError,
31 SnapshotCredentials, SnapshotInfo, SnapshotOptions, DATABASE_SNAPSHOT_MAGIC,
32};
33use lora_store::{InMemoryGraph, SnapshotMeta, SnapshotPayload};
34
35use crate::durable_io::{sync_dir, sync_file};
36use crate::error::{LoraError, LoraErrorCode};
37use crate::Database;
38
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
44pub enum SnapshotByteFormat {
45 Database,
46}
47
48impl SnapshotByteFormat {
49 pub fn detect(bytes: &[u8]) -> Option<Self> {
50 if bytes.starts_with(DATABASE_SNAPSHOT_MAGIC) {
51 Some(Self::Database)
52 } else {
53 None
54 }
55 }
56}
57
58pub(crate) fn snapshot_info_to_meta(info: SnapshotInfo) -> SnapshotMeta {
59 SnapshotMeta {
60 format_version: info.format_version,
61 node_count: info.node_count,
62 relationship_count: info.relationship_count,
63 wal_lsn: info.wal_lsn,
64 }
65}
66
67pub(crate) fn snapshot_tmp_path(target: &Path) -> PathBuf {
72 let mut tmp = target.as_os_str().to_owned();
73 tmp.push(".tmp");
74 PathBuf::from(tmp)
75}
76
77pub(crate) fn sync_parent_dir(path: &Path) -> Result<()> {
78 let Some(parent) = path.parent() else {
79 return Ok(());
80 };
81 Ok(sync_dir(parent)?)
82}
83
84pub(crate) struct TempFileGuard {
93 path: Option<PathBuf>,
94}
95
96impl TempFileGuard {
97 pub(crate) fn new(path: PathBuf) -> Self {
98 Self { path: Some(path) }
99 }
100
101 pub(crate) fn commit(mut self) {
105 self.path.take();
106 }
107}
108
109impl Drop for TempFileGuard {
110 fn drop(&mut self) {
111 if let Some(path) = self.path.take() {
112 let _ = std::fs::remove_file(path);
116 }
117 }
118}
119
120pub(crate) fn decode_snapshot_bytes(
124 bytes: &[u8],
125 credentials: Option<&SnapshotCredentials>,
126) -> Result<(SnapshotPayload, SnapshotInfo), SnapshotCodecError> {
127 decode_database_snapshot(bytes, credentials)
128}
129
130pub(crate) fn read_snapshot_from<R: Read>(
133 reader: R,
134 credentials: Option<&SnapshotCredentials>,
135) -> Result<(SnapshotPayload, SnapshotInfo), SnapshotCodecError> {
136 read_database_snapshot(reader, credentials)
137}
138
139pub(crate) fn encode_snapshot_to<W: Write>(
141 writer: W,
142 payload: &SnapshotPayload,
143 wal_lsn: Option<u64>,
144 options: &SnapshotOptions,
145) -> Result<SnapshotInfo, SnapshotCodecError> {
146 write_database_snapshot(writer, payload, wal_lsn, options)
147}
148
149impl Database<InMemoryGraph> {
154 pub fn save_snapshot_to(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
161 let options = SnapshotOptions {
162 compression: Compression::None,
163 encryption: None,
164 };
165 self.save_snapshot_to_with_options(path, &options)
166 }
167
168 pub fn load_snapshot_from(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
172 self.load_snapshot_from_with_credentials(path, None)
173 }
174
175 pub fn in_memory_from_snapshot(path: impl AsRef<Path>) -> Result<Self, LoraError> {
179 let db = Self::in_memory();
180 db.load_snapshot_from_with_credentials(path, None)?;
181 Ok(db)
182 }
183
184 pub fn save_snapshot_to_bytes(&self) -> Result<Vec<u8>, LoraError> {
193 let options = SnapshotOptions {
194 compression: Compression::None,
195 encryption: None,
196 };
197 let (bytes, _) = self.save_snapshot_to_bytes_with_options(&options)?;
198 Ok(bytes)
199 }
200
201 pub fn save_snapshot_to_bytes_with_options(
204 &self,
205 options: &SnapshotOptions,
206 ) -> Result<(Vec<u8>, SnapshotInfo), LoraError> {
207 let guard = self.read_store();
208 let payload = guard.snapshot_payload();
209 let mut bytes = Vec::new();
210 let info = encode_snapshot_to(&mut bytes, &payload, None, options)?;
211 Ok((bytes, info))
212 }
213
214 pub fn save_snapshot_to_with_options(
219 &self,
220 path: impl AsRef<Path>,
221 options: &SnapshotOptions,
222 ) -> Result<SnapshotMeta, LoraError> {
223 let path = path.as_ref();
224 let tmp = snapshot_tmp_path(path);
225 let guard = self.read_store();
226
227 let file = OpenOptions::new()
228 .write(true)
229 .create(true)
230 .truncate(true)
231 .open(&tmp)?;
232 let tmp_guard = TempFileGuard::new(tmp.clone());
233 let mut writer = BufWriter::new(file);
234
235 let payload = guard.snapshot_payload();
236 let info = encode_snapshot_to(&mut writer, &payload, None, options)?;
237
238 writer.flush()?;
239 let file = writer.into_inner().map_err(|e| e.into_error())?;
240 sync_file(&file)?;
241 drop(file);
242
243 std::fs::rename(&tmp, path)?;
244 tmp_guard.commit();
245
246 sync_parent_dir(path).map_err(|e| LoraError::new(LoraErrorCode::Io, e.to_string()))?;
247
248 Ok(snapshot_info_to_meta(info))
249 }
250
251 pub fn load_snapshot_from_bytes(&self, bytes: &[u8]) -> Result<SnapshotMeta, LoraError> {
254 self.load_snapshot_from_bytes_with_credentials(bytes, None)
255 }
256
257 pub fn load_snapshot_from_bytes_with_credentials(
260 &self,
261 bytes: &[u8],
262 credentials: Option<&SnapshotCredentials>,
263 ) -> Result<SnapshotMeta, LoraError> {
264 if SnapshotByteFormat::detect(bytes).is_none() {
265 return Err(LoraError::new(
266 LoraErrorCode::SnapshotCodec,
267 "snapshot bytes have unrecognized magic",
268 ));
269 }
270 let mut guard = self.write_store();
271 let (payload, info) = decode_snapshot_bytes(bytes, credentials)?;
272 let meta = snapshot_info_to_meta(info);
273 guard
274 .staged_mut_or_error()?
275 .load_snapshot_payload(payload)?;
276 guard.publish_in_place();
280 self.publish_reset();
283 drop(guard);
284 Ok(meta)
285 }
286
287 pub fn load_snapshot_from_with_credentials(
290 &self,
291 path: impl AsRef<Path>,
292 credentials: Option<&SnapshotCredentials>,
293 ) -> Result<SnapshotMeta, LoraError> {
294 let bytes = std::fs::read(path.as_ref())?;
295 self.load_snapshot_from_bytes_with_credentials(&bytes, credentials)
296 }
297
298 pub fn checkpoint_to(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
312 let recorder = self.wal.as_ref().ok_or_else(|| {
313 LoraError::new(LoraErrorCode::Internal, "checkpoint requires WAL enabled")
314 })?;
315 let path = path.as_ref();
316 let tmp = snapshot_tmp_path(path);
317
318 let guard = self.write_store();
319
320 recorder.force_fsync()?;
323 let snapshot_lsn = recorder.wal().durable_lsn();
324
325 let file = OpenOptions::new()
326 .write(true)
327 .create(true)
328 .truncate(true)
329 .open(&tmp)?;
330 let tmp_guard = TempFileGuard::new(tmp.clone());
331 let mut writer = BufWriter::new(file);
332 let payload = guard.staged_or_error()?.snapshot_payload();
333 let options = SnapshotOptions {
334 compression: Compression::None,
335 encryption: None,
336 };
337 let info = encode_snapshot_to(&mut writer, &payload, Some(snapshot_lsn.raw()), &options)?;
338 let meta = snapshot_info_to_meta(info);
339
340 writer.flush()?;
341 let file = writer.into_inner().map_err(|e| e.into_error())?;
342 sync_file(&file)?;
343 drop(file);
344
345 std::fs::rename(&tmp, path)?;
346 tmp_guard.commit();
347
348 sync_parent_dir(path).map_err(|e| LoraError::new(LoraErrorCode::Io, e.to_string()))?;
349
350 recorder.checkpoint_marker(snapshot_lsn)?;
354 recorder.force_fsync()?;
355
356 if let Err(err) = recorder.truncate_up_to(snapshot_lsn) {
359 tracing::warn!(
360 lsn = snapshot_lsn.raw(),
361 error = %err,
362 "WAL truncation after checkpoint failed; will retry later"
363 );
364 }
365
366 Ok(meta)
367 }
368
369 pub fn checkpoint_managed(&self) -> Result<SnapshotMeta, LoraError> {
372 let recorder = self.wal.as_ref().ok_or_else(|| {
373 LoraError::new(
374 LoraErrorCode::Internal,
375 "managed checkpoint requires WAL enabled",
376 )
377 })?;
378 let snapshots = self.snapshots.as_ref().ok_or_else(|| {
379 LoraError::new(
380 LoraErrorCode::Internal,
381 "managed checkpoint requires snapshots enabled",
382 )
383 })?;
384 let guard = self.write_store();
385 snapshots
386 .checkpoint(guard.staged_or_error()?, recorder)
387 .map_err(Into::into)
388 }
389}
390
391pub trait SnapshotAdmin: Send + Sync + 'static {
403 fn save_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError>;
404 fn load_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError>;
405}
406
407impl SnapshotAdmin for Database<InMemoryGraph> {
408 fn save_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError> {
409 self.save_snapshot_to(path)
410 }
411
412 fn load_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError> {
413 self.load_snapshot_from(path)
414 }
415}