From cd474044811db0f4087dafe69f864faffb964dbf Mon Sep 17 00:00:00 2001 From: Samuel Dare Date: Mon, 15 Jul 2024 11:58:29 +0400 Subject: [PATCH] fix: total coldkey stakes --- pallets/subtensor/src/benchmarks.rs | 54 +++++------ pallets/subtensor/src/lib.rs | 15 ++- pallets/subtensor/src/swap.rs | 45 ++++++++- pallets/subtensor/tests/swap.rs | 137 ++++++++++++++++++++++++++++ 4 files changed, 222 insertions(+), 29 deletions(-) diff --git a/pallets/subtensor/src/benchmarks.rs b/pallets/subtensor/src/benchmarks.rs index 80a375d10..03e087a92 100644 --- a/pallets/subtensor/src/benchmarks.rs +++ b/pallets/subtensor/src/benchmarks.rs @@ -314,33 +314,33 @@ benchmarks! { assert_ok!(Subtensor::::register_network(RawOrigin::Signed(coldkey.clone()).into())); }: dissolve_network(RawOrigin::Signed(coldkey), 1) - swap_hotkey { - let seed: u32 = 1; - let coldkey: T::AccountId = account("Alice", 0, seed); - let old_hotkey: T::AccountId = account("Bob", 0, seed); - let new_hotkey: T::AccountId = account("Charlie", 0, seed); - - let netuid = 1u16; - Subtensor::::init_new_network(netuid, 100); - Subtensor::::set_min_burn(netuid, 1); - Subtensor::::set_max_burn(netuid, 1); - Subtensor::::set_target_registrations_per_interval(netuid, 256); - Subtensor::::set_max_registrations_per_block(netuid, 256); - - Subtensor::::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64); - assert_ok!(Subtensor::::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, old_hotkey.clone())); - assert_ok!(Subtensor::::become_delegate(RawOrigin::Signed(coldkey.clone()).into(), old_hotkey.clone())); - - let max_uids = Subtensor::::get_max_allowed_uids(netuid) as u32; - for i in 0..max_uids - 1 { - let coldkey: T::AccountId = account("Axon", 0, i); - let hotkey: T::AccountId = account("Hotkey", 0, i); - - Subtensor::::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64); - assert_ok!(Subtensor::::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey)); - assert_ok!(Subtensor::::add_stake(RawOrigin::Signed(coldkey).into(), old_hotkey.clone(), 1_000_000_000)); - } - }: _(RawOrigin::Signed(coldkey), old_hotkey, new_hotkey) + // swap_hotkey { + // let seed: u32 = 1; + // let coldkey: T::AccountId = account("Alice", 0, seed); + // let old_hotkey: T::AccountId = account("Bob", 0, seed); + // let new_hotkey: T::AccountId = account("Charlie", 0, seed); + + // let netuid = 1u16; + // Subtensor::::init_new_network(netuid, 100); + // Subtensor::::set_min_burn(netuid, 1); + // Subtensor::::set_max_burn(netuid, 1); + // Subtensor::::set_target_registrations_per_interval(netuid, 256); + // Subtensor::::set_max_registrations_per_block(netuid, 256); + + // Subtensor::::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64); + // assert_ok!(Subtensor::::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, old_hotkey.clone())); + // assert_ok!(Subtensor::::become_delegate(RawOrigin::Signed(coldkey.clone()).into(), old_hotkey.clone())); + + // let max_uids = Subtensor::::get_max_allowed_uids(netuid) as u32; + // for i in 0..max_uids - 1 { + // let coldkey: T::AccountId = account("Axon", 0, i); + // let hotkey: T::AccountId = account("Hotkey", 0, i); + + // Subtensor::::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64); + // assert_ok!(Subtensor::::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey)); + // assert_ok!(Subtensor::::add_stake(RawOrigin::Signed(coldkey).into(), old_hotkey.clone(), 1_000_000_000)); + // } + // }: _(RawOrigin::Signed(coldkey), old_hotkey, new_hotkey) commit_weights { let tempo: u16 = 1; diff --git a/pallets/subtensor/src/lib.rs b/pallets/subtensor/src/lib.rs index 83b6a4005..d42cb495e 100644 --- a/pallets/subtensor/src/lib.rs +++ b/pallets/subtensor/src/lib.rs @@ -2254,7 +2254,7 @@ pub mod pallet { Self::user_remove_network(origin, netuid) } - /// Sets values for liquid alpha + /// fix delegate stake after coldkey swap #[pallet::call_index(64)] #[pallet::weight((0, DispatchClass::Operational, Pays::No))] pub fn sudo_hotfix_swap_coldkey_delegates( @@ -2266,6 +2266,19 @@ pub mod pallet { Self::swap_hotfix(&old_coldkey, &new_coldkey); Ok(()) } + + /// fix total coldkey stake after hotfix + #[pallet::call_index(68)] + #[pallet::weight((0, DispatchClass::Operational, Pays::No))] + pub fn sudo_hotfix_swap_coldkey_total_stake( + origin: OriginFor, + old_coldkey: T::AccountId, + new_coldkey: T::AccountId, + ) -> DispatchResult { + ensure_root(origin)?; + Self::fix_total_coldkey_stake_after_hotfix(&old_coldkey, &new_coldkey); + Ok(()) + } } // ---- Subtensor helper functions. diff --git a/pallets/subtensor/src/swap.rs b/pallets/subtensor/src/swap.rs index 50e3ed944..6ba65c4d1 100644 --- a/pallets/subtensor/src/swap.rs +++ b/pallets/subtensor/src/swap.rs @@ -973,7 +973,7 @@ impl Pallet { } pub fn swap_hotfix(old_coldkey: &T::AccountId, new_coldkey: &T::AccountId) { - let weight = T::DbWeight::get().reads_writes(2, 1); + let _weight = T::DbWeight::get().reads_writes(2, 1); let staking_hotkeys = StakingHotkeys::::get(old_coldkey); for staking_hotkey in staking_hotkeys { if Stake::::contains_key(staking_hotkey.clone(), old_coldkey) { @@ -1001,6 +1001,49 @@ impl Pallet { StakingHotkeys::::remove(old_coldkey); } + pub fn fix_total_coldkey_stake_after_hotfix( + old_coldkey: &T::AccountId, + new_coldkey: &T::AccountId, + ) { + log::info!( + "Starting fix_total_coldkey_stake_after_hotfix: old_coldkey: {:?}, new_coldkey: {:?}", + old_coldkey, + new_coldkey + ); + + let staking_hotkeys = StakingHotkeys::::get(new_coldkey); + let mut total_stake_for_new_coldkey: u64 = 0; + + for hotkey in &staking_hotkeys { + let stake = Stake::::get(hotkey, new_coldkey); + total_stake_for_new_coldkey = total_stake_for_new_coldkey.saturating_add(stake); + + log::info!( + "Counted stake for hotkey {:?}: new_coldkey: {}", + hotkey, + stake + ); + } + + // Set the old coldkey stake to 0 as all stakes should have been moved + TotalColdkeyStake::::insert(old_coldkey, 0); + + // Set the new coldkey stake to the total we just calculated + TotalColdkeyStake::::insert(new_coldkey, total_stake_for_new_coldkey); + + // Read the values from storage for logging + let old_coldkey_stake = TotalColdkeyStake::::get(old_coldkey); + let new_coldkey_stake = TotalColdkeyStake::::get(new_coldkey); + + log::info!( + "Updated TotalColdkeyStake - Old coldkey: {}, New coldkey: {}", + old_coldkey_stake, + new_coldkey_stake + ); + + log::info!("Completed fix_total_coldkey_stake_after_hotfix"); + } + /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey. /// /// # Arguments diff --git a/pallets/subtensor/tests/swap.rs b/pallets/subtensor/tests/swap.rs index 6acf0b820..e1bbe78cd 100644 --- a/pallets/subtensor/tests/swap.rs +++ b/pallets/subtensor/tests/swap.rs @@ -1977,3 +1977,140 @@ fn test_sudo_hotfix_swap_coldkey_delegates_with_broken_stake() { assert_eq!(Stake::::get(h2, new_coldkey), 100); }); } + +// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test swap -- test_coldkey_swap_with_hotfix_and_total_stake_fix --exact --nocapture +#[test] +fn test_coldkey_swap_with_hotfix_and_total_stake_fix() { + new_test_ext(1).execute_with(|| { + let new_coldkey = U256::from(0); + let old_coldkey = U256::from(4); + let h1 = U256::from(5); + let h2 = U256::from(6); + let h3 = U256::from(7); + let stake_amount = 100u64; + + // Setup initial state to simulate after a broken coldkey swap + Stake::::insert(h3, old_coldkey, stake_amount); + Stake::::insert(h2, old_coldkey, stake_amount); + Stake::::insert(h1, new_coldkey, stake_amount); + assert_eq!(Stake::::get(h3, old_coldkey), stake_amount); + assert_eq!(Stake::::get(h2, old_coldkey), stake_amount); + assert_eq!(Stake::::get(h1, new_coldkey), stake_amount); + + StakingHotkeys::::insert(new_coldkey, vec![h1, h2]); + StakingHotkeys::::insert(old_coldkey, vec![h3, h2]); + + // Set up total coldkey stake (this would be incorrect after a broken swap) + TotalColdkeyStake::::insert(old_coldkey, 2 * stake_amount); + TotalColdkeyStake::::insert(new_coldkey, stake_amount); + + // Log initial state + log::info!("Initial state (simulating after broken swap):"); + log::info!( + "Old coldkey stake: {}", + SubtensorModule::get_total_stake_for_coldkey(&old_coldkey) + ); + log::info!( + "New coldkey stake: {}", + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey) + ); + log::info!( + "Old coldkey staking hotkeys: {:?}", + StakingHotkeys::::get(old_coldkey) + ); + log::info!( + "New coldkey staking hotkeys: {:?}", + StakingHotkeys::::get(new_coldkey) + ); + log::info!("h1 stake (new): {}", Stake::::get(h1, new_coldkey)); + log::info!("h2 stake (old): {}", Stake::::get(h2, old_coldkey)); + log::info!("h2 stake (new): {}", Stake::::get(h2, new_coldkey)); + log::info!("h3 stake (old): {}", Stake::::get(h3, old_coldkey)); + + // Apply the hotfix for StakingHotkeys using the sudo extrinsic + assert_ok!(SubtensorModule::sudo_hotfix_swap_coldkey_delegates( + RuntimeOrigin::root(), + old_coldkey, + new_coldkey + )); + + // Log state after hotfix + log::info!("After hotfix:"); + log::info!( + "Old coldkey stake: {}", + SubtensorModule::get_total_stake_for_coldkey(&old_coldkey) + ); + log::info!( + "New coldkey stake: {}", + SubtensorModule::get_total_stake_for_coldkey(&new_coldkey) + ); + log::info!( + "Old coldkey staking hotkeys: {:?}", + StakingHotkeys::::get(old_coldkey) + ); + log::info!( + "New coldkey staking hotkeys: {:?}", + StakingHotkeys::::get(new_coldkey) + ); + log::info!("h1 stake (new): {}", Stake::::get(h1, new_coldkey)); + log::info!("h2 stake (old): {}", Stake::::get(h2, old_coldkey)); + log::info!("h2 stake (new): {}", Stake::::get(h2, new_coldkey)); + log::info!("h3 stake (new): {}", Stake::::get(h3, new_coldkey)); + + // Fix the TotalColdkeyStake using the sudo extrinsic + assert_ok!(SubtensorModule::sudo_hotfix_swap_coldkey_total_stake( + RuntimeOrigin::root(), + old_coldkey, + new_coldkey + )); + + // Log final state + log::info!("Final state:"); + let final_old_coldkey_stake = SubtensorModule::get_total_stake_for_coldkey(&old_coldkey); + let final_new_coldkey_stake = SubtensorModule::get_total_stake_for_coldkey(&new_coldkey); + log::info!("Old coldkey stake: {}", final_old_coldkey_stake); + log::info!("New coldkey stake: {}", final_new_coldkey_stake); + log::info!( + "Old coldkey staking hotkeys: {:?}", + StakingHotkeys::::get(old_coldkey) + ); + log::info!( + "New coldkey staking hotkeys: {:?}", + StakingHotkeys::::get(new_coldkey) + ); + log::info!("h1 stake (new): {}", Stake::::get(h1, new_coldkey)); + log::info!("h2 stake (old): {}", Stake::::get(h2, old_coldkey)); + log::info!("h2 stake (new): {}", Stake::::get(h2, new_coldkey)); + log::info!("h3 stake (new): {}", Stake::::get(h3, new_coldkey)); + + // Verify final state + assert_eq!( + final_old_coldkey_stake, 0, + "Expected old coldkey stake to be 0, but got {}", + final_old_coldkey_stake + ); + assert_eq!( + final_new_coldkey_stake, + 3 * stake_amount, + "Expected total stake for new coldkey to be {}, but got {}", + 3 * stake_amount, + final_new_coldkey_stake + ); + assert_eq!(StakingHotkeys::::get(old_coldkey), Vec::::new()); + assert_eq!(StakingHotkeys::::get(new_coldkey), vec![h1, h2, h3]); + + // Verify individual stakes + assert_eq!(Stake::::get(h1, new_coldkey), stake_amount); + assert_eq!( + Stake::::get(h2, new_coldkey), + stake_amount, + "Expected h2 stake to be {}, but got {}", + stake_amount, + Stake::::get(h2, new_coldkey) + ); + assert_eq!(Stake::::get(h3, new_coldkey), stake_amount); + assert_eq!(Stake::::get(h1, old_coldkey), 0); + assert_eq!(Stake::::get(h2, old_coldkey), 0); + assert_eq!(Stake::::get(h3, old_coldkey), 0); + }); +}