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
22#[cfg(unix)]
23use std::fs::File;
24use std::fs::OpenOptions;
25use std::io::{BufWriter, Read, Write};
26use std::path::{Path, PathBuf};
27
28#[cfg(unix)]
29use anyhow::Context;
30use anyhow::Result;
31
32use lora_snapshot::{
33 decode_snapshot as decode_database_snapshot, read_snapshot as read_database_snapshot,
34 write_snapshot as write_database_snapshot, Compression, SnapshotCodecError,
35 SnapshotCredentials, SnapshotInfo, SnapshotOptions, DATABASE_SNAPSHOT_MAGIC,
36};
37use lora_store::{InMemoryGraph, SnapshotMeta, SnapshotPayload};
38
39use crate::error::{LoraError, LoraErrorCode};
40use crate::Database;
41
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum SnapshotByteFormat {
48 Database,
49}
50
51impl SnapshotByteFormat {
52 pub fn detect(bytes: &[u8]) -> Option<Self> {
53 if bytes.starts_with(DATABASE_SNAPSHOT_MAGIC) {
54 Some(Self::Database)
55 } else {
56 None
57 }
58 }
59}
60
61pub(crate) fn snapshot_info_to_meta(info: SnapshotInfo) -> SnapshotMeta {
62 SnapshotMeta {
63 format_version: info.format_version,
64 node_count: info.node_count,
65 relationship_count: info.relationship_count,
66 wal_lsn: info.wal_lsn,
67 }
68}
69
70pub(crate) fn snapshot_tmp_path(target: &Path) -> PathBuf {
75 let mut tmp = target.as_os_str().to_owned();
76 tmp.push(".tmp");
77 PathBuf::from(tmp)
78}
79
80#[cfg(unix)]
81pub(crate) fn sync_parent_dir(path: &Path) -> Result<()> {
82 let Some(parent) = path.parent() else {
83 return Ok(());
84 };
85 let dir = File::open(parent).with_context(|| format!("open dir {}", parent.display()))?;
86 dir.sync_all()
87 .with_context(|| format!("sync dir {}", parent.display()))
88}
89
90#[cfg(not(unix))]
91pub(crate) fn sync_parent_dir(_path: &Path) -> Result<()> {
92 Ok(())
93}
94
95pub(crate) struct TempFileGuard {
104 path: Option<PathBuf>,
105}
106
107impl TempFileGuard {
108 pub(crate) fn new(path: PathBuf) -> Self {
109 Self { path: Some(path) }
110 }
111
112 pub(crate) fn commit(mut self) {
116 self.path.take();
117 }
118}
119
120impl Drop for TempFileGuard {
121 fn drop(&mut self) {
122 if let Some(path) = self.path.take() {
123 let _ = std::fs::remove_file(path);
127 }
128 }
129}
130
131pub(crate) fn decode_snapshot_bytes(
135 bytes: &[u8],
136 credentials: Option<&SnapshotCredentials>,
137) -> Result<(SnapshotPayload, SnapshotInfo), SnapshotCodecError> {
138 decode_database_snapshot(bytes, credentials)
139}
140
141pub(crate) fn read_snapshot_from<R: Read>(
144 reader: R,
145 credentials: Option<&SnapshotCredentials>,
146) -> Result<(SnapshotPayload, SnapshotInfo), SnapshotCodecError> {
147 read_database_snapshot(reader, credentials)
148}
149
150pub(crate) fn encode_snapshot_to<W: Write>(
152 writer: W,
153 payload: &SnapshotPayload,
154 wal_lsn: Option<u64>,
155 options: &SnapshotOptions,
156) -> Result<SnapshotInfo, SnapshotCodecError> {
157 write_database_snapshot(writer, payload, wal_lsn, options)
158}
159
160impl Database<InMemoryGraph> {
165 pub fn save_snapshot_to(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
172 let options = SnapshotOptions {
173 compression: Compression::None,
174 encryption: None,
175 };
176 self.save_snapshot_to_with_options(path, &options)
177 }
178
179 pub fn load_snapshot_from(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
183 self.load_snapshot_from_with_credentials(path, None)
184 }
185
186 pub fn in_memory_from_snapshot(path: impl AsRef<Path>) -> Result<Self, LoraError> {
190 let db = Self::in_memory();
191 db.load_snapshot_from_with_credentials(path, None)?;
192 Ok(db)
193 }
194
195 pub fn save_snapshot_to_bytes(&self) -> Result<Vec<u8>, LoraError> {
204 let options = SnapshotOptions {
205 compression: Compression::None,
206 encryption: None,
207 };
208 let (bytes, _) = self.save_snapshot_to_bytes_with_options(&options)?;
209 Ok(bytes)
210 }
211
212 pub fn save_snapshot_to_bytes_with_options(
215 &self,
216 options: &SnapshotOptions,
217 ) -> Result<(Vec<u8>, SnapshotInfo), LoraError> {
218 let guard = self.read_store();
219 let payload = guard.snapshot_payload();
220 let mut bytes = Vec::new();
221 let info = encode_snapshot_to(&mut bytes, &payload, None, options)?;
222 Ok((bytes, info))
223 }
224
225 pub fn save_snapshot_to_with_options(
230 &self,
231 path: impl AsRef<Path>,
232 options: &SnapshotOptions,
233 ) -> Result<SnapshotMeta, LoraError> {
234 let path = path.as_ref();
235 let tmp = snapshot_tmp_path(path);
236 let guard = self.read_store();
237
238 let file = OpenOptions::new()
239 .write(true)
240 .create(true)
241 .truncate(true)
242 .open(&tmp)?;
243 let tmp_guard = TempFileGuard::new(tmp.clone());
244 let mut writer = BufWriter::new(file);
245
246 let payload = guard.snapshot_payload();
247 let info = encode_snapshot_to(&mut writer, &payload, None, options)?;
248
249 writer.flush()?;
250 let file = writer.into_inner().map_err(|e| e.into_error())?;
251 file.sync_all()?;
252 drop(file);
253
254 std::fs::rename(&tmp, path)?;
255 tmp_guard.commit();
256
257 sync_parent_dir(path).map_err(|e| LoraError::new(LoraErrorCode::Io, e.to_string()))?;
258
259 Ok(snapshot_info_to_meta(info))
260 }
261
262 pub fn load_snapshot_from_bytes(&self, bytes: &[u8]) -> Result<SnapshotMeta, LoraError> {
265 self.load_snapshot_from_bytes_with_credentials(bytes, None)
266 }
267
268 pub fn load_snapshot_from_bytes_with_credentials(
271 &self,
272 bytes: &[u8],
273 credentials: Option<&SnapshotCredentials>,
274 ) -> Result<SnapshotMeta, LoraError> {
275 if SnapshotByteFormat::detect(bytes).is_none() {
276 return Err(LoraError::new(
277 LoraErrorCode::SnapshotCodec,
278 "snapshot bytes have unrecognized magic",
279 ));
280 }
281 let mut guard = self.write_store();
282 let (payload, info) = decode_snapshot_bytes(bytes, credentials)?;
283 let meta = snapshot_info_to_meta(info);
284 guard.load_snapshot_payload(payload)?;
285 guard.publish();
289 Ok(meta)
290 }
291
292 pub fn load_snapshot_from_with_credentials(
295 &self,
296 path: impl AsRef<Path>,
297 credentials: Option<&SnapshotCredentials>,
298 ) -> Result<SnapshotMeta, LoraError> {
299 let bytes = std::fs::read(path.as_ref())?;
300 self.load_snapshot_from_bytes_with_credentials(&bytes, credentials)
301 }
302
303 pub fn checkpoint_to(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
317 let recorder = self.wal.as_ref().ok_or_else(|| {
318 LoraError::new(LoraErrorCode::Internal, "checkpoint requires WAL enabled")
319 })?;
320 let path = path.as_ref();
321 let tmp = snapshot_tmp_path(path);
322
323 let guard = self.write_store();
324
325 recorder.force_fsync()?;
328 let snapshot_lsn = recorder.wal().durable_lsn();
329
330 let file = OpenOptions::new()
331 .write(true)
332 .create(true)
333 .truncate(true)
334 .open(&tmp)?;
335 let tmp_guard = TempFileGuard::new(tmp.clone());
336 let mut writer = BufWriter::new(file);
337 let payload = guard.snapshot_payload();
338 let options = SnapshotOptions {
339 compression: Compression::None,
340 encryption: None,
341 };
342 let info = encode_snapshot_to(&mut writer, &payload, Some(snapshot_lsn.raw()), &options)?;
343 let meta = snapshot_info_to_meta(info);
344
345 writer.flush()?;
346 let file = writer.into_inner().map_err(|e| e.into_error())?;
347 file.sync_all()?;
348 drop(file);
349
350 std::fs::rename(&tmp, path)?;
351 tmp_guard.commit();
352
353 sync_parent_dir(path).map_err(|e| LoraError::new(LoraErrorCode::Io, e.to_string()))?;
354
355 recorder.checkpoint_marker(snapshot_lsn)?;
359 recorder.force_fsync()?;
360
361 if let Err(err) = recorder.truncate_up_to(snapshot_lsn) {
364 tracing::warn!(
365 lsn = snapshot_lsn.raw(),
366 error = %err,
367 "WAL truncation after checkpoint failed; will retry later"
368 );
369 }
370
371 Ok(meta)
372 }
373
374 pub fn checkpoint_managed(&self) -> Result<SnapshotMeta, LoraError> {
377 let recorder = self.wal.as_ref().ok_or_else(|| {
378 LoraError::new(
379 LoraErrorCode::Internal,
380 "managed checkpoint requires WAL enabled",
381 )
382 })?;
383 let snapshots = self.snapshots.as_ref().ok_or_else(|| {
384 LoraError::new(
385 LoraErrorCode::Internal,
386 "managed checkpoint requires snapshots enabled",
387 )
388 })?;
389 let guard = self.write_store();
390 snapshots.checkpoint(&guard, recorder).map_err(Into::into)
391 }
392}
393
394pub trait SnapshotAdmin: Send + Sync + 'static {
406 fn save_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError>;
407 fn load_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError>;
408}
409
410impl SnapshotAdmin for Database<InMemoryGraph> {
411 fn save_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError> {
412 self.save_snapshot_to(path)
413 }
414
415 fn load_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError> {
416 self.load_snapshot_from(path)
417 }
418}