mod json;
pub(crate) mod store;
pub use json::{snapshot_credentials_from_json, snapshot_options_from_json};
pub(crate) use store::ManagedSnapshotStore;
pub use store::SnapshotConfig;
use std::fs::OpenOptions;
use std::io::{BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use anyhow::Result;
use lora_snapshot::{
decode_snapshot as decode_database_snapshot, read_snapshot as read_database_snapshot,
write_snapshot as write_database_snapshot, Compression, SnapshotCodecError,
SnapshotCredentials, SnapshotInfo, SnapshotOptions, DATABASE_SNAPSHOT_MAGIC,
};
use lora_store::{InMemoryGraph, SnapshotMeta, SnapshotPayload};
use crate::durable_io::{sync_dir, sync_file};
use crate::error::{LoraError, LoraErrorCode};
use crate::Database;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SnapshotByteFormat {
Database,
}
impl SnapshotByteFormat {
pub fn detect(bytes: &[u8]) -> Option<Self> {
if bytes.starts_with(DATABASE_SNAPSHOT_MAGIC) {
Some(Self::Database)
} else {
None
}
}
}
pub(crate) fn snapshot_info_to_meta(info: SnapshotInfo) -> SnapshotMeta {
SnapshotMeta {
format_version: info.format_version,
node_count: info.node_count,
relationship_count: info.relationship_count,
wal_lsn: info.wal_lsn,
}
}
pub(crate) fn snapshot_tmp_path(target: &Path) -> PathBuf {
let mut tmp = target.as_os_str().to_owned();
tmp.push(".tmp");
PathBuf::from(tmp)
}
pub(crate) fn sync_parent_dir(path: &Path) -> Result<()> {
let Some(parent) = path.parent() else {
return Ok(());
};
Ok(sync_dir(parent)?)
}
pub(crate) struct TempFileGuard {
path: Option<PathBuf>,
}
impl TempFileGuard {
pub(crate) fn new(path: PathBuf) -> Self {
Self { path: Some(path) }
}
pub(crate) fn commit(mut self) {
self.path.take();
}
}
impl Drop for TempFileGuard {
fn drop(&mut self) {
if let Some(path) = self.path.take() {
let _ = std::fs::remove_file(path);
}
}
}
pub(crate) fn decode_snapshot_bytes(
bytes: &[u8],
credentials: Option<&SnapshotCredentials>,
) -> Result<(SnapshotPayload, SnapshotInfo), SnapshotCodecError> {
decode_database_snapshot(bytes, credentials)
}
pub(crate) fn read_snapshot_from<R: Read>(
reader: R,
credentials: Option<&SnapshotCredentials>,
) -> Result<(SnapshotPayload, SnapshotInfo), SnapshotCodecError> {
read_database_snapshot(reader, credentials)
}
pub(crate) fn encode_snapshot_to<W: Write>(
writer: W,
payload: &SnapshotPayload,
wal_lsn: Option<u64>,
options: &SnapshotOptions,
) -> Result<SnapshotInfo, SnapshotCodecError> {
write_database_snapshot(writer, payload, wal_lsn, options)
}
impl Database<InMemoryGraph> {
pub fn save_snapshot_to(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
let options = SnapshotOptions {
compression: Compression::None,
encryption: None,
};
self.save_snapshot_to_with_options(path, &options)
}
pub fn load_snapshot_from(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
self.load_snapshot_from_with_credentials(path, None)
}
pub fn in_memory_from_snapshot(path: impl AsRef<Path>) -> Result<Self, LoraError> {
let db = Self::in_memory();
db.load_snapshot_from_with_credentials(path, None)?;
Ok(db)
}
pub fn save_snapshot_to_bytes(&self) -> Result<Vec<u8>, LoraError> {
let options = SnapshotOptions {
compression: Compression::None,
encryption: None,
};
let (bytes, _) = self.save_snapshot_to_bytes_with_options(&options)?;
Ok(bytes)
}
pub fn save_snapshot_to_bytes_with_options(
&self,
options: &SnapshotOptions,
) -> Result<(Vec<u8>, SnapshotInfo), LoraError> {
let guard = self.read_store();
let payload = guard.snapshot_payload();
let mut bytes = Vec::new();
let info = encode_snapshot_to(&mut bytes, &payload, None, options)?;
Ok((bytes, info))
}
pub fn save_snapshot_to_with_options(
&self,
path: impl AsRef<Path>,
options: &SnapshotOptions,
) -> Result<SnapshotMeta, LoraError> {
let path = path.as_ref();
let tmp = snapshot_tmp_path(path);
let guard = self.read_store();
let file = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(&tmp)?;
let tmp_guard = TempFileGuard::new(tmp.clone());
let mut writer = BufWriter::new(file);
let payload = guard.snapshot_payload();
let info = encode_snapshot_to(&mut writer, &payload, None, options)?;
writer.flush()?;
let file = writer.into_inner().map_err(|e| e.into_error())?;
sync_file(&file)?;
drop(file);
std::fs::rename(&tmp, path)?;
tmp_guard.commit();
sync_parent_dir(path).map_err(|e| LoraError::new(LoraErrorCode::Io, e.to_string()))?;
Ok(snapshot_info_to_meta(info))
}
pub fn load_snapshot_from_bytes(&self, bytes: &[u8]) -> Result<SnapshotMeta, LoraError> {
self.load_snapshot_from_bytes_with_credentials(bytes, None)
}
pub fn load_snapshot_from_bytes_with_credentials(
&self,
bytes: &[u8],
credentials: Option<&SnapshotCredentials>,
) -> Result<SnapshotMeta, LoraError> {
if SnapshotByteFormat::detect(bytes).is_none() {
return Err(LoraError::new(
LoraErrorCode::SnapshotCodec,
"snapshot bytes have unrecognized magic",
));
}
let mut guard = self.write_store();
let (payload, info) = decode_snapshot_bytes(bytes, credentials)?;
let meta = snapshot_info_to_meta(info);
guard
.staged_mut_or_error()?
.load_snapshot_payload(payload)?;
guard.publish();
Ok(meta)
}
pub fn load_snapshot_from_with_credentials(
&self,
path: impl AsRef<Path>,
credentials: Option<&SnapshotCredentials>,
) -> Result<SnapshotMeta, LoraError> {
let bytes = std::fs::read(path.as_ref())?;
self.load_snapshot_from_bytes_with_credentials(&bytes, credentials)
}
pub fn checkpoint_to(&self, path: impl AsRef<Path>) -> Result<SnapshotMeta, LoraError> {
let recorder = self.wal.as_ref().ok_or_else(|| {
LoraError::new(LoraErrorCode::Internal, "checkpoint requires WAL enabled")
})?;
let path = path.as_ref();
let tmp = snapshot_tmp_path(path);
let guard = self.write_store();
recorder.force_fsync()?;
let snapshot_lsn = recorder.wal().durable_lsn();
let file = OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(&tmp)?;
let tmp_guard = TempFileGuard::new(tmp.clone());
let mut writer = BufWriter::new(file);
let payload = guard.staged_or_error()?.snapshot_payload();
let options = SnapshotOptions {
compression: Compression::None,
encryption: None,
};
let info = encode_snapshot_to(&mut writer, &payload, Some(snapshot_lsn.raw()), &options)?;
let meta = snapshot_info_to_meta(info);
writer.flush()?;
let file = writer.into_inner().map_err(|e| e.into_error())?;
sync_file(&file)?;
drop(file);
std::fs::rename(&tmp, path)?;
tmp_guard.commit();
sync_parent_dir(path).map_err(|e| LoraError::new(LoraErrorCode::Io, e.to_string()))?;
recorder.checkpoint_marker(snapshot_lsn)?;
recorder.force_fsync()?;
if let Err(err) = recorder.truncate_up_to(snapshot_lsn) {
tracing::warn!(
lsn = snapshot_lsn.raw(),
error = %err,
"WAL truncation after checkpoint failed; will retry later"
);
}
Ok(meta)
}
pub fn checkpoint_managed(&self) -> Result<SnapshotMeta, LoraError> {
let recorder = self.wal.as_ref().ok_or_else(|| {
LoraError::new(
LoraErrorCode::Internal,
"managed checkpoint requires WAL enabled",
)
})?;
let snapshots = self.snapshots.as_ref().ok_or_else(|| {
LoraError::new(
LoraErrorCode::Internal,
"managed checkpoint requires snapshots enabled",
)
})?;
let guard = self.write_store();
snapshots
.checkpoint(guard.staged_or_error()?, recorder)
.map_err(Into::into)
}
}
pub trait SnapshotAdmin: Send + Sync + 'static {
fn save_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError>;
fn load_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError>;
}
impl SnapshotAdmin for Database<InMemoryGraph> {
fn save_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError> {
self.save_snapshot_to(path)
}
fn load_snapshot(&self, path: &Path) -> Result<SnapshotMeta, LoraError> {
self.load_snapshot_from(path)
}
}