fix: total coldkey stakes

opentensor · Jul 15, 2024 · cd47404 · cd47404
1 parent ca51109
commit cd47404
Show file tree

Hide file tree

Showing 4 changed files with 222 additions and 29 deletions.
diff --git a/pallets/subtensor/src/benchmarks.rs b/pallets/subtensor/src/benchmarks.rs
@@ -314,33 +314,33 @@ benchmarks! {
     assert_ok!(Subtensor::<T>::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::<T>::init_new_network(netuid, 100);
-    Subtensor::<T>::set_min_burn(netuid, 1);
-    Subtensor::<T>::set_max_burn(netuid, 1);
-    Subtensor::<T>::set_target_registrations_per_interval(netuid, 256);
-    Subtensor::<T>::set_max_registrations_per_block(netuid, 256);
-
-    Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64);
-    assert_ok!(Subtensor::<T>::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, old_hotkey.clone()));
-    assert_ok!(Subtensor::<T>::become_delegate(RawOrigin::Signed(coldkey.clone()).into(), old_hotkey.clone()));
-
-    let max_uids = Subtensor::<T>::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::<T>::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64);
-        assert_ok!(Subtensor::<T>::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey));
-        assert_ok!(Subtensor::<T>::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::<T>::init_new_network(netuid, 100);
+  //   Subtensor::<T>::set_min_burn(netuid, 1);
+  //   Subtensor::<T>::set_max_burn(netuid, 1);
+  //   Subtensor::<T>::set_target_registrations_per_interval(netuid, 256);
+  //   Subtensor::<T>::set_max_registrations_per_block(netuid, 256);
+
+  //   Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64);
+  //   assert_ok!(Subtensor::<T>::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, old_hotkey.clone()));
+  //   assert_ok!(Subtensor::<T>::become_delegate(RawOrigin::Signed(coldkey.clone()).into(), old_hotkey.clone()));
+
+  //   let max_uids = Subtensor::<T>::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::<T>::add_balance_to_coldkey_account(&coldkey.clone(), 10_000_000_000u64);
+  //       assert_ok!(Subtensor::<T>::burned_register(RawOrigin::Signed(coldkey.clone()).into(), netuid, hotkey));
+  //       assert_ok!(Subtensor::<T>::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
@@ -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<T>,
+            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
@@ -973,7 +973,7 @@ impl<T: Config> Pallet<T> {
     }
 
     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::<T>::get(old_coldkey);
         for staking_hotkey in staking_hotkeys {
             if Stake::<T>::contains_key(staking_hotkey.clone(), old_coldkey) {
@@ -1001,6 +1001,49 @@ impl<T: Config> Pallet<T> {
         StakingHotkeys::<T>::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::<T>::get(new_coldkey);
+        let mut total_stake_for_new_coldkey: u64 = 0;
+
+        for hotkey in &staking_hotkeys {
+            let stake = Stake::<T>::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::<T>::insert(old_coldkey, 0);
+
+        // Set the new coldkey stake to the total we just calculated
+        TotalColdkeyStake::<T>::insert(new_coldkey, total_stake_for_new_coldkey);
+
+        // Read the values from storage for logging
+        let old_coldkey_stake = TotalColdkeyStake::<T>::get(old_coldkey);
+        let new_coldkey_stake = TotalColdkeyStake::<T>::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
@@ -1977,3 +1977,140 @@ fn test_sudo_hotfix_swap_coldkey_delegates_with_broken_stake() {
         assert_eq!(Stake::<Test>::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::<Test>::insert(h3, old_coldkey, stake_amount);
+        Stake::<Test>::insert(h2, old_coldkey, stake_amount);
+        Stake::<Test>::insert(h1, new_coldkey, stake_amount);
+        assert_eq!(Stake::<Test>::get(h3, old_coldkey), stake_amount);
+        assert_eq!(Stake::<Test>::get(h2, old_coldkey), stake_amount);
+        assert_eq!(Stake::<Test>::get(h1, new_coldkey), stake_amount);
+
+        StakingHotkeys::<Test>::insert(new_coldkey, vec![h1, h2]);
+        StakingHotkeys::<Test>::insert(old_coldkey, vec![h3, h2]);
+
+        // Set up total coldkey stake (this would be incorrect after a broken swap)
+        TotalColdkeyStake::<Test>::insert(old_coldkey, 2 * stake_amount);
+        TotalColdkeyStake::<Test>::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::<Test>::get(old_coldkey)
+        );
+        log::info!(
+            "New coldkey staking hotkeys: {:?}",
+            StakingHotkeys::<Test>::get(new_coldkey)
+        );
+        log::info!("h1 stake (new): {}", Stake::<Test>::get(h1, new_coldkey));
+        log::info!("h2 stake (old): {}", Stake::<Test>::get(h2, old_coldkey));
+        log::info!("h2 stake (new): {}", Stake::<Test>::get(h2, new_coldkey));
+        log::info!("h3 stake (old): {}", Stake::<Test>::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::<Test>::get(old_coldkey)
+        );
+        log::info!(
+            "New coldkey staking hotkeys: {:?}",
+            StakingHotkeys::<Test>::get(new_coldkey)
+        );
+        log::info!("h1 stake (new): {}", Stake::<Test>::get(h1, new_coldkey));
+        log::info!("h2 stake (old): {}", Stake::<Test>::get(h2, old_coldkey));
+        log::info!("h2 stake (new): {}", Stake::<Test>::get(h2, new_coldkey));
+        log::info!("h3 stake (new): {}", Stake::<Test>::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::<Test>::get(old_coldkey)
+        );
+        log::info!(
+            "New coldkey staking hotkeys: {:?}",
+            StakingHotkeys::<Test>::get(new_coldkey)
+        );
+        log::info!("h1 stake (new): {}", Stake::<Test>::get(h1, new_coldkey));
+        log::info!("h2 stake (old): {}", Stake::<Test>::get(h2, old_coldkey));
+        log::info!("h2 stake (new): {}", Stake::<Test>::get(h2, new_coldkey));
+        log::info!("h3 stake (new): {}", Stake::<Test>::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::<Test>::get(old_coldkey), Vec::<U256>::new());
+        assert_eq!(StakingHotkeys::<Test>::get(new_coldkey), vec![h1, h2, h3]);
+
+        // Verify individual stakes
+        assert_eq!(Stake::<Test>::get(h1, new_coldkey), stake_amount);
+        assert_eq!(
+            Stake::<Test>::get(h2, new_coldkey),
+            stake_amount,
+            "Expected h2 stake to be {}, but got {}",
+            stake_amount,
+            Stake::<Test>::get(h2, new_coldkey)
+        );
+        assert_eq!(Stake::<Test>::get(h3, new_coldkey), stake_amount);
+        assert_eq!(Stake::<Test>::get(h1, old_coldkey), 0);
+        assert_eq!(Stake::<Test>::get(h2, old_coldkey), 0);
+        assert_eq!(Stake::<Test>::get(h3, old_coldkey), 0);
+    });
+}