diff --git a/pallets/subtensor/src/macros/hooks.rs b/pallets/subtensor/src/macros/hooks.rs index 3203ae312c..834e6c86bb 100644 --- a/pallets/subtensor/src/macros/hooks.rs +++ b/pallets/subtensor/src/macros/hooks.rs @@ -83,7 +83,9 @@ mod hooks { // Remove unused maps entries .saturating_add(migrations::migrate_remove_unused_maps_and_values::migrate_remove_unused_maps_and_values::()) // Set last emission block number for all existed subnets before start call feature applied - .saturating_add(migrations::migrate_set_first_emission_block_number::migrate_set_first_emission_block_number::()); + .saturating_add(migrations::migrate_set_first_emission_block_number::migrate_set_first_emission_block_number::()) + // Remove all zero value entries in TotalHotkeyAlpha + .saturating_add(migrations::migrate_remove_zero_total_hotkey_alpha::migrate_remove_zero_total_hotkey_alpha::()); weight } diff --git a/pallets/subtensor/src/migrations/migrate_remove_zero_total_hotkey_alpha.rs b/pallets/subtensor/src/migrations/migrate_remove_zero_total_hotkey_alpha.rs new file mode 100644 index 0000000000..3b45615bf4 --- /dev/null +++ b/pallets/subtensor/src/migrations/migrate_remove_zero_total_hotkey_alpha.rs @@ -0,0 +1,60 @@ +use super::*; +use frame_support::{traits::Get, weights::Weight}; +use log; +use scale_info::prelude::string::String; + +pub fn migrate_remove_zero_total_hotkey_alpha() -> Weight { + let migration_name = b"migrate_remove_zero_total_hotkey_alpha".to_vec(); + let mut weight = T::DbWeight::get().reads(1); + + // ------------------------------ + // Step 0: Check if already run + // ------------------------------ + if HasMigrationRun::::get(&migration_name) { + log::info!( + "Migration '{:?}' has already run. Skipping.", + migration_name + ); + return weight; + } + + log::info!( + "Running migration '{}'", + String::from_utf8_lossy(&migration_name) + ); + + // ------------------------------ + // Step 1: Remove any zero entries in TotalHotkeyAlpha + // ------------------------------ + + let mut removed_entries_count = 0u64; + + // For each (hotkey, netuid, alpha) entry, remove if alpha == 0 + for (hotkey, netuid, alpha) in TotalHotkeyAlpha::::iter() { + if alpha == 0 { + TotalHotkeyAlpha::::remove(&hotkey, netuid); + removed_entries_count = removed_entries_count.saturating_add(1); + } + } + + weight = weight.saturating_add(T::DbWeight::get().reads(removed_entries_count)); + weight = weight.saturating_add(T::DbWeight::get().writes(removed_entries_count)); + + log::info!( + "Removed {} zero entries from TotalHotkeyAlpha.", + removed_entries_count + ); + + // ------------------------------ + // 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 19e057e3ec..b342d54979 100644 --- a/pallets/subtensor/src/migrations/mod.rs +++ b/pallets/subtensor/src/migrations/mod.rs @@ -11,6 +11,7 @@ pub mod migrate_populate_owned_hotkeys; pub mod migrate_rao; pub mod migrate_remove_stake_map; pub mod migrate_remove_unused_maps_and_values; +pub mod migrate_remove_zero_total_hotkey_alpha; pub mod migrate_set_first_emission_block_number; pub mod migrate_set_min_burn; pub mod migrate_set_min_difficulty; diff --git a/pallets/subtensor/src/staking/recycle_alpha.rs b/pallets/subtensor/src/staking/recycle_alpha.rs index 109d26c850..cb5e740e84 100644 --- a/pallets/subtensor/src/staking/recycle_alpha.rs +++ b/pallets/subtensor/src/staking/recycle_alpha.rs @@ -42,7 +42,14 @@ impl Pallet { Error::::InsufficientLiquidity ); - TotalHotkeyAlpha::::mutate(&hotkey, netuid, |v| *v = v.saturating_sub(amount)); + if TotalHotkeyAlpha::::mutate(&hotkey, netuid, |v| { + *v = v.saturating_sub(amount); + *v + }) == 0 + { + TotalHotkeyAlpha::::remove(&hotkey, netuid); + } + SubnetAlphaOut::::mutate(netuid, |total| { *total = total.saturating_sub(amount); }); @@ -92,7 +99,13 @@ impl Pallet { Error::::InsufficientLiquidity ); - TotalHotkeyAlpha::::mutate(&hotkey, netuid, |v| *v = v.saturating_sub(amount)); + if TotalHotkeyAlpha::::mutate(&hotkey, netuid, |v| { + *v = v.saturating_sub(amount); + *v + }) == 0 + { + TotalHotkeyAlpha::::remove(&hotkey, netuid); + } // Deposit event Self::deposit_event(Event::AlphaBurned(coldkey, hotkey, amount, netuid)); diff --git a/pallets/subtensor/src/tests/migration.rs b/pallets/subtensor/src/tests/migration.rs index 2fac8c5db5..0628127413 100644 --- a/pallets/subtensor/src/tests/migration.rs +++ b/pallets/subtensor/src/tests/migration.rs @@ -438,3 +438,45 @@ fn test_migrate_set_first_emission_block_number() { } }); } + +#[test] +fn test_migrate_remove_zero_total_hotkey_alpha() { + new_test_ext(1).execute_with(|| { + const MIGRATION_NAME: &str = "migrate_remove_zero_total_hotkey_alpha"; + let netuid = 1u16; + + let hotkey_zero = U256::from(100u64); + let hotkey_nonzero = U256::from(101u64); + + // Insert one zero-alpha entry and one non-zero entry + TotalHotkeyAlpha::::insert(hotkey_zero, netuid, 0u64); + TotalHotkeyAlpha::::insert(hotkey_nonzero, netuid, 123u64); + + assert_eq!(TotalHotkeyAlpha::::get(hotkey_zero, netuid), 0u64); + assert_eq!(TotalHotkeyAlpha::::get(hotkey_nonzero, netuid), 123u64); + + assert!( + !HasMigrationRun::::get(MIGRATION_NAME.as_bytes().to_vec()), + "Migration should not have run yet." + ); + + let weight = crate::migrations::migrate_remove_zero_total_hotkey_alpha::migrate_remove_zero_total_hotkey_alpha::(); + + assert!( + HasMigrationRun::::get(MIGRATION_NAME.as_bytes().to_vec()), + "Migration should be marked as run." + ); + + assert!( + !TotalHotkeyAlpha::::contains_key(hotkey_zero, netuid), + "Zero-alpha entry should have been removed." + ); + + assert_eq!(TotalHotkeyAlpha::::get(hotkey_nonzero, netuid), 123u64); + + assert!( + !weight.is_zero(), + "Migration weight should be non-zero." + ); + }); +}