// Copyright 2018-2019 Mozilla // // Licensed under the Apache License, Version 2.0 (the "License"); you may not use // this file except in compliance with the License. You may obtain a copy of the // License at http://www.apache.org/licenses/LICENSE-2.0 // Unless required by applicable law or agreed to in writing, software distributed // under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR // CONDITIONS OF ANY KIND, either express or implied. See the License for the // specific language governing permissions and limitations under the License. use std::{ fs, path::{Path, PathBuf}, }; use lmdb::Error as LmdbError; use super::{ DatabaseFlagsImpl, DatabaseImpl, EnvironmentFlagsImpl, ErrorImpl, InfoImpl, RoTransactionImpl, RwTransactionImpl, StatImpl, }; use crate::backend::traits::{ BackendEnvironment, BackendEnvironmentBuilder, BackendInfo, BackendIter, BackendRoCursor, BackendRoCursorTransaction, BackendStat, }; #[derive(Debug, PartialEq, Eq, Copy, Clone)] pub struct EnvironmentBuilderImpl { builder: lmdb::EnvironmentBuilder, env_path_type: EnvironmentPathType, env_lock_type: EnvironmentLockType, env_db_type: EnvironmentDefaultDbType, make_dir_if_needed: bool, } impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl { type Environment = EnvironmentImpl; type Error = ErrorImpl; type Flags = EnvironmentFlagsImpl; fn new() -> EnvironmentBuilderImpl { EnvironmentBuilderImpl { builder: lmdb::Environment::new(), env_path_type: EnvironmentPathType::SubDir, env_lock_type: EnvironmentLockType::Lockfile, env_db_type: EnvironmentDefaultDbType::SingleDatabase, make_dir_if_needed: false, } } fn set_flags(&mut self, flags: T) -> &mut Self where T: Into, { let flags = flags.into(); if flags.0 == lmdb::EnvironmentFlags::NO_SUB_DIR { self.env_path_type = EnvironmentPathType::NoSubDir; } if flags.0 == lmdb::EnvironmentFlags::NO_LOCK { self.env_lock_type = EnvironmentLockType::NoLockfile; } self.builder.set_flags(flags.0); self } fn set_max_readers(&mut self, max_readers: u32) -> &mut Self { self.builder.set_max_readers(max_readers); self } fn set_max_dbs(&mut self, max_dbs: u32) -> &mut Self { if max_dbs > 0 { self.env_db_type = EnvironmentDefaultDbType::MultipleNamedDatabases } self.builder.set_max_dbs(max_dbs); self } fn set_map_size(&mut self, size: usize) -> &mut Self { self.builder.set_map_size(size); self } fn set_make_dir_if_needed(&mut self, make_dir_if_needed: bool) -> &mut Self { self.make_dir_if_needed = make_dir_if_needed; self } fn set_discard_if_corrupted(&mut self, _discard_if_corrupted: bool) -> &mut Self { // Unfortunately, when opening a database, LMDB doesn't handle all the ways it could have // been corrupted. Prefer using the `SafeMode` backend if this is important. unimplemented!(); } fn open(&self, path: &Path) -> Result { match self.env_path_type { EnvironmentPathType::NoSubDir => { if !path.is_file() { return Err(ErrorImpl::UnsuitableEnvironmentPath(path.into())); } } EnvironmentPathType::SubDir => { if !path.is_dir() { if !self.make_dir_if_needed { return Err(ErrorImpl::UnsuitableEnvironmentPath(path.into())); } fs::create_dir_all(path)?; } } } self.builder .open(path) .map_err(ErrorImpl::LmdbError) .and_then(|lmdbenv| { EnvironmentImpl::new( path, self.env_path_type, self.env_lock_type, self.env_db_type, lmdbenv, ) }) } } #[derive(Debug, PartialEq, Eq, Copy, Clone)] pub enum EnvironmentPathType { SubDir, NoSubDir, } #[derive(Debug, PartialEq, Eq, Copy, Clone)] pub enum EnvironmentLockType { Lockfile, NoLockfile, } #[derive(Debug, PartialEq, Eq, Copy, Clone)] pub enum EnvironmentDefaultDbType { SingleDatabase, MultipleNamedDatabases, } #[derive(Debug)] pub struct EnvironmentImpl { path: PathBuf, env_path_type: EnvironmentPathType, env_lock_type: EnvironmentLockType, env_db_type: EnvironmentDefaultDbType, lmdbenv: lmdb::Environment, } impl EnvironmentImpl { pub(crate) fn new( path: &Path, env_path_type: EnvironmentPathType, env_lock_type: EnvironmentLockType, env_db_type: EnvironmentDefaultDbType, lmdbenv: lmdb::Environment, ) -> Result { Ok(EnvironmentImpl { path: path.to_path_buf(), env_path_type, env_lock_type, env_db_type, lmdbenv, }) } } impl<'e> BackendEnvironment<'e> for EnvironmentImpl { type Database = DatabaseImpl; type Error = ErrorImpl; type Flags = DatabaseFlagsImpl; type Info = InfoImpl; type RoTransaction = RoTransactionImpl<'e>; type RwTransaction = RwTransactionImpl<'e>; type Stat = StatImpl; fn get_dbs(&self) -> Result>, Self::Error> { if self.env_db_type == EnvironmentDefaultDbType::SingleDatabase { return Ok(vec![None]); } let db = self .lmdbenv .open_db(None) .map(DatabaseImpl) .map_err(ErrorImpl::LmdbError)?; let reader = self.begin_ro_txn()?; let cursor = reader.open_ro_cursor(&db)?; let mut iter = cursor.into_iter(); let mut store = vec![]; while let Some(result) = iter.next() { let (key, _) = result?; let name = String::from_utf8(key.to_owned()) .map_err(|_| ErrorImpl::LmdbError(lmdb::Error::Corrupted))?; store.push(Some(name)); } Ok(store) } fn open_db(&self, name: Option<&str>) -> Result { self.lmdbenv .open_db(name) .map(DatabaseImpl) .map_err(ErrorImpl::LmdbError) } fn create_db( &self, name: Option<&str>, flags: Self::Flags, ) -> Result { self.lmdbenv .create_db(name, flags.0) .map(DatabaseImpl) .map_err(ErrorImpl::LmdbError) } fn begin_ro_txn(&'e self) -> Result { self.lmdbenv .begin_ro_txn() .map(RoTransactionImpl) .map_err(ErrorImpl::LmdbError) } fn begin_rw_txn(&'e self) -> Result { self.lmdbenv .begin_rw_txn() .map(RwTransactionImpl) .map_err(ErrorImpl::LmdbError) } fn sync(&self, force: bool) -> Result<(), Self::Error> { self.lmdbenv.sync(force).map_err(ErrorImpl::LmdbError) } fn stat(&self) -> Result { self.lmdbenv .stat() .map(StatImpl) .map_err(ErrorImpl::LmdbError) } fn info(&self) -> Result { self.lmdbenv .info() .map(InfoImpl) .map_err(ErrorImpl::LmdbError) } fn freelist(&self) -> Result { self.lmdbenv.freelist().map_err(ErrorImpl::LmdbError) } fn load_ratio(&self) -> Result, Self::Error> { let stat = self.stat()?; let info = self.info()?; let freelist = self.freelist()?; let last_pgno = info.last_pgno() + 1; // pgno is 0 based. let total_pgs = info.map_size() / stat.page_size(); if freelist > last_pgno { return Err(ErrorImpl::LmdbError(LmdbError::Corrupted)); } let used_pgs = last_pgno - freelist; Ok(Some(used_pgs as f32 / total_pgs as f32)) } fn set_map_size(&self, size: usize) -> Result<(), Self::Error> { self.lmdbenv .set_map_size(size) .map_err(ErrorImpl::LmdbError) } fn get_files_on_disk(&self) -> Vec { let mut store = vec![]; if self.env_path_type == EnvironmentPathType::NoSubDir { // The option NO_SUB_DIR could change the default directory layout; therefore this should // probably return the path used to create environment, along with the custom lockfile // when available. unimplemented!(); } let mut db_filename = self.path.clone(); db_filename.push("data.mdb"); store.push(db_filename); if self.env_lock_type == EnvironmentLockType::Lockfile { let mut lock_filename = self.path.clone(); lock_filename.push("lock.mdb"); store.push(lock_filename); } store } }