From e1bc8891e76dd115810810b45465e2c8975e4c2d Mon Sep 17 00:00:00 2001
From: acerone85 <andrea.cerone@gmail.com>
Date: Thu, 3 Oct 2024 16:50:26 +0100
Subject: [PATCH] Test that migration can complete during normal DB operations

---
 .../fuel-core/src/state/historical_rocksdb.rs | 104 ++++++++++++++++++
 1 file changed, 104 insertions(+)

diff --git a/crates/fuel-core/src/state/historical_rocksdb.rs b/crates/fuel-core/src/state/historical_rocksdb.rs
index 7c7a7c6c95f..2800c751d55 100644
--- a/crates/fuel-core/src/state/historical_rocksdb.rs
+++ b/crates/fuel-core/src/state/historical_rocksdb.rs
@@ -170,6 +170,23 @@ where
     }
 }
 
+impl<Description> TryFrom<InnerHistoricalRocksDB<Description>>
+    for HistoricalRocksDB<Description>
+where
+    Description: DatabaseDescription,
+{
+    type Error = DatabaseError;
+
+    fn try_from(inner: InnerHistoricalRocksDB<Description>) -> DatabaseResult<Self> {
+        let InnerHistoricalRocksDB {
+            db,
+            state_rewind_policy,
+            ..
+        } = inner;
+        Self::new(db, state_rewind_policy)
+    }
+}
+
 impl<Description> KeyValueInspect for HistoricalRocksDB<Description>
 where
     Description: DatabaseDescription,
@@ -1424,6 +1441,93 @@ mod tests {
         assert!(!historical_rocks_db.is_migration_in_progress());
     }
 
+    #[tokio::test]
+    async fn historical_rocksdb_perform_migration_on_new_instance() {
+        // Given
+        let rocks_db = RocksDb::<Historical<OnChain>>::default_open_temp(None).unwrap();
+
+        let historical_rocks_db =
+            InnerHistoricalRocksDB::new(rocks_db, StateRewindPolicy::RewindFullRange)
+                .unwrap();
+
+        // Commit 100 blocks
+        for i in 1..=100u32 {
+            let mut transaction = historical_rocks_db.read_transaction();
+            transaction
+                .storage_as_mut::<ContractsAssets>()
+                .insert(&key(), &(123 + i as u64))
+                .unwrap();
+            historical_rocks_db
+                .commit_changes(Some(i.into()), transaction.into_changes())
+                .unwrap();
+        }
+
+        let mut revert_migration_transaction = StorageTransaction::transaction(
+            &historical_rocks_db.db,
+            ConflictPolicy::Overwrite,
+            Changes::default(),
+        );
+
+        // Revert the modification history from v2 to v1
+        historical_rocks_db
+            .db
+            .iter_all::<ModificationsHistoryV2<OnChain>>(None)
+            .map(Result::unwrap)
+            .for_each(|(height, changes)| {
+                revert_migration_transaction
+                    .storage_as_mut::<ModificationsHistoryV1<OnChain>>()
+                    .insert(&height, &changes)
+                    .unwrap();
+                revert_migration_transaction
+                    .storage_as_mut::<ModificationsHistoryV2<OnChain>>()
+                    .remove(&height)
+                    .unwrap();
+            });
+
+        historical_rocks_db
+            .db
+            .commit_changes(&revert_migration_transaction.into_changes())
+            .unwrap();
+
+        let v1_changes = historical_rocks_db
+            .db
+            .iter_all::<ModificationsHistoryV1<OnChain>>(None)
+            .count();
+        let v2_changes = historical_rocks_db
+            .db
+            .iter_all::<ModificationsHistoryV2<OnChain>>(None)
+            .count();
+
+        assert_eq!(v1_changes, 100);
+        assert_eq!(v2_changes, 0);
+
+        // When
+        // Wrap the inner historical rocksdb in a new instance to start the migration.
+        let historical_rocks_db_with_migration =
+            HistoricalRocksDB::try_from(historical_rocks_db).unwrap();
+
+        // Keep writing to the database until the migration is complete.
+        let mut i = 101;
+        while historical_rocks_db_with_migration
+            .inner
+            .is_migration_in_progress()
+        {
+            let mut transaction = historical_rocks_db_with_migration.read_transaction();
+            transaction
+                .storage_as_mut::<ContractsAssets>()
+                .insert(&key(), &(123 + i as u64))
+                .unwrap();
+            historical_rocks_db_with_migration
+                .commit_changes(Some(i.into()), transaction.into_changes())
+                .unwrap();
+            i += 1;
+        }
+        // Then
+        assert!(!historical_rocks_db_with_migration
+            .inner
+            .is_migration_in_progress());
+    }
+
     #[test]
     fn state_rewind_policy__rewind_range_1__second_rollback_fails() {
         // Given