use core::ops::Bound; use std::collections::HashMap; use std::convert::TryInto; use std::marker::PhantomPinned; use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::{Arc, RwLock}; use heed::types::Bytes; use heed::{BytesDecode, Env, EnvFlags, RoTxn, RwTxn, WithTls}; type Database = heed::Database; use crate::{ open::{Engine, OpenOpt}, Db, DbError, DbResult, Error, IDb, ITx, ITxFn, OnCommit, TxError, TxFnResult, TxOpError, TxResult, TxValueIter, Value, ValueIter, }; pub use heed; // ---- top-level open function pub(crate) fn open_db(path: &PathBuf, opt: &OpenOpt) -> DbResult { info!("Opening LMDB database at: {}", path.display()); if let Err(e) = std::fs::create_dir_all(path) { return Err(DbError( format!("Unable to create LMDB data directory: {}", e).into(), )); } let map_size = match opt.lmdb_map_size { None => recommended_map_size(), Some(v) => v - (v % 4096), }; let mut env_builder = heed::EnvOpenOptions::new(); env_builder.max_dbs(100); env_builder.map_size(map_size); env_builder.max_readers(2048); let mut env_flags = EnvFlags::NO_READ_AHEAD | EnvFlags::NO_META_SYNC; if !opt.fsync { env_flags |= EnvFlags::NO_SYNC; } let open_res = unsafe { env_builder.flags(env_flags); env_builder.open(path) }; match open_res { Err(heed::Error::Io(e)) if e.kind() == std::io::ErrorKind::OutOfMemory => Err(DbError( "OutOfMemory error while trying to open LMDB database. This can happen \ if your operating system is not allowing you to use sufficient virtual \ memory address space. Please check that no limit is set (ulimit -v). \ You may also try to set a smaller `lmdb_map_size` configuration parameter. \ On 32-bit machines, you should probably switch to another database engine." .into(), )), Err(e) => Err(DbError(format!("Cannot open LMDB database: {}", e).into())), Ok(db) => Ok(LmdbDb::init(db)), } } // -- err impl From for DbError { fn from(e: heed::Error) -> DbError { DbError(format!("LMDB: {}", e).into()) } } impl From for Error { fn from(e: heed::Error) -> Error { DbError::from(e).into() } } impl From for TxOpError { fn from(e: heed::Error) -> TxOpError { DbError::from(e).into() } } // -- db pub struct LmdbDb { db: heed::Env, trees: RwLock<(Vec, HashMap)>, } impl LmdbDb { pub fn init(db: Env) -> Db { let s = Self { db, trees: RwLock::new((Vec::new(), HashMap::new())), }; Db(Arc::new(s)) } fn get_tree(&self, i: usize) -> DbResult { self.trees .read() .unwrap() .0 .get(i) .cloned() .ok_or_else(|| DbError("invalid tree id".into())) } } impl IDb for LmdbDb { fn engine(&self) -> String { "LMDB (using Heed crate)".into() } fn open_tree(&self, name: &str) -> DbResult { let mut trees = self.trees.write().unwrap(); if let Some(i) = trees.1.get(name) { Ok(*i) } else { let mut wtxn = self.db.write_txn()?; let tree = self.db.create_database(&mut wtxn, Some(name))?; wtxn.commit()?; let i = trees.0.len(); trees.0.push(tree); trees.1.insert(name.to_string(), i); Ok(i) } } fn list_trees(&self) -> DbResult> { let rtxn = self.db.read_txn()?; let tree0 = match self .db .open_database::(&rtxn, None)? { Some(x) => x, None => return Ok(vec![]), }; let mut ret = vec![]; for item in tree0.iter(&rtxn)? { let (tree_name, _) = item?; ret.push(tree_name.to_string()); } let mut ret2 = vec![]; for tree_name in ret { if self .db .open_database::(&rtxn, Some(&tree_name))? .is_some() { ret2.push(tree_name); } } drop(rtxn); Ok(ret2) } fn snapshot(&self, base_path: &Path) -> DbResult<()> { std::fs::create_dir_all(base_path)?; let path = Engine::Lmdb.db_path(base_path); self.db .copy_to_path(path, heed::CompactionOption::Enabled)?; Ok(()) } // ---- fn get(&self, tree: usize, key: &[u8]) -> DbResult> { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; let val = tree.get(&tx, key)?; match val { None => Ok(None), Some(v) => Ok(Some(v.to_vec())), } } fn approximate_len(&self, tree: usize) -> DbResult { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; Ok(tree.len(&tx)?.try_into().unwrap()) } fn is_empty(&self, tree: usize) -> DbResult { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; Ok(tree.is_empty(&tx)?) } fn insert(&self, tree: usize, key: &[u8], value: &[u8]) -> DbResult<()> { let tree = self.get_tree(tree)?; let mut tx = self.db.write_txn()?; tree.put(&mut tx, key, value)?; tx.commit()?; Ok(()) } fn remove(&self, tree: usize, key: &[u8]) -> DbResult<()> { let tree = self.get_tree(tree)?; let mut tx = self.db.write_txn()?; tree.delete(&mut tx, key)?; tx.commit()?; Ok(()) } fn clear(&self, tree: usize) -> DbResult<()> { let tree = self.get_tree(tree)?; let mut tx = self.db.write_txn()?; tree.clear(&mut tx)?; tx.commit()?; Ok(()) } fn iter(&self, tree: usize) -> DbResult> { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; // Safety: the cloture does not store its argument anywhere, unsafe { TxAndIterator::make(tx, |tx| Ok(tree.iter(tx)?)) } } fn iter_rev(&self, tree: usize) -> DbResult> { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; // Safety: the cloture does not store its argument anywhere, unsafe { TxAndIterator::make(tx, |tx| Ok(tree.rev_iter(tx)?)) } } fn range<'r>( &self, tree: usize, low: Bound<&'r [u8]>, high: Bound<&'r [u8]>, ) -> DbResult> { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; // Safety: the cloture does not store its argument anywhere, unsafe { TxAndIterator::make(tx, |tx| Ok(tree.range(tx, &(low, high))?)) } } fn range_rev<'r>( &self, tree: usize, low: Bound<&'r [u8]>, high: Bound<&'r [u8]>, ) -> DbResult> { let tree = self.get_tree(tree)?; let tx = self.db.read_txn()?; // Safety: the cloture does not store its argument anywhere, unsafe { TxAndIterator::make(tx, |tx| Ok(tree.rev_range(tx, &(low, high))?)) } } // ---- fn transaction(&self, f: &dyn ITxFn) -> TxResult { let trees = self.trees.read().unwrap(); let mut tx = LmdbTx { trees: &trees.0[..], tx: self .db .write_txn() .map_err(Error::from) .map_err(TxError::Db)?, }; let res = f.try_on(&mut tx); match res { TxFnResult::Ok(on_commit) => { tx.tx.commit().map_err(Error::from).map_err(TxError::Db)?; Ok(on_commit) } TxFnResult::Abort => { tx.tx.abort(); Err(TxError::Abort(())) } TxFnResult::DbErr => { tx.tx.abort(); Err(TxError::Db( DbError("(this message will be discarded)".into()).into(), )) } } } } // ---- struct LmdbTx<'a> { trees: &'a [Database], tx: RwTxn<'a>, } impl<'a> LmdbTx<'a> { fn get_tree(&self, i: usize) -> DbResult<&Database> { self.trees.get(i).ok_or_else(|| { DbError( "invalid tree id (it might have been opened after the transaction started)".into(), ) }) } } impl<'a> ITx for LmdbTx<'a> { fn get(&self, tree: usize, key: &[u8]) -> DbResult> { let tree = self.get_tree(tree)?; match tree.get(&self.tx, key)? { Some(v) => Ok(Some(v.to_vec())), None => Ok(None), } } fn len(&self, tree: usize) -> DbResult { let tree = self.get_tree(tree)?; Ok(tree.len(&self.tx)? as usize) } fn insert(&mut self, tree: usize, key: &[u8], value: &[u8]) -> DbResult<()> { let tree = *self.get_tree(tree)?; tree.put(&mut self.tx, key, value)?; Ok(()) } fn remove(&mut self, tree: usize, key: &[u8]) -> DbResult<()> { let tree = *self.get_tree(tree)?; tree.delete(&mut self.tx, key)?; Ok(()) } fn clear(&mut self, tree: usize) -> DbResult<()> { let tree = *self.get_tree(tree)?; tree.clear(&mut self.tx)?; Ok(()) } fn iter(&self, tree: usize) -> DbResult> { let tree = *self.get_tree(tree)?; Ok(Box::new(tree.iter(&self.tx)?.map(tx_iter_item))) } fn iter_rev(&self, tree: usize) -> DbResult> { let tree = *self.get_tree(tree)?; Ok(Box::new(tree.rev_iter(&self.tx)?.map(tx_iter_item))) } fn range<'r>( &self, tree: usize, low: Bound<&'r [u8]>, high: Bound<&'r [u8]>, ) -> DbResult> { let tree = *self.get_tree(tree)?; Ok(Box::new( tree.range(&self.tx, &(low, high))?.map(tx_iter_item), )) } fn range_rev<'r>( &self, tree: usize, low: Bound<&'r [u8]>, high: Bound<&'r [u8]>, ) -> DbResult> { let tree = *self.get_tree(tree)?; Ok(Box::new( tree.rev_range(&self.tx, &(low, high))?.map(tx_iter_item), )) } } // ---- iterators outside transactions ---- // complicated, they must hold the transaction object // therefore a bit of unsafe code (it is a self-referential struct) type IteratorItem<'a> = heed::Result<( >::DItem, >::DItem, )>; struct TxAndIterator<'a, I> where I: Iterator> + 'a, { tx: RoTxn<'a, WithTls>, iter: Option, _pin: PhantomPinned, } impl<'a, I> TxAndIterator<'a, I> where I: Iterator> + 'a, { fn iter(self: Pin<&mut Self>) -> &mut Option { // Safety: iter is not structural unsafe { &mut self.get_unchecked_mut().iter } } /// Safety: iterfun must not store its argument anywhere but in its result. unsafe fn make(tx: RoTxn<'a, WithTls>, iterfun: F) -> DbResult> where F: FnOnce(&'a RoTxn<'a>) -> DbResult, { let res = TxAndIterator { tx, iter: None, _pin: PhantomPinned, }; let mut boxed = Box::pin(res); let tx_lifetime_overextended: &'a RoTxn<'a> = { let tx = &boxed.tx; // Safety: Artificially extending the lifetime because // this reference will only be stored and accessed from the // returned ValueIter which guarantees that it is destroyed // before the tx it is pointing to. #[expect(clippy::deref_addrof)] unsafe { &*&raw const *tx } }; let iter = iterfun(tx_lifetime_overextended)?; *boxed.as_mut().iter() = Some(iter); Ok(Box::new(TxAndIteratorPin(boxed))) } } impl<'a, I> Drop for TxAndIterator<'a, I> where I: Iterator> + 'a, { fn drop(&mut self) { // Safety: `new_unchecked` is okay because we know this value is never // used again after being dropped. let this = unsafe { Pin::new_unchecked(self) }; drop(this.iter().take()); } } struct TxAndIteratorPin<'a, I>(Pin>>) where I: Iterator> + 'a; impl<'a, I> Iterator for TxAndIteratorPin<'a, I> where I: Iterator> + 'a, { type Item = DbResult<(Value, Value)>; fn next(&mut self) -> Option { let mut_ref = Pin::as_mut(&mut self.0); let next = mut_ref.iter().as_mut()?.next()?; let res = match next { Err(e) => Err(DbError::from(e)), Ok((k, v)) => Ok((k.to_vec(), v.to_vec())), }; Some(res) } } // ---- iterators within transactions ---- fn tx_iter_item<'a>( item: std::result::Result<(&'a [u8], &'a [u8]), heed::Error>, ) -> DbResult<(Vec, Vec)> { item.map(|(k, v)| (k.to_vec(), v.to_vec())) .map_err(DbError::from) } // ---- utility ---- #[cfg(target_pointer_width = "64")] pub fn recommended_map_size() -> usize { 1usize << 40 } #[cfg(target_pointer_width = "32")] pub fn recommended_map_size() -> usize { tracing::warn!("LMDB is not recommended on 32-bit systems, database size will be limited"); 1usize << 30 }