Migration to fix StakingHotKeys added

Author: grbIzl

pallets/subtensor/src/macros/hooks.rs

@@ -162,7 +162,9 @@ mod hooks {
                 // Migrate pending emissions
                 .saturating_add(migrations::migrate_pending_emissions::migrate_pending_emissions::<T>())
                 // Reset unactive subnets
-                .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::<T>());
+                .saturating_add(migrations::migrate_reset_unactive_sn::migrate_reset_unactive_sn::<T>())
+                // Fix staking hot keys
+                .saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::<T>());
             weight
         }
 

pallets/subtensor/src/migrations/migrate_fix_staking_hot_keys.rs

@@ -0,0 +1,55 @@
+use super::*;
+use frame_support::{traits::Get, weights::Weight};
+use log;
+use scale_info::prelude::string::String;
+use sp_std::collections::btree_map::BTreeMap;
+
+pub fn migrate_fix_staking_hot_keys<T: Config>() -> Weight {
+    let migration_name = b"migrate_fix_staking_hot_keys".to_vec();
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Skip if already executed
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            target: "runtime",
+            "Migration '{}' already run - skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+
+    let mut cache: BTreeMap<T::AccountId, Vec<T::AccountId>> = BTreeMap::new();
+    let mut storage_reads: u64 = 0;
+    let mut storage_writes: u64 = 0;
+
+    for ((hotkey, coldkey, _netuid), alpha) in Alpha::<T>::iter() {
+        if alpha == 0 {
+            continue;
+        }
+
+        let staking_hotkeys = cache.entry(coldkey.clone())
+            .or_insert_with(|| {
+                storage_reads = storage_reads.saturating_add(1);
+                StakingHotkeys::<T>::get(&coldkey)
+            });
+
+        if !staking_hotkeys.contains(&hotkey) {
+            staking_hotkeys.push(hotkey.clone());
+            storage_writes = storage_writes.saturating_add(1);
+            StakingHotkeys::<T>::insert(&coldkey, staking_hotkeys.clone());
+        }
+    }
+    weight = weight.saturating_add(T::DbWeight::get().reads_writes(storage_reads, storage_writes));
+
+    // Mark migration done
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        target: "runtime",
+        "Migration '{}' completed.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    weight
+}

pallets/subtensor/src/migrations/mod.rs

@@ -58,6 +58,7 @@ pub mod migrate_to_v1_separate_emission;
 pub mod migrate_to_v2_fixed_total_stake;
 pub mod migrate_transfer_ownership_to_foundation;
 pub mod migrate_upgrade_revealed_commitments;
+pub mod migrate_fix_staking_hot_keys;
 
 pub(crate) fn migrate_storage<T: Config>(
     migration_name: &'static str,

pallets/subtensor/src/tests/migration.rs

@@ -1138,6 +1138,43 @@ fn test_migrate_rate_limit_keys() {
     });
 }
 
+#[test]
+fn test_migrate_fix_staking_hot_keys() {
+    new_test_ext(1).execute_with(|| {
+        const MIGRATION_NAME: &[u8] = b"migrate_fix_staking_hot_keys";
+
+        assert!(
+            !HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should not have run yet"
+        );
+
+        // Add some data
+        Alpha::<Test>::insert(
+            (U256::from(1),
+            U256::from(2),
+            NetUid::ROOT),
+            U64F64::from(1_u64)
+        );
+        // Run migration
+        let weight = migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::<Test>();
+
+        assert!(
+            HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()),
+            "Migration should be marked as completed"
+        );
+
+        // Check migration has been marked as run
+        assert!(HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()));
+
+        // Verify results
+        assert_eq!(
+            StakingHotkeys::<Test>::get(U256::from(2)),
+            vec![U256::from(1)]
+        );
+    });
+}
+
+
 #[test]
 fn test_migrate_fix_root_subnet_tao() {
     new_test_ext(1).execute_with(|| {