From 2d5541b0e7ed0939e99710f40d3859c05bc31175 Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Fri, 3 Jun 2022 16:40:38 +0200
Subject: Integrate LMDB with rest

---
 src/db/bin/convert.rs  | 11 +++++++++++
 src/db/lmdb_adapter.rs | 26 +++++++++++++++++++++++++-
 src/garage/server.rs   | 22 +++++++++++++++++++++-
 src/util/config.rs     |  2 +-
 4 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/src/db/bin/convert.rs b/src/db/bin/convert.rs
index 7525bcc9..5dd9e034 100644
--- a/src/db/bin/convert.rs
+++ b/src/db/bin/convert.rs
@@ -48,6 +48,17 @@ fn open_db(path: PathBuf, engine: String) -> Result<Db> {
 			let db = sqlite_adapter::rusqlite::Connection::open(&path)?;
 			Ok(sqlite_adapter::SqliteDb::init(db))
 		}
+		"lmdb" | "heed" => {
+			std::fs::create_dir_all(&path).map_err(|e| {
+				Error(format!("Unable to create LMDB data directory: {}", e).into())
+			})?;
+			let db = lmdb_adapter::heed::EnvOpenOptions::new()
+				.max_dbs(100)
+				.map_size(1usize << 30)
+				.open(&path)
+				.unwrap();
+			Ok(lmdb_adapter::LmdbDb::init(db))
+		}
 		e => Err(Error(format!("Invalid DB engine: {}", e).into())),
 	}
 }
diff --git a/src/db/lmdb_adapter.rs b/src/db/lmdb_adapter.rs
index d568b5c5..095f512b 100644
--- a/src/db/lmdb_adapter.rs
+++ b/src/db/lmdb_adapter.rs
@@ -72,7 +72,31 @@ impl IDb for LmdbDb {
 	}
 
 	fn list_trees(&self) -> Result<Vec<String>> {
-		unimplemented!()
+		let tree0 = match self.db.open_database::<heed::types::Str, ByteSlice>(None)? {
+			Some(x) => x,
+			None => return Ok(vec![]),
+		};
+
+		let mut ret = vec![];
+		let tx = self.db.read_txn()?;
+		for item in tree0.iter(&tx)? {
+			let (tree_name, _) = item?;
+			ret.push(tree_name.to_string());
+		}
+		drop(tx);
+
+		let mut ret2 = vec![];
+		for tree_name in ret {
+			if self
+				.db
+				.open_database::<ByteSlice, ByteSlice>(Some(&tree_name))?
+				.is_some()
+			{
+				ret2.push(tree_name);
+			}
+		}
+
+		Ok(ret2)
 	}
 
 	// ----
diff --git a/src/garage/server.rs b/src/garage/server.rs
index a56f124a..9a1aa975 100644
--- a/src/garage/server.rs
+++ b/src/garage/server.rs
@@ -36,6 +36,7 @@ pub async fn run_server(config_file: PathBuf) -> Result<(), Error> {
 	let db = match config.db_engine.as_str() {
 		"sled" => {
 			db_path.push("db");
+			info!("Opening Sled database at: {}", db_path.display());
 			let db = db::sled_adapter::sled::Config::default()
 				.path(&db_path)
 				.cache_capacity(config.sled_cache_capacity)
@@ -46,13 +47,32 @@ pub async fn run_server(config_file: PathBuf) -> Result<(), Error> {
 		}
 		"sqlite" | "sqlite3" | "rusqlite" => {
 			db_path.push("db.sqlite");
+			info!("Opening Sqlite database at: {}", db_path.display());
 			let db = db::sqlite_adapter::rusqlite::Connection::open(db_path)
 				.expect("Unable to open sqlite DB");
 			db::sqlite_adapter::SqliteDb::init(db)
 		}
+		"lmdb" | "heed" => {
+			db_path.push("db.lmdb");
+			info!("Opening LMDB database at: {}", db_path.display());
+			std::fs::create_dir_all(&db_path).expect("Unable to create LMDB data directory");
+			let map_size = if u32::MAX as usize == usize::MAX {
+				warn!("LMDB is not recommended on 32-bit systems, database size will be limited");
+				1usize << 30 // 1GB for 32-bit systems
+			} else {
+				1usize << 40 // 1TB for 64-bit systems
+			};
+
+			let db = db::lmdb_adapter::heed::EnvOpenOptions::new()
+				.max_dbs(100)
+				.map_size(map_size)
+				.open(&db_path)
+				.expect("Unable to open LMDB DB");
+			db::lmdb_adapter::LmdbDb::init(db)
+		}
 		e => {
 			return Err(Error::Message(format!(
-				"Unsupported DB engine: {} (options: sled, sqlite)",
+				"Unsupported DB engine: {} (options: sled, sqlite, lmdb)",
 				e
 			)));
 		}
diff --git a/src/util/config.rs b/src/util/config.rs
index 3b37adbb..e8ef4fdd 100644
--- a/src/util/config.rs
+++ b/src/util/config.rs
@@ -65,7 +65,7 @@ pub struct Config {
 	pub kubernetes_skip_crd: bool,
 
 	// -- DB
-	/// Database engine to use for metadata (options: sled, sqlite)
+	/// Database engine to use for metadata (options: sled, sqlite, lmdb)
 	#[serde(default = "default_db_engine")]
 	pub db_engine: String,
 
-- 
cgit v1.2.3