Optimized cleanup

AstarNetwork · Dec 21, 2023 · 6016046 · 6016046
1 parent b487f00
commit 6016046
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/pallets/dapp-staking-migration/src/lib.rs b/pallets/dapp-staking-migration/src/lib.rs
@@ -271,9 +271,15 @@ pub mod pallet {
                 Self::deposit_event(Event::<T>::EntriesDeleted(entries_deleted));
             }
 
-            // Put the pallet back into maintenance mode.
+            // Put the pallet back into maintenance mode in case we're still migration the old storage over,
+            // otherwise disable the maintenance mode.
             pallet_dapp_staking_v3::ActiveProtocolState::<T>::mutate(|state| {
-                state.maintenance = true;
+                state.maintenance = match migration_state {
+                    MigrationState::NotInProgress
+                    | MigrationState::RegisteredDApps
+                    | MigrationState::Ledgers => true,
+                    MigrationState::Cleanup | MigrationState::Finished => false,
+                };
             });
 
             if migration_state != init_migration_state {

diff --git a/pallets/dapp-staking-migration/src/tests.rs b/pallets/dapp-staking-migration/src/tests.rs
@@ -148,10 +148,20 @@ fn migrate_call_works() {
                 Some(Weight::from_parts(1, 1))
             ));
 
-            assert!(
-                pallet_dapp_staking_v3::ActiveProtocolState::<Test>::get().maintenance,
-                "Maintenance must always be returned after migrate call finishes."
-            );
+            match MigrationStateStorage::<Test>::get() {
+                MigrationState::RegisteredDApps | MigrationState::Ledgers => {
+                    assert!(
+                        pallet_dapp_staking_v3::ActiveProtocolState::<Test>::get().maintenance,
+                        "Pallet must be in the maintenance mode during old storage migration."
+                    );
+                }
+                _ => {
+                    assert!(
+                        !pallet_dapp_staking_v3::ActiveProtocolState::<Test>::get().maintenance,
+                        "Maintenance mode is disabled during old storage cleanup."
+                    );
+                }
+            }
         }
 
         // Check post-state