Merge pull request #206 from mozilla/migration-bug

Migration updates

src/backend/impl_safe/environment.rs

@@ -12,6 +12,7 @@ use std::{
     borrow::Cow,
     collections::HashMap,
     fs,
+    ops::DerefMut,
     path::{
         Path,
         PathBuf,
@@ -115,32 +116,51 @@ impl<'b> BackendEnvironmentBuilder<'b> for EnvironmentBuilderImpl {
     }
 }
 
+#[derive(Debug)]
+pub(crate) struct EnvironmentDbs {
+    pub(crate) arena: DatabaseArena,
+    pub(crate) name_map: DatabaseNameMap,
+}
+
+#[derive(Debug)]
+pub(crate) struct EnvironmentDbsRefMut<'a> {
+    pub(crate) arena: &'a mut DatabaseArena,
+    pub(crate) name_map: &'a mut DatabaseNameMap,
+}
+
+impl<'a> From<&'a mut EnvironmentDbs> for EnvironmentDbsRefMut<'a> {
+    fn from(dbs: &mut EnvironmentDbs) -> EnvironmentDbsRefMut {
+        EnvironmentDbsRefMut {
+            arena: &mut dbs.arena,
+            name_map: &mut dbs.name_map,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct EnvironmentImpl {
     path: PathBuf,
     max_dbs: usize,
-    arena: RwLock<DatabaseArena>,
-    dbs: RwLock<DatabaseNameMap>,
+    dbs: RwLock<EnvironmentDbs>,
     ro_txns: Arc<()>,
     rw_txns: Arc<()>,
 }
 
 impl EnvironmentImpl {
     fn serialize(&self) -> Result<Vec<u8>, ErrorImpl> {
-        let arena = self.arena.read().map_err(|_| ErrorImpl::EnvPoisonError)?;
         let dbs = self.dbs.read().map_err(|_| ErrorImpl::EnvPoisonError)?;
-        let data: HashMap<_, _> = dbs.iter().map(|(name, id)| (name, &arena[id.0])).collect();
+        let data: HashMap<_, _> = dbs.name_map.iter().map(|(name, id)| (name, &dbs.arena[id.0])).collect();
         Ok(bincode::serialize(&data)?)
     }
 
     fn deserialize(bytes: &[u8]) -> Result<(DatabaseArena, DatabaseNameMap), ErrorImpl> {
         let mut arena = DatabaseArena::new();
-        let mut dbs = HashMap::new();
+        let mut name_map = HashMap::new();
         let data: HashMap<_, _> = bincode::deserialize(&bytes)?;
         for (name, db) in data {
-            dbs.insert(name, DatabaseImpl(arena.alloc(db)));
+            name_map.insert(name, DatabaseImpl(arena.alloc(db)));
         }
-        Ok((arena, dbs))
+        Ok((arena, name_map))
     }
 }
 
@@ -165,8 +185,10 @@ impl EnvironmentImpl {
         Ok(EnvironmentImpl {
             path: path.to_path_buf(),
             max_dbs: max_dbs.unwrap_or(std::usize::MAX),
-            arena: RwLock::new(DatabaseArena::new()),
-            dbs: RwLock::new(HashMap::new()),
+            dbs: RwLock::new(EnvironmentDbs {
+                arena: DatabaseArena::new(),
+                name_map: HashMap::new(),
+            }),
             ro_txns: Arc::new(()),
             rw_txns: Arc::new(()),
         })
@@ -180,9 +202,11 @@ impl EnvironmentImpl {
         if fs::metadata(&path).is_err() {
             return Ok(());
         };
-        let (arena, dbs) = Self::deserialize(&fs::read(&path)?)?;
-        self.arena = RwLock::new(arena);
-        self.dbs = RwLock::new(dbs);
+        let (arena, name_map) = Self::deserialize(&fs::read(&path)?)?;
+        self.dbs = RwLock::new(EnvironmentDbs {
+            arena,
+            name_map,
+        });
         Ok(())
     }
 
@@ -195,12 +219,12 @@ impl EnvironmentImpl {
         Ok(())
     }
 
-    pub(crate) fn dbs(&self) -> Result<RwLockReadGuard<DatabaseArena>, ErrorImpl> {
-        self.arena.read().map_err(|_| ErrorImpl::EnvPoisonError)
+    pub(crate) fn dbs(&self) -> Result<RwLockReadGuard<EnvironmentDbs>, ErrorImpl> {
+        self.dbs.read().map_err(|_| ErrorImpl::EnvPoisonError)
     }
 
-    pub(crate) fn dbs_mut(&self) -> Result<RwLockWriteGuard<DatabaseArena>, ErrorImpl> {
-        self.arena.write().map_err(|_| ErrorImpl::EnvPoisonError)
+    pub(crate) fn dbs_mut(&self) -> Result<RwLockWriteGuard<EnvironmentDbs>, ErrorImpl> {
+        self.dbs.write().map_err(|_| ErrorImpl::EnvPoisonError)
     }
 }
 
@@ -215,7 +239,7 @@ impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
 
     fn get_dbs(&self) -> Result<Vec<Option<String>>, Self::Error> {
         let dbs = self.dbs.read().map_err(|_| ErrorImpl::EnvPoisonError)?;
-        Ok(dbs.keys().map(|key| key.to_owned()).collect())
+        Ok(dbs.name_map.keys().map(|key| key.to_owned()).collect())
     }
 
     fn open_db(&self, name: Option<&str>) -> Result<Self::Database, Self::Error> {
@@ -225,8 +249,8 @@ impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
         // TOOD: don't reallocate `name`.
         let key = name.map(String::from);
         let dbs = self.dbs.read().map_err(|_| ErrorImpl::EnvPoisonError)?;
-        let id = dbs.get(&key).ok_or(ErrorImpl::DbNotFoundError)?;
-        Ok(*id)
+        let db = dbs.name_map.get(&key).ok_or(ErrorImpl::DbNotFoundError)?;
+        Ok(*db)
     }
 
     fn create_db(&self, name: Option<&str>, flags: Self::Flags) -> Result<Self::Database, Self::Error> {
@@ -236,11 +260,13 @@ impl<'e> BackendEnvironment<'e> for EnvironmentImpl {
         // TOOD: don't reallocate `name`.
         let key = name.map(String::from);
         let mut dbs = self.dbs.write().map_err(|_| ErrorImpl::EnvPoisonError)?;
-        let mut arena = self.arena.write().map_err(|_| ErrorImpl::EnvPoisonError)?;
-        if dbs.keys().filter_map(|k| k.as_ref()).count() >= self.max_dbs && name != None {
+        if dbs.name_map.keys().filter_map(|k| k.as_ref()).count() >= self.max_dbs && name != None {
             return Err(ErrorImpl::DbsFull);
         }
-        let id = dbs.entry(key).or_insert_with(|| DatabaseImpl(arena.alloc(Database::new(Some(flags), None))));
+        let parts = EnvironmentDbsRefMut::from(dbs.deref_mut());
+        let arena = parts.arena;
+        let name_map = parts.name_map;
+        let id = name_map.entry(key).or_insert_with(|| DatabaseImpl(arena.alloc(Database::new(Some(flags), None))));
         Ok(*id)
     }
 

src/backend/impl_safe/transaction.rs

@@ -37,7 +37,7 @@ pub struct RoTransactionImpl<'t> {
 
 impl<'t> RoTransactionImpl<'t> {
     pub(crate) fn new(env: &'t EnvironmentImpl, idx: Arc<()>) -> Result<RoTransactionImpl<'t>, ErrorImpl> {
-        let snapshots = env.dbs()?.iter().map(|(id, db)| (DatabaseImpl(id), db.snapshot())).collect();
+        let snapshots = env.dbs()?.arena.iter().map(|(id, db)| (DatabaseImpl(id), db.snapshot())).collect();
         Ok(RoTransactionImpl {
             env,
             snapshots,
@@ -78,7 +78,7 @@ pub struct RwTransactionImpl<'t> {
 
 impl<'t> RwTransactionImpl<'t> {
     pub(crate) fn new(env: &'t EnvironmentImpl, idx: Arc<()>) -> Result<RwTransactionImpl<'t>, ErrorImpl> {
-        let snapshots = env.dbs()?.iter().map(|(id, db)| (DatabaseImpl(id), db.snapshot())).collect();
+        let snapshots = env.dbs()?.arena.iter().map(|(id, db)| (DatabaseImpl(id), db.snapshot())).collect();
         Ok(RwTransactionImpl {
             env,
             snapshots,
@@ -144,7 +144,7 @@ impl<'t> BackendRwTransaction for RwTransactionImpl<'t> {
         let mut dbs = self.env.dbs_mut()?;
 
         for (id, snapshot) in self.snapshots {
-            let db = dbs.get_mut(id.0).ok_or_else(|| ErrorImpl::DbIsForeignError)?;
+            let db = dbs.arena.get_mut(id.0).ok_or_else(|| ErrorImpl::DbIsForeignError)?;
             db.replace(snapshot);
         }
 

src/migrator.rs

@@ -28,8 +28,8 @@
 //!   handling all errors.
 //! * `easy_migrate_<src>_to_<dst>` which is similar to the above, but ignores the
 //!   migration and doesn't delete any files if the source environment is invalid
-//!   (corrupted), unavailable (path not found or incompatible with configuration), or
-//!   empty (database has no records).
+//!   (corrupted), unavailable (path not accessible or incompatible with configuration),
+//!   or empty (database has no records).
 //!
 //! The tool currently has these limitations:
 //!
@@ -100,9 +100,9 @@ macro_rules! fn_migrator {
     };
 
     (open $migrate:tt, $name:tt, $builder:tt, $src_env:ty, $dst_env:ty) => {
-        /// Same as the non `open_*` migration method, but automatically attempts to open the
-        /// source environment. Finally, deletes all of its supporting files if there's no other
-        /// environment open at that path.
+        /// Same as the the `migrate_x_to_y` migration method above, but automatically attempts
+        /// to open the source environment. Finally, deletes all of its supporting files if
+        /// there's no other environment open at that path and the migration succeeded.
         pub fn $name<F, D>(path: &std::path::Path, build: F, dst_env: D) -> Result<(), MigrateError>
         where
             F: FnOnce(crate::backend::$builder) -> crate::backend::$builder,
@@ -126,8 +126,13 @@ macro_rules! fn_migrator {
     };
 
     (easy $migrate:tt, $name:tt, $src_env:ty, $dst_env:ty) => {
-        /// Same as the `open_*` migration method, but ignores the migration and doesn't delete
-        /// any files if the source environment is invalid (corrupted), unavailable, or empty.
+        /// Same as the `open_and_migrate_x_to_y` migration method above, but ignores the
+        /// migration and doesn't delete any files if the following conditions apply:
+        /// - Source environment is invalid (corrupted), unavailable, or empty.
+        /// - Destination environment is not empty.
+        /// Use this instead of the other migration methods if:
+        /// - You're not concerned by throwing away old data and starting fresh with a new store.
+        /// - You'll never want to overwrite data in the new store from the old store.
         pub fn $name<D>(path: &std::path::Path, dst_env: D) -> Result<(), MigrateError>
         where
             D: std::ops::Deref<Target = Rkv<$dst_env>>,
@@ -137,6 +142,7 @@ macro_rules! fn_migrator {
                 Err(crate::MigrateError::StoreError(crate::StoreError::IoError(_))) => Ok(()),
                 Err(crate::MigrateError::StoreError(crate::StoreError::UnsuitableEnvironmentPath(_))) => Ok(()),
                 Err(crate::MigrateError::SourceEmpty) => Ok(()),
+                Err(crate::MigrateError::DestinationNotEmpty) => Ok(()),
                 result => result,
             }?;
 

tests/env-migration.rs

@@ -18,8 +18,11 @@ use tempfile::Builder;
 use rkv::{
     backend::{
         Lmdb,
+        LmdbEnvironment,
         SafeMode,
+        SafeModeEnvironment,
     },
+    Manager,
     Migrator,
     Rkv,
     StoreOptions,
@@ -38,8 +41,8 @@ macro_rules! populate_store {
 }
 
 #[test]
-fn test_simple_migrator_lmdb_to_safe() {
-    let root = Builder::new().prefix("test_simple_migrator_lmdb_to_safe").tempdir().expect("tempdir");
+fn test_open_migrator_lmdb_to_safe() {
+    let root = Builder::new().prefix("test_open_migrator_lmdb_to_safe").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     // Populate source environment and persist to disk.
@@ -92,8 +95,8 @@ fn test_simple_migrator_lmdb_to_safe() {
 }
 
 #[test]
-fn test_simple_migrator_safe_to_lmdb() {
-    let root = Builder::new().prefix("test_simple_migrator_safe_to_lmdb").tempdir().expect("tempdir");
+fn test_open_migrator_safe_to_lmdb() {
+    let root = Builder::new().prefix("test_open_migrator_safe_to_lmdb").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     // Populate source environment and persist to disk.
@@ -140,8 +143,8 @@ fn test_simple_migrator_safe_to_lmdb() {
 }
 
 #[test]
-fn test_migrator_round_trip() {
-    let root = Builder::new().prefix("test_simple_migrator_lmdb_to_safe").tempdir().expect("tempdir");
+fn test_open_migrator_round_trip() {
+    let root = Builder::new().prefix("test_open_migrator_lmdb_to_safe").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     // Populate source environment and persist to disk.
@@ -184,8 +187,8 @@ fn test_migrator_round_trip() {
 }
 
 #[test]
-fn test_migrator_no_dir_1() {
-    let root = Builder::new().prefix("test_migrator_no_dir").tempdir().expect("tempdir");
+fn test_easy_migrator_no_dir_1() {
+    let root = Builder::new().prefix("test_easy_migrator_no_dir").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     // This won't fail with IoError even though the path is a bogus path, because this
@@ -205,8 +208,8 @@ fn test_migrator_no_dir_1() {
 }
 
 #[test]
-fn test_migrator_no_dir_2() {
-    let root = Builder::new().prefix("test_migrator_no_dir").tempdir().expect("tempdir");
+fn test_easy_migrator_no_dir_2() {
+    let root = Builder::new().prefix("test_easy_migrator_no_dir").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     // This won't fail with IoError even though the path is a bogus path, because this
@@ -226,8 +229,8 @@ fn test_migrator_no_dir_2() {
 }
 
 #[test]
-fn test_migrator_invalid_1() {
-    let root = Builder::new().prefix("test_migrator_invalid").tempdir().expect("tempdir");
+fn test_easy_migrator_invalid_1() {
+    let root = Builder::new().prefix("test_easy_migrator_invalid").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     let dbfile = root.path().join("data.mdb");
@@ -250,8 +253,8 @@ fn test_migrator_invalid_1() {
 }
 
 #[test]
-fn test_migrator_invalid_2() {
-    let root = Builder::new().prefix("test_migrator_invalid").tempdir().expect("tempdir");
+fn test_easy_migrator_invalid_2() {
+    let root = Builder::new().prefix("test_easy_migrator_invalid").tempdir().expect("tempdir");
     fs::create_dir_all(root.path()).expect("dir created");
 
     let dbfile = root.path().join("data.safe.bin");
@@ -354,3 +357,115 @@ fn test_migrator_safe_to_lmdb_3() {
     assert_eq!(store.get(&reader, "bar").expect("read"), Some(Value::Bool(true)));
     assert_eq!(store.get(&reader, "baz").expect("read"), Some(Value::Str("héllo, yöu")));
 }
+
+#[test]
+fn test_easy_migrator_failed_migration_1() {
+    let root = Builder::new().prefix("test_easy_migrator_failed_migration_1").tempdir().expect("tempdir");
+    fs::create_dir_all(root.path()).expect("dir created");
+
+    let dbfile = root.path().join("data.mdb");
+    fs::write(&dbfile, "bogus").expect("bogus dbfile created");
+
+    // This won't fail with FileInvalid even though the database is a bogus file, because this
+    // is the "easy mode" migration which automatically handles (ignores) this error.
+    let dst_env = Rkv::new::<SafeMode>(root.path()).expect("new succeeded");
+    Migrator::easy_migrate_lmdb_to_safe_mode(root.path(), &dst_env).expect("migrated");
+
+    // Populate destination environment and persist to disk.
+    populate_store!(&dst_env);
+    dst_env.sync(true).expect("synced");
+
+    // Delete bogus file and create a valid source environment in its place.
+    fs::remove_file(&dbfile).expect("bogus dbfile removed");
+    let src_env = Rkv::new::<Lmdb>(root.path()).expect("new succeeded");
+    populate_store!(&src_env);
+    src_env.sync(true).expect("synced");
+
+    // Attempt to migrate again. This should *NOT* fail with DestinationNotEmpty.
+    Migrator::easy_migrate_lmdb_to_safe_mode(root.path(), &dst_env).expect("migrated");
+}
+
+#[test]
+fn test_easy_migrator_failed_migration_2() {
+    let root = Builder::new().prefix("test_easy_migrator_failed_migration_2").tempdir().expect("tempdir");
+    fs::create_dir_all(root.path()).expect("dir created");
+
+    let dbfile = root.path().join("data.safe.bin");
+    fs::write(&dbfile, "bogus").expect("bogus dbfile created");
+
+    // This won't fail with FileInvalid even though the database is a bogus file, because this
+    // is the "easy mode" migration which automatically handles (ignores) this error.
+    let dst_env = Rkv::new::<Lmdb>(root.path()).expect("new succeeded");
+    Migrator::easy_migrate_safe_mode_to_lmdb(root.path(), &dst_env).expect("migrated");
+
+    // Populate destination environment and persist to disk.
+    populate_store!(&dst_env);
+    dst_env.sync(true).expect("synced");
+
+    // Delete bogus file and create a valid source environment in its place.
+    fs::remove_file(&dbfile).expect("bogus dbfile removed");
+    let src_env = Rkv::new::<SafeMode>(root.path()).expect("new succeeded");
+    populate_store!(&src_env);
+    src_env.sync(true).expect("synced");
+
+    // Attempt to migrate again. This should *NOT* fail with DestinationNotEmpty.
+    Migrator::easy_migrate_safe_mode_to_lmdb(root.path(), &dst_env).expect("migrated");
+}
+
+fn test_easy_migrator_from_manager_failed_migration_1() {
+    let root = Builder::new().prefix("test_easy_migrator_from_manager_failed_migration_1").tempdir().expect("tempdir");
+    fs::create_dir_all(root.path()).expect("dir created");
+
+    {
+        let mut src_manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
+        let created_src_arc = src_manager.get_or_create(root.path(), Rkv::new::<Lmdb>).unwrap();
+        let src_env = created_src_arc.read().unwrap();
+        populate_store!(&src_env);
+        src_env.sync(true).expect("synced");
+    }
+    {
+        let mut dst_manager = Manager::<SafeModeEnvironment>::singleton().write().unwrap();
+        let created_dst_arc_1 = dst_manager.get_or_create(root.path(), Rkv::new::<SafeMode>).unwrap();
+        let dst_env_1 = created_dst_arc_1.read().unwrap();
+        populate_store!(&dst_env_1);
+        dst_env_1.sync(true).expect("synced");
+    }
+
+    // Attempt to migrate again in a new env. This should *NOT* fail with DestinationNotEmpty.
+    let dst_manager = Manager::<SafeModeEnvironment>::singleton().read().unwrap();
+    let created_dst_arc_2 = dst_manager.get(root.path()).unwrap().unwrap();
+    let dst_env_2 = created_dst_arc_2.read().unwrap();
+    Migrator::easy_migrate_lmdb_to_safe_mode(root.path(), dst_env_2).expect("migrated");
+}
+
+fn test_easy_migrator_from_manager_failed_migration_2() {
+    let root = Builder::new().prefix("test_easy_migrator_from_manager_failed_migration_2").tempdir().expect("tempdir");
+    fs::create_dir_all(root.path()).expect("dir created");
+
+    {
+        let mut src_manager = Manager::<SafeModeEnvironment>::singleton().write().unwrap();
+        let created_src_arc = src_manager.get_or_create(root.path(), Rkv::new::<SafeMode>).unwrap();
+        let src_env = created_src_arc.read().unwrap();
+        populate_store!(&src_env);
+        src_env.sync(true).expect("synced");
+    }
+    {
+        let mut dst_manager = Manager::<LmdbEnvironment>::singleton().write().unwrap();
+        let created_dst_arc_1 = dst_manager.get_or_create(root.path(), Rkv::new::<Lmdb>).unwrap();
+        let dst_env_1 = created_dst_arc_1.read().unwrap();
+        populate_store!(&dst_env_1);
+        dst_env_1.sync(true).expect("synced");
+    }
+
+    // Attempt to migrate again in a new env. This should *NOT* fail with DestinationNotEmpty.
+    let dst_manager = Manager::<LmdbEnvironment>::singleton().read().unwrap();
+    let created_dst_arc_2 = dst_manager.get(root.path()).unwrap().unwrap();
+    let dst_env_2 = created_dst_arc_2.read().unwrap();
+    Migrator::easy_migrate_safe_mode_to_lmdb(root.path(), dst_env_2).expect("migrated");
+}
+
+#[test]
+fn test_easy_migrator_from_manager_failed_migration() {
+    test_easy_migrator_from_manager_failed_migration_1();
+    test_easy_migrator_from_manager_failed_migration_2();
+}