test: test

Author: ardocrat

store/src/lmdb.rs

@@ -15,7 +15,9 @@
 //! Storage of core types using LMDB.
 
 use heed::types::Bytes;
-use heed::{Database, Env, EnvOpenOptions, RoTxn, RwTxn, WithoutTls};
+use heed::{Database, Env, EnvOpenOptions, RoTxn, RwTxn, WithTls, WithoutTls};
+use lazy_static::lazy_static;
+use std::path::Path;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
@@ -26,14 +28,19 @@ use crate::grin_core::ser::{self, DeserializationMode, ProtocolVersion};
 use crate::util::RwLock;
 
 /// number of bytes to grow the database by when needed
-pub const ALLOC_CHUNK_SIZE_DEFAULT: usize = 134_217_728 / 32; //128 MB
+pub const ALLOC_CHUNK_SIZE_DEFAULT: usize = 134_217_728; //128 MB
 /// And for test mode, to avoid too much disk allocation on windows
 pub const ALLOC_CHUNK_SIZE_DEFAULT_TEST: usize = 1_048_576; //1 MB
 const RESIZE_PERCENT: f32 = 0.9;
 /// Want to ensure that each resize gives us at least this %
 /// of total space free
 const RESIZE_MIN_TARGET_PERCENT: f32 = 0.65;
 
+lazy_static! {
+	/// Check if environment is resizing.
+	static ref RESIZING: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
+}
+
 /// Main error type for this lmdb
 #[derive(Clone, Eq, PartialEq, Debug, thiserror::Error)]
 pub enum Error {
@@ -80,51 +87,54 @@ where
 
 const DEFAULT_DB_VERSION: ProtocolVersion = ProtocolVersion(3);
 
+const ENV_NAME: &'static str = "lmdb";
+
 /// LMDB-backed store facilitating data access and serialization. All writes
 /// are done through a Batch abstraction providing atomicity.
 pub struct Store {
-	env: Arc<Env<WithoutTls>>,
 	db: Arc<RwLock<Option<Arc<Database<Bytes, Bytes>>>>>,
+	env: Arc<Env<WithTls>>,
 	name: String,
 	version: ProtocolVersion,
 	alloc_chunk_size: usize,
-	resizing: Arc<AtomicBool>,
 }
 
 impl Store {
 	/// Create a new LMDB env under the provided directory.
 	/// By default creates an environment named "lmdb".
 	/// Be aware of transactional semantics in lmdb
 	/// (transactions are per environment, not per database).
+	/// db with non-default `env_name` will be migrated into default environment.
 	pub fn new(
 		root_path: &str,
 		env_name: Option<&str>,
 		db_name: Option<&str>,
 		max_readers: Option<u32>,
 	) -> Result<Store, Error> {
-		let name = match env_name {
-			Some(n) => n.to_owned(),
-			None => "lmdb".to_owned(),
-		};
-		let db_name = match db_name {
-			Some(n) => n.to_owned(),
-			None => "lmdb".to_owned(),
-		};
-		let full_path = [root_path.to_owned(), name].join("/");
+		let name = env_name.unwrap_or_else(|| ENV_NAME);
+		let db_name = db_name.unwrap_or_else(|| ENV_NAME);
+		let mut full_path = Path::new(root_path).join(name);
+
+		// Fix path to use default environment.
+		if name != ENV_NAME {
+			full_path = Path::new(root_path).join(ENV_NAME);
+		}
+
 		fs::create_dir_all(&full_path).map_err(|e| {
 			Error::FileErr(format!(
-				"Unable to create directory 'db_root' to store chain_data: {:?}",
-				e
+				"Unable to create {:?} to store data: {:?}",
+				full_path, e
 			))
 		})?;
 
 		let alloc_chunk_size = match global::is_production_mode() {
 			true => ALLOC_CHUNK_SIZE_DEFAULT,
-			false => ALLOC_CHUNK_SIZE_DEFAULT_TEST,
+			false => ALLOC_CHUNK_SIZE_DEFAULT,
 		};
 
+		// Environment setup.
 		let env = unsafe {
-			let mut options = EnvOpenOptions::new().read_txn_without_tls();
+			let mut options = EnvOpenOptions::new();
 			let mut env_options = options.map_size(alloc_chunk_size).max_dbs(8);
 			if let Some(max_readers) = max_readers {
 				env_options = env_options.max_readers(max_readers);
@@ -137,13 +147,15 @@ impl Store {
 				env.resize(new_size)?;
 			};
 		}
+		debug!("DB Mapsize for {:?} is {}", full_path, env.info().map_size);
+
+		// Database setup.
 		let s = Store {
 			env: Arc::new(env),
 			db: Arc::new(RwLock::new(None)),
-			name: db_name,
+			name: db_name.to_string(),
 			version: DEFAULT_DB_VERSION,
 			alloc_chunk_size,
-			resizing: Default::default(),
 		};
 		{
 			let mut write = s.env.write_txn()?;
@@ -152,9 +164,46 @@ impl Store {
 			let mut w_db = s.db.write();
 			*w_db = Some(Arc::new(db));
 		}
+
+		// Migrate to default environment if needed.
+		let migrate_from = Path::new(root_path).join(name);
+		if name != ENV_NAME && migrate_from.exists() {
+			debug!("Migrating {} to {:?}", name, full_path);
+			if s.migrate_to_default_env(&migrate_from, db_name).is_ok() {
+				let _ = fs::remove_dir_all(&migrate_from);
+			} else {
+				error!("Migrating {} failed", name);
+			}
+		}
 		Ok(s)
+	}
 
-		// debug!("DB Mapsize for {} is {}", full_path, env.info().map_size);
+	/// Migrate db from provided path to store environment.
+	fn migrate_to_default_env(&self, path: &Path, db_name: &str) -> Result<(), Error> {
+		let env = unsafe {
+			let mut options = EnvOpenOptions::new().read_txn_without_tls();
+			let env_options = options.map_size(self.alloc_chunk_size).max_dbs(1);
+			env_options.open(path)?
+		};
+		let db_from = {
+			let mut write = env.write_txn()?;
+			let db: Database<Bytes, Bytes> = env.create_database(&mut write, Some(db_name))?;
+			write.commit()?;
+			db
+		};
+		let db_to = self.db.read();
+		let mut write_to = self.env.write_txn()?;
+		let read_from = env.read_txn()?;
+		let mut count = 0;
+		for kv in db_from.iter(&read_from)? {
+			count += 1;
+			if let Ok((k, v)) = kv {
+				db_to.as_ref().unwrap().put(&mut write_to, &k, &v)?;
+			}
+		}
+		write_to.commit()?;
+		debug!("Migrated {} records of {}", count, db_name);
+		Ok(())
 	}
 
 	/// Construct a new store using a specific protocol version.
@@ -166,7 +215,6 @@ impl Store {
 			name: self.name.clone(),
 			version,
 			alloc_chunk_size: self.alloc_chunk_size,
-			resizing: self.resizing.clone(),
 		}
 	}
 
@@ -176,7 +224,7 @@ impl Store {
 	}
 
 	/// Determines whether the environment needs a resize based on a simple percentage threshold.
-	pub fn needs_resize(env: Arc<Env<WithoutTls>>, alloc_chunk_size: usize) -> (bool, usize) {
+	pub fn needs_resize(env: Arc<Env<WithTls>>, alloc_chunk_size: usize) -> (bool, usize) {
 		let env_info = env.info();
 		let stat = env.stat();
 		let size_used = stat.page_size as usize * env_info.last_page_number;
@@ -274,19 +322,19 @@ impl Store {
 
 	/// Builds a new batch to be used with this store.
 	pub fn batch(&self) -> Result<Batch<'_>, Error> {
-		while self.resizing.load(Ordering::SeqCst) {
+		while RESIZING.load(Ordering::SeqCst) {
 			debug!("Wait resizing {} DB", self.name);
 			thread::sleep(Duration::from_millis(100));
 		}
 		let (resize, new_size) = Self::needs_resize(self.env.clone(), self.alloc_chunk_size);
 		if resize {
-			self.resizing.store(true, Ordering::SeqCst);
+			RESIZING.store(true, Ordering::SeqCst);
 			debug!("Start resizing {} DB", self.name);
 			unsafe {
 				thread::sleep(Duration::from_millis(3000));
 				self.env.resize(new_size)?;
 			}
-			self.resizing.store(false, Ordering::SeqCst);
+			RESIZING.store(false, Ordering::SeqCst);
 			debug!("End resizing {} DB", self.name);
 		}
 		Ok(Batch::new(self)?)
@@ -419,7 +467,7 @@ where
 	F: Fn(&[u8], &[u8]) -> Result<T, Error>,
 {
 	db: Arc<Option<Arc<Database<Bytes, Bytes>>>>,
-	read: Arc<RoTxn<'a, WithoutTls>>,
+	read: Arc<RoTxn<'a, WithTls>>,
 	skip: usize,
 	prefix: Vec<u8>,
 	deserialize: F,
@@ -464,7 +512,7 @@ where
 	/// Initialize a new prefix iterator.
 	pub fn new(
 		db: Arc<Option<Arc<Database<Bytes, Bytes>>>>,
-		read: RoTxn<'a, WithoutTls>,
+		read: RoTxn<'a, WithTls>,
 		prefix: &[u8],
 		deserialize: F,
 	) -> PrefixIterator<'a, F, T> {