From 0ee04e7a846ba2fa822da1b4575cc37619b81048 Mon Sep 17 00:00:00 2001 From: Shamil Gadelshin Date: Tue, 27 Jan 2026 16:55:27 +0300 Subject: [PATCH 1/3] Add an event to slow coldkey migration --- pallets/subtensor/src/macros/events.rs | 8 ++++++++ pallets/subtensor/src/staking/claim_root.rs | 6 +++++- pallets/subtensor/src/staking/helpers.rs | 10 +++++++++- 3 files changed, 22 insertions(+), 2 deletions(-) diff --git a/pallets/subtensor/src/macros/events.rs b/pallets/subtensor/src/macros/events.rs index c86cc1a1e5..bd100b9fc5 100644 --- a/pallets/subtensor/src/macros/events.rs +++ b/pallets/subtensor/src/macros/events.rs @@ -481,5 +481,13 @@ mod events { /// The amount of alpha distributed alpha: AlphaCurrency, }, + + /// Root claim coldkey migration event. + RootClaimColdkeyMigrated { + /// Keys migrated this block + processed_keys: u64, + /// Added coldkeys + added_coldkeys: u64, + }, } } diff --git a/pallets/subtensor/src/staking/claim_root.rs b/pallets/subtensor/src/staking/claim_root.rs index 24a26d154c..ccb0c7263e 100644 --- a/pallets/subtensor/src/staking/claim_root.rs +++ b/pallets/subtensor/src/staking/claim_root.rs @@ -316,13 +316,17 @@ impl Pallet { .saturating_add(Weight::from_parts(100_412, 0).saturating_mul(k.into())) } - pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) { + pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) -> bool { if !StakingColdkeys::::contains_key(coldkey) { let n = NumStakingColdkeys::::get(); StakingColdkeysByIndex::::insert(n, coldkey.clone()); StakingColdkeys::::insert(coldkey.clone(), n); NumStakingColdkeys::::mutate(|n| *n = n.saturating_add(1)); + + return true; } + + false } pub fn run_auto_claim_root_divs(last_block_hash: T::Hash) -> Weight { diff --git a/pallets/subtensor/src/staking/helpers.rs b/pallets/subtensor/src/staking/helpers.rs index 099f8e26b6..8d2df436f9 100644 --- a/pallets/subtensor/src/staking/helpers.rs +++ b/pallets/subtensor/src/staking/helpers.rs @@ -367,6 +367,8 @@ impl Pallet { /// alphas. It keeps the alpha value stored when it's >= than MIN_ALPHA. /// The function uses AlphaMapLastKey as a storage for key iterator between runs. pub fn populate_root_coldkey_staking_maps() { + let mut processed_keys = 0u64; + let mut added_coldkeys = 0u64; // Get starting key for the batch. Get the first key if we restart the process. let mut new_starting_raw_key = AlphaMapLastKey::::get(); let mut starting_key = None; @@ -381,6 +383,8 @@ impl Pallet { .take(ALPHA_MAP_BATCH_SIZE) .collect::>(); + processed_keys = processed_keys.saturating_add(keys.len() as u64); + // New iteration: insert the starting key in the batch if it's a new iteration // iter_keys_from() skips the starting key if let Some(starting_key) = starting_key { @@ -398,7 +402,9 @@ impl Pallet { let (_, coldkey, netuid) = key.clone(); if netuid == NetUid::ROOT { - Self::maybe_add_coldkey_index(&coldkey); + if Self::maybe_add_coldkey_index(&coldkey) { + added_coldkeys = added_coldkeys.saturating_add(1); + } } new_starting_key = Some(Alpha::::hashed_key_for(key)); @@ -411,6 +417,8 @@ impl Pallet { AlphaMapLastKey::::put(new_starting_key); } + + Self::deposit_event(Event::::RootClaimColdkeyMigrated { processed_keys, added_coldkeys }) } pub fn burn_subnet_alpha(_netuid: NetUid, _amount: AlphaCurrency) { From 643d079d20c6b11026227a6ae47c93189db787af Mon Sep 17 00:00:00 2001 From: Shamil Gadelshin Date: Tue, 27 Jan 2026 17:39:19 +0300 Subject: [PATCH 2/3] Reset coldkey migration for root claim --- pallets/subtensor/src/macros/hooks.rs | 5 +- .../migrate_reset_coldkey_for_root_claim.rs | 42 ++++++++++++ pallets/subtensor/src/migrations/mod.rs | 1 + pallets/subtensor/src/staking/helpers.rs | 5 +- pallets/subtensor/src/tests/migration.rs | 67 +++++++++++++++++++ 5 files changed, 118 insertions(+), 2 deletions(-) create mode 100644 pallets/subtensor/src/migrations/migrate_reset_coldkey_for_root_claim.rs diff --git a/pallets/subtensor/src/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs index ed57d52c8b..e779392787 100644 --- a/pallets/subtensor/src/macros/hooks.rs +++ b/pallets/subtensor/src/macros/hooks.rs @@ -164,7 +164,10 @@ mod hooks { // Remove unknown neuron axon, certificate 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::()); + .saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::()) + // Reset coldkey migration + .saturating_add(migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::()); + weight } diff --git a/pallets/subtensor/src/migrations/migrate_reset_coldkey_for_root_claim.rs b/pallets/subtensor/src/migrations/migrate_reset_coldkey_for_root_claim.rs new file mode 100644 index 0000000000..bdf4b79298 --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_reset_coldkey_for_root_claim.rs @@ -0,0 +1,42 @@ +use super::*; +use crate::HasMigrationRun; +use frame_support::{traits::Get, weights::Weight}; +use scale_info::prelude::string::String; + +pub fn migrate_reset_coldkey_for_root_claim() -> Weight { + let migration_name = b"migrate_reset_coldkey_for_root_claim".to_vec(); + let mut weight = T::DbWeight::get().reads(1); + + if HasMigrationRun::::get(&migration_name) { + log::info!( + "Migration '{:?}' has already run. Skipping.", + String::from_utf8_lossy(&migration_name) + ); + return weight; + } + + log::info!( + "Running migration '{}'", + String::from_utf8_lossy(&migration_name) + ); + + // ------------------------------ + // Step 1: Reset AlphaMapLastKey + // ------------------------------ + + AlphaMapLastKey::::kill(); + + // ------------------------------ + // Step 2: Mark Migration as Completed + // ------------------------------ + + HasMigrationRun::::insert(&migration_name, true); + weight = weight.saturating_add(T::DbWeight::get().writes(1)); + + log::info!( + "Migration '{:?}' completed successfully.", + String::from_utf8_lossy(&migration_name) + ); + + weight +} diff --git a/pallets/subtensor/src/migrations/mod.rs b/pallets/subtensor/src/migrations/mod.rs index a03da9289e..d76a94f8b0 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -42,6 +42,7 @@ pub mod migrate_remove_unknown_neuron_axon_cert_prom; pub mod migrate_remove_unused_maps_and_values; pub mod migrate_remove_zero_total_hotkey_alpha; pub mod migrate_reset_bonds_moving_average; +pub mod migrate_reset_coldkey_for_root_claim; pub mod migrate_reset_max_burn; pub mod migrate_reset_unactive_sn; pub mod migrate_set_first_emission_block_number; diff --git a/pallets/subtensor/src/staking/helpers.rs b/pallets/subtensor/src/staking/helpers.rs index 8d2df436f9..2ebcbb8536 100644 --- a/pallets/subtensor/src/staking/helpers.rs +++ b/pallets/subtensor/src/staking/helpers.rs @@ -418,7 +418,10 @@ impl Pallet { AlphaMapLastKey::::put(new_starting_key); } - Self::deposit_event(Event::::RootClaimColdkeyMigrated { processed_keys, added_coldkeys }) + Self::deposit_event(Event::::RootClaimColdkeyMigrated { + processed_keys, + added_coldkeys, + }) } pub fn burn_subnet_alpha(_netuid: NetUid, _amount: AlphaCurrency) { diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index bed77e797f..a0653a189e 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -2968,3 +2968,70 @@ fn test_migrate_remove_unknown_neuron_axon_cert_prom() { } } } + +#[test] +fn test_migrate_reset_coldkey_for_root_claim() { + new_test_ext(1).execute_with(|| { + const MIGRATION_NAME: &[u8] = b"migrate_reset_coldkey_for_root_claim"; + + // Step 1: Set up initial state - put a value in AlphaMapLastKey + let test_key: Option> = Some(vec![1, 2, 3, 4, 5]); + AlphaMapLastKey::::put(test_key.clone()); + + // Verify the value was set + assert_eq!( + AlphaMapLastKey::::get(), + test_key, + "AlphaMapLastKey should have the test value." + ); + + // Step 2: Verify migration hasn't run yet + assert!( + !HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should not have run yet." + ); + + // Step 3: Run the migration + let weight = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::(); + + // Step 4: Verify AlphaMapLastKey was killed (reset to default None) + assert_eq!( + AlphaMapLastKey::::get(), + None, + "AlphaMapLastKey should be reset to None after migration." + ); + + // Step 5: Verify migration is marked as completed + assert!( + HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should be marked as run." + ); + + // Step 6: Verify weight is non-zero + assert!( + !weight.is_zero(), + "Migration weight should be non-zero." + ); + + // Step 7: Run migration again - should be idempotent + // Set the value again to verify migration doesn't run twice + let test_key_2: Option> = Some(vec![6, 7, 8, 9, 10]); + AlphaMapLastKey::::put(test_key_2.clone()); + + let weight_second_run = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::(); + + // Value should NOT be reset since migration already ran + assert_eq!( + AlphaMapLastKey::::get(), + test_key_2, + "AlphaMapLastKey should not be changed on second migration run." + ); + + // Second run should only have read weight (checking if migration ran) + assert_eq!( + weight_second_run, + ::DbWeight::get().reads(1), + "Second migration run should only have read weight." + ); + }); +} From 2e713a67f1398bf77523a79a30fd18c4e37a4653 Mon Sep 17 00:00:00 2001 From: Shamil Gadelshin Date: Fri, 30 Jan 2026 18:08:05 +0300 Subject: [PATCH 3/3] Fix merge and clippy issues --- pallets/subtensor/src/staking/helpers.rs | 6 +- pallets/subtensor/src/tests/migration.rs | 132 ++++++++++++----------- 2 files changed, 69 insertions(+), 69 deletions(-) diff --git a/pallets/subtensor/src/staking/helpers.rs b/pallets/subtensor/src/staking/helpers.rs index 2ebcbb8536..8854be0fbd 100644 --- a/pallets/subtensor/src/staking/helpers.rs +++ b/pallets/subtensor/src/staking/helpers.rs @@ -401,10 +401,8 @@ impl Pallet { for key in keys { let (_, coldkey, netuid) = key.clone(); - if netuid == NetUid::ROOT { - if Self::maybe_add_coldkey_index(&coldkey) { - added_coldkeys = added_coldkeys.saturating_add(1); - } + if netuid == NetUid::ROOT && Self::maybe_add_coldkey_index(&coldkey) { + added_coldkeys = added_coldkeys.saturating_add(1); } new_starting_key = Some(Alpha::::hashed_key_for(key)); diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index 07d7f750ad..d28fe93a7e 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -2969,71 +2969,6 @@ fn test_migrate_remove_unknown_neuron_axon_cert_prom() { } #[test] -fn test_migrate_reset_coldkey_for_root_claim() { - new_test_ext(1).execute_with(|| { - const MIGRATION_NAME: &[u8] = b"migrate_reset_coldkey_for_root_claim"; - - // Step 1: Set up initial state - put a value in AlphaMapLastKey - let test_key: Option> = Some(vec![1, 2, 3, 4, 5]); - AlphaMapLastKey::::put(test_key.clone()); - - // Verify the value was set - assert_eq!( - AlphaMapLastKey::::get(), - test_key, - "AlphaMapLastKey should have the test value." - ); - - // Step 2: Verify migration hasn't run yet - assert!( - !HasMigrationRun::::get(MIGRATION_NAME.to_vec()), - "Migration should not have run yet." - ); - - // Step 3: Run the migration - let weight = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::(); - - // Step 4: Verify AlphaMapLastKey was killed (reset to default None) - assert_eq!( - AlphaMapLastKey::::get(), - None, - "AlphaMapLastKey should be reset to None after migration." - ); - - // Step 5: Verify migration is marked as completed - assert!( - HasMigrationRun::::get(MIGRATION_NAME.to_vec()), - "Migration should be marked as run." - ); - - // Step 6: Verify weight is non-zero - assert!( - !weight.is_zero(), - "Migration weight should be non-zero." - ); - - // Step 7: Run migration again - should be idempotent - // Set the value again to verify migration doesn't run twice - let test_key_2: Option> = Some(vec![6, 7, 8, 9, 10]); - AlphaMapLastKey::::put(test_key_2.clone()); - - let weight_second_run = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::(); - - // Value should NOT be reset since migration already ran - assert_eq!( - AlphaMapLastKey::::get(), - test_key_2, - "AlphaMapLastKey should not be changed on second migration run." - ); - - // Second run should only have read weight (checking if migration ran) - assert_eq!( - weight_second_run, - ::DbWeight::get().reads(1), - "Second migration run should only have read weight." - ); - }); -} fn test_migrate_coldkey_swap_scheduled_to_announcements() { new_test_ext(1000).execute_with(|| { const MIGRATION_NAME: &[u8] = b"migrate_coldkey_swap_scheduled_to_announcements"; @@ -3104,3 +3039,70 @@ fn test_migrate_coldkey_swap_scheduled_to_announcements() { ); }); } + +#[test] +fn test_migrate_reset_coldkey_for_root_claim() { + new_test_ext(1).execute_with(|| { + const MIGRATION_NAME: &[u8] = b"migrate_reset_coldkey_for_root_claim"; + + // Step 1: Set up initial state - put a value in AlphaMapLastKey + let test_key: Option> = Some(vec![1, 2, 3, 4, 5]); + AlphaMapLastKey::::put(test_key.clone()); + + // Verify the value was set + assert_eq!( + AlphaMapLastKey::::get(), + test_key, + "AlphaMapLastKey should have the test value." + ); + + // Step 2: Verify migration hasn't run yet + assert!( + !HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should not have run yet." + ); + + // Step 3: Run the migration + let weight = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::(); + + // Step 4: Verify AlphaMapLastKey was killed (reset to default None) + assert_eq!( + AlphaMapLastKey::::get(), + None, + "AlphaMapLastKey should be reset to None after migration." + ); + + // Step 5: Verify migration is marked as completed + assert!( + HasMigrationRun::::get(MIGRATION_NAME.to_vec()), + "Migration should be marked as run." + ); + + // Step 6: Verify weight is non-zero + assert!( + !weight.is_zero(), + "Migration weight should be non-zero." + ); + + // Step 7: Run migration again - should be idempotent + // Set the value again to verify migration doesn't run twice + let test_key_2: Option> = Some(vec![6, 7, 8, 9, 10]); + AlphaMapLastKey::::put(test_key_2.clone()); + + let weight_second_run = crate::migrations::migrate_reset_coldkey_for_root_claim::migrate_reset_coldkey_for_root_claim::(); + + // Value should NOT be reset since migration already ran + assert_eq!( + AlphaMapLastKey::::get(), + test_key_2, + "AlphaMapLastKey should not be changed on second migration run." + ); + + // Second run should only have read weight (checking if migration ran) + assert_eq!( + weight_second_run, + ::DbWeight::get().reads(1), + "Second migration run should only have read weight." + ); + }); +}