diff --git a/pallets/subtensor/src/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs index 31be3e5e4f..ed57d52c8b 100644 --- a/pallets/subtensor/src/macros/hooks.rs +++ b/pallets/subtensor/src/macros/hooks.rs @@ -162,7 +162,9 @@ mod hooks { // Remove old identity map entries(Identities, SubnetIdentities, SubnetIdentitiesV2) .saturating_add(migrations::migrate_remove_old_identity_maps::migrate_remove_old_identity_maps::()) // Remove unknown neuron axon, certificate prom - .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::()); + .saturating_add(migrations::migrate_remove_unknown_neuron_axon_cert_prom::migrate_remove_unknown_neuron_axon_cert_prom::()) + // Fix staking hot keys + .saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::()); weight } diff --git a/pallets/subtensor/src/migrations/migrate_fix_staking_hot_keys.rs b/pallets/subtensor/src/migrations/migrate_fix_staking_hot_keys.rs new file mode 100644 index 0000000000..8c0358614d --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_fix_staking_hot_keys.rs @@ -0,0 +1,56 @@ +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() -> 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::::get(&migration_name) { + log::info!( + target: "runtime", + "Migration '{}' already run - skipping.", + String::from_utf8_lossy(&migration_name) + ); + return weight; + } + + let mut cache: BTreeMap> = BTreeMap::new(); + let mut storage_reads: u64 = 0; + let mut storage_writes: u64 = 0; + + for ((hotkey, coldkey, _netuid), alpha) in Alpha::::iter() { + storage_reads = storage_reads.saturating_add(1); + + if alpha == 0 { + continue; + } + + let staking_hotkeys = cache.entry(coldkey.clone()).or_insert_with(|| { + storage_reads = storage_reads.saturating_add(1); + StakingHotkeys::::get(&coldkey) + }); + + if !staking_hotkeys.contains(&hotkey) { + staking_hotkeys.push(hotkey.clone()); + storage_writes = storage_writes.saturating_add(1); + StakingHotkeys::::insert(&coldkey, staking_hotkeys.clone()); + } + } + weight = weight.saturating_add(T::DbWeight::get().reads_writes(storage_reads, storage_writes)); + + // Mark migration done + HasMigrationRun::::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 +} diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs index 4c9d5f01d1..a03da9289e 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -19,6 +19,7 @@ pub mod migrate_fix_childkeys; pub mod migrate_fix_is_network_member; pub mod migrate_fix_root_subnet_tao; pub mod migrate_fix_root_tao_and_alpha_in; +pub mod migrate_fix_staking_hot_keys; pub mod migrate_init_tao_flow; pub mod migrate_init_total_issuance; pub mod migrate_kappa_map_to_default; diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index 1d78b4dffd..bed77e797f 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -1138,6 +1138,41 @@ 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::::get(MIGRATION_NAME.to_vec()), + "Migration should not have run yet" + ); + + // Add some data + Alpha::::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::(); + + assert!( + HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should be marked as completed" + ); + + // Check migration has been marked as run + assert!(HasMigrationRun::::get(MIGRATION_NAME.to_vec())); + + // Verify results + assert_eq!( + StakingHotkeys::::get(U256::from(2)), + vec![U256::from(1)] + ); + }); +} + #[test] fn test_migrate_fix_root_subnet_tao() { new_test_ext(1).execute_with(|| {