diff --git a/pallets/subtensor/src/macros/events.rs b/pallets/subtensor/src/macros/events.rs index 4fa69b100a..3da3a344ab 100644 --- a/pallets/subtensor/src/macros/events.rs +++ b/pallets/subtensor/src/macros/events.rs @@ -487,5 +487,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/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs index 899e8d32f2..811aa0f335 100644 --- a/pallets/subtensor/src/macros/hooks.rs +++ b/pallets/subtensor/src/macros/hooks.rs @@ -166,7 +166,10 @@ mod hooks { // Fix staking hot keys .saturating_add(migrations::migrate_fix_staking_hot_keys::migrate_fix_staking_hot_keys::()) // Migrate coldkey swap scheduled to announcements - .saturating_add(migrations::migrate_coldkey_swap_scheduled_to_announcements::migrate_coldkey_swap_scheduled_to_announcements::()); + .saturating_add(migrations::migrate_coldkey_swap_scheduled_to_announcements::migrate_coldkey_swap_scheduled_to_announcements::()) + // 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 23a2899b94..4e7cc36e8c 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -43,6 +43,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/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..8854be0fbd 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 { @@ -397,8 +401,8 @@ impl Pallet { for key in keys { let (_, coldkey, netuid) = key.clone(); - if netuid == NetUid::ROOT { - Self::maybe_add_coldkey_index(&coldkey); + 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)); @@ -411,6 +415,11 @@ 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) { diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index beec7a3cba..d28fe93a7e 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -3039,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." + ); + }); +}