From 8dff278b7227007d8a2c2f6c6d6f76c0c5d7a1ac Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Thu, 14 Mar 2024 17:24:53 +0100
Subject: [PATCH] [db-snapshot] Implement db snapshotting logic in garage_db

---
 Cargo.nix                |  1 +
 src/db/Cargo.toml        |  2 +-
 src/db/lib.rs            | 12 ++++++++++++
 src/db/lmdb_adapter.rs   | 10 ++++++++++
 src/db/sled_adapter.rs   |  8 ++++++++
 src/db/sqlite_adapter.rs | 12 ++++++++++++
 6 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/Cargo.nix b/Cargo.nix
index 0a4ca99a..453d83bb 100644
--- a/Cargo.nix
+++ b/Cargo.nix
@@ -4769,6 +4769,7 @@ in
     registry = "registry+https://github.com/rust-lang/crates.io-index";
     src = fetchCratesIo { inherit name version; sha256 = "a78046161564f5e7cd9008aff3b2990b3850dc8e0349119b98e8f251e099f24d"; };
     features = builtins.concatLists [
+      (lib.optional (rootFeatures' ? "garage/bundled-libs" || rootFeatures' ? "garage/default" || rootFeatures' ? "garage/sqlite" || rootFeatures' ? "garage_db/bundled-libs" || rootFeatures' ? "garage_db/default" || rootFeatures' ? "garage_db/rusqlite" || rootFeatures' ? "garage_db/sqlite" || rootFeatures' ? "garage_model/default" || rootFeatures' ? "garage_model/sqlite") "backup")
       (lib.optional (rootFeatures' ? "garage/bundled-libs" || rootFeatures' ? "garage/default" || rootFeatures' ? "garage_db/bundled-libs") "bundled")
       (lib.optional (rootFeatures' ? "garage/bundled-libs" || rootFeatures' ? "garage/default" || rootFeatures' ? "garage_db/bundled-libs") "modern_sqlite")
     ];
diff --git a/src/db/Cargo.toml b/src/db/Cargo.toml
index d7c89620..324de74c 100644
--- a/src/db/Cargo.toml
+++ b/src/db/Cargo.toml
@@ -17,7 +17,7 @@ hexdump.workspace = true
 tracing.workspace = true
 
 heed = { workspace = true, optional = true }
-rusqlite = { workspace = true, optional = true }
+rusqlite = { workspace = true, optional = true, features = ["backup"] }
 sled = { workspace = true, optional = true }
 
 [dev-dependencies]
diff --git a/src/db/lib.rs b/src/db/lib.rs
index 0fb457ce..7f19172f 100644
--- a/src/db/lib.rs
+++ b/src/db/lib.rs
@@ -19,6 +19,7 @@ use core::ops::{Bound, RangeBounds};
 
 use std::borrow::Cow;
 use std::cell::Cell;
+use std::path::PathBuf;
 use std::sync::Arc;
 
 use err_derive::Error;
@@ -48,6 +49,12 @@ pub type TxValueIter<'a> = Box<dyn std::iter::Iterator<Item = TxOpResult<(Value,
 #[error(display = "{}", _0)]
 pub struct Error(pub Cow<'static, str>);
 
+impl From<std::io::Error> for Error {
+	fn from(e: std::io::Error) -> Error {
+		Error(format!("IO: {}", e).into())
+	}
+}
+
 pub type Result<T> = std::result::Result<T, Error>;
 
 #[derive(Debug, Error)]
@@ -129,6 +136,10 @@ impl Db {
 		}
 	}
 
+	pub fn snapshot(&self, path: &PathBuf) -> Result<()> {
+		self.0.snapshot(path)
+	}
+
 	pub fn import(&self, other: &Db) -> Result<()> {
 		let existing_trees = self.list_trees()?;
 		if !existing_trees.is_empty() {
@@ -325,6 +336,7 @@ pub(crate) trait IDb: Send + Sync {
 	fn engine(&self) -> String;
 	fn open_tree(&self, name: &str) -> Result<usize>;
 	fn list_trees(&self) -> Result<Vec<String>>;
+	fn snapshot(&self, path: &PathBuf) -> Result<()>;
 
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value>>;
 	fn len(&self, tree: usize) -> Result<usize>;
diff --git a/src/db/lmdb_adapter.rs b/src/db/lmdb_adapter.rs
index 59fa132d..4b131aff 100644
--- a/src/db/lmdb_adapter.rs
+++ b/src/db/lmdb_adapter.rs
@@ -3,6 +3,7 @@ use core::ptr::NonNull;
 
 use std::collections::HashMap;
 use std::convert::TryInto;
+use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
 
 use heed::types::ByteSlice;
@@ -102,6 +103,15 @@ impl IDb for LmdbDb {
 		Ok(ret2)
 	}
 
+	fn snapshot(&self, to: &PathBuf) -> Result<()> {
+		std::fs::create_dir_all(to)?;
+		let mut path = to.clone();
+		path.push("data.mdb");
+		self.db
+			.copy_to_path(path, heed::CompactionOption::Disabled)?;
+		Ok(())
+	}
+
 	// ----
 
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
diff --git a/src/db/sled_adapter.rs b/src/db/sled_adapter.rs
index 84f2001b..c34b4d81 100644
--- a/src/db/sled_adapter.rs
+++ b/src/db/sled_adapter.rs
@@ -2,6 +2,7 @@ use core::ops::Bound;
 
 use std::cell::Cell;
 use std::collections::HashMap;
+use std::path::PathBuf;
 use std::sync::{Arc, RwLock};
 
 use sled::transaction::{
@@ -96,6 +97,13 @@ impl IDb for SledDb {
 		Ok(trees)
 	}
 
+	fn snapshot(&self, to: &PathBuf) -> Result<()> {
+		let to_db = sled::open(to)?;
+		let export = self.db.export();
+		to_db.import(export);
+		Ok(())
+	}
+
 	// ----
 
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value>> {
diff --git a/src/db/sqlite_adapter.rs b/src/db/sqlite_adapter.rs
index 9f967c66..827f3cc3 100644
--- a/src/db/sqlite_adapter.rs
+++ b/src/db/sqlite_adapter.rs
@@ -2,6 +2,7 @@ use core::ops::Bound;
 
 use std::borrow::BorrowMut;
 use std::marker::PhantomPinned;
+use std::path::PathBuf;
 use std::pin::Pin;
 use std::ptr::NonNull;
 use std::sync::{Arc, Mutex, MutexGuard};
@@ -119,6 +120,17 @@ impl IDb for SqliteDb {
 		Ok(trees)
 	}
 
+	fn snapshot(&self, to: &PathBuf) -> Result<()> {
+		fn progress(p: rusqlite::backup::Progress) {
+			let percent = (p.pagecount - p.remaining) * 100 / p.pagecount;
+			info!("Sqlite snapshot progres: {}%", percent);
+		}
+		let this = self.0.lock().unwrap();
+		this.db
+			.backup(rusqlite::DatabaseName::Main, to, Some(progress))?;
+		Ok(())
+	}
+
 	// ----
 
 	fn get(&self, tree: usize, key: &[u8]) -> Result<Option<Value>> {