added migration to remove total hk/ck stakes this interval

Author: l0r1s

pallets/subtensor/src/macros/hooks.rs

@@ -87,7 +87,9 @@ mod hooks {
                 // Remove all zero value entries in TotalHotkeyAlpha
                 .saturating_add(migrations::migrate_remove_zero_total_hotkey_alpha::migrate_remove_zero_total_hotkey_alpha::<T>())
                 // Wipe existing items to prevent bad decoding for new type
-                .saturating_add(migrations::migrate_upgrade_revealed_commitments::migrate_upgrade_revealed_commitments::<T>());
+                .saturating_add(migrations::migrate_upgrade_revealed_commitments::migrate_upgrade_revealed_commitments::<T>())
+                // Remove all entries in TotalHotkeyColdkeyStakesThisInterval
+                .saturating_add(migrations::migrate_remove_total_hotkey_coldkey_stakes_this_interval::migrate_remove_total_hotkey_coldkey_stakes_this_interval::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_remove_total_hotkey_coldkey_stakes_this_interval.rs

@@ -0,0 +1,51 @@
+use super::*;
+use crate::HasMigrationRun;
+use frame_support::{traits::Get, weights::Weight};
+use sp_io::{KillStorageResult, hashing::twox_128, storage::clear_prefix};
+
+pub fn migrate_remove_total_hotkey_coldkey_stakes_this_interval<T: Config>() -> Weight {
+    let migration_name = "migrate_remove_total_hotkey_coldkey_stakes_this_interval";
+    let migration_name_bytes = migration_name.as_bytes().to_vec();
+
+    let mut weight = T::DbWeight::get().reads(1);
+    if HasMigrationRun::<T>::get(&migration_name_bytes) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            migration_name
+        );
+        return weight;
+    }
+
+    log::info!("Running migration '{}'", migration_name);
+
+    let pallet_name = twox_128(b"SubtensorModule");
+    let storage_name = twox_128(b"TotalHotkeyColdkeyStakesThisInterval");
+    let prefix = [pallet_name, storage_name].concat();
+
+    // Try to remove all entries from the storage, if some entries are remaining,
+    // the migration will re-run again on next blocks until all entries are removed.
+    let removed_entries_count = match clear_prefix(&prefix, Some(u32::MAX)) {
+        KillStorageResult::AllRemoved(removed) => {
+            // Mark migration as completed
+            HasMigrationRun::<T>::insert(&migration_name_bytes, true);
+            weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+            log::info!(
+                "Migration '{:?}' completed successfully. {:?} entries removed.",
+                migration_name,
+                removed
+            );
+            removed
+        }
+        KillStorageResult::SomeRemaining(removed) => {
+            log::info!(
+                "Migration '{:?}' completed partially. {:?} entries removed.",
+                migration_name,
+                removed
+            );
+            removed
+        }
+    };
+
+    weight.saturating_add(T::DbWeight::get().writes(removed_entries_count as u64))
+}

pallets/subtensor/src/migrations/mod.rs

@@ -10,6 +10,7 @@ pub mod migrate_init_total_issuance;
 pub mod migrate_populate_owned_hotkeys;
 pub mod migrate_rao;
 pub mod migrate_remove_stake_map;
+pub mod migrate_remove_total_hotkey_coldkey_stakes_this_interval;
 pub mod migrate_remove_unused_maps_and_values;
 pub mod migrate_remove_zero_total_hotkey_alpha;
 pub mod migrate_set_first_emission_block_number;

pallets/subtensor/src/tests/migration.rs

@@ -555,3 +555,39 @@ fn test_migrate_revealed_commitments() {
         assert!(!weight.is_zero(), "Migration weight should be non-zero");
     });
 }
+
+#[test]
+fn test_migrate_remove_total_hotkey_coldkey_stakes_this_interval() {
+    new_test_ext(1).execute_with(|| {
+        const MIGRATION_NAME: &str = "migrate_remove_total_hotkey_coldkey_stakes_this_interval";
+
+        let pallet_name = twox_128(b"SubtensorModule");
+        let storage_name = twox_128(b"TotalHotkeyColdkeyStakesThisInterval");
+        let prefix = [pallet_name, storage_name].concat();
+
+        // Set up 200 000 entries to be deleted.
+        for i in 0..200_000{
+            let hotkey = U256::from(i as u64);
+            let coldkey = U256::from(i as u64);
+            let key = [prefix.clone(), hotkey.encode(), coldkey.encode()].concat();
+            let value = (100 + i, 200 + i);
+            put_raw(&key, &value.encode());
+        }
+
+        assert!(frame_support::storage::unhashed::contains_prefixed_key(&prefix), "Entries should exist before migration.");
+        assert!(
+            !HasMigrationRun::<Test>::get(MIGRATION_NAME.as_bytes().to_vec()),
+            "Migration should not have run yet."
+        );
+
+        // Run migration
+        let weight = crate::migrations::migrate_remove_total_hotkey_coldkey_stakes_this_interval::migrate_remove_total_hotkey_coldkey_stakes_this_interval::<Test>();
+
+        assert!(!frame_support::storage::unhashed::contains_prefixed_key(&prefix), "All entries should have been removed.");
+        assert!(
+            HasMigrationRun::<Test>::get(MIGRATION_NAME.as_bytes().to_vec()),
+            "Migration should be marked as run."
+        );
+        assert!(!weight.is_zero(),"Migration weight should be non-zero.");
+    });
+}