Merge pull request #636 from opentensor/total_coldkey_migration

Total coldkey migration

Author: unconst

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;

pallets/subtensor/src/lib.rs

@@ -1419,7 +1419,9 @@ pub mod pallet {
                 // Populate OwnedHotkeys map for coldkey swap. Doesn't update storage vesion.
                 .saturating_add(migration::migrate_populate_owned::<T>())
                 // Populate StakingHotkeys map for coldkey swap. Doesn't update storage vesion.
-                .saturating_add(migration::migrate_populate_staking_hotkeys::<T>());
+                .saturating_add(migration::migrate_populate_staking_hotkeys::<T>())
+                // Fix total coldkey stake.
+                .saturating_add(migration::migrate_fix_total_coldkey_stake::<T>());
 
             weight
         }
@@ -2061,17 +2063,17 @@ pub mod pallet {
         }
 
         /// The extrinsic for user to change its hotkey
-        #[pallet::call_index(70)]
-        #[pallet::weight((Weight::from_parts(1_940_000_000, 0)
-		.saturating_add(T::DbWeight::get().reads(272))
-		.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
-        pub fn swap_hotkey(
-            origin: OriginFor<T>,
-            hotkey: T::AccountId,
-            new_hotkey: T::AccountId,
-        ) -> DispatchResultWithPostInfo {
-            Self::do_swap_hotkey(origin, &hotkey, &new_hotkey)
-        }
+        ///#[pallet::call_index(70)]
+        ///#[pallet::weight((Weight::from_parts(1_940_000_000, 0)
+        ///.saturating_add(T::DbWeight::get().reads(272))
+        ///.saturating_add(T::DbWeight::get().writes(527)), DispatchClass::Operational, Pays::No))]
+        ///pub fn swap_hotkey(
+        ///    origin: OriginFor<T>,
+        ///    hotkey: T::AccountId,
+        ///    new_hotkey: T::AccountId,
+        ///) -> DispatchResultWithPostInfo {
+        ///    Self::do_swap_hotkey(origin, &hotkey, &new_hotkey)
+        ///}
 
         /// The extrinsic for user to change the coldkey associated with their account.
         ///
@@ -2253,6 +2255,17 @@ pub mod pallet {
         pub fn dissolve_network(origin: OriginFor<T>, netuid: u16) -> DispatchResult {
             Self::user_remove_network(origin, netuid)
         }
+
+        /// Sets values for liquid alpha
+        #[pallet::call_index(64)]
+        #[pallet::weight((0, DispatchClass::Operational, Pays::No))]
+        pub fn sudo_hotfix_swap_coldkey_delegates(
+            _origin: OriginFor<T>,
+            _old_coldkey: T::AccountId,
+            _new_coldkey: T::AccountId,
+        ) -> DispatchResult {
+            Ok(())
+        }
     }
 
     // ---- Subtensor helper functions.

pallets/subtensor/src/migration.rs

@@ -24,6 +24,62 @@ pub mod deprecated_loaded_emission_format {
         StorageMap<Pallet<T>, Identity, u16, Vec<(AccountIdOf<T>, u64)>, OptionQuery>;
 }
 
+/// Migrates and fixes the total coldkey stake.
+///
+/// This function iterates through all staking hotkeys, calculates the total stake for each coldkey,
+/// and updates the `TotalColdkeyStake` storage accordingly. The migration is only performed if the
+/// on-chain storage version is 6.
+///
+/// # Returns
+/// The weight of the migration process.
+pub fn do_migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
+    // Initialize the weight with one read operation.
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Iterate through all staking hotkeys.
+    for (coldkey, hotkey_vec) in StakingHotkeys::<T>::iter() {
+        // Init the zero value.
+        let mut coldkey_stake_sum: u64 = 0;
+        weight = weight.saturating_add(T::DbWeight::get().reads(1));
+
+        // Calculate the total stake for the current coldkey.
+        for hotkey in hotkey_vec {
+            // Cant fail on retrieval.
+            coldkey_stake_sum =
+                coldkey_stake_sum.saturating_add(Stake::<T>::get(hotkey, coldkey.clone()));
+            weight = weight.saturating_add(T::DbWeight::get().reads(1));
+        }
+        // Update the `TotalColdkeyStake` storage with the calculated stake sum.
+        // Cant fail on insert.
+        TotalColdkeyStake::<T>::insert(coldkey.clone(), coldkey_stake_sum);
+        weight = weight.saturating_add(T::DbWeight::get().writes(1));
+    }
+    weight
+}
+// Public migrate function to be called by Lib.rs on upgrade.
+pub fn migrate_fix_total_coldkey_stake<T: Config>() -> Weight {
+    let current_storage_version: u16 = 7;
+    let next_storage_version: u16 = 8;
+
+    // Initialize the weight with one read operation.
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Grab the current on-chain storage version.
+    // Cant fail on retrieval.
+    let onchain_version = Pallet::<T>::on_chain_storage_version();
+
+    // Only run this migration on storage version 6.
+    if onchain_version == current_storage_version {
+        weight = weight.saturating_add(do_migrate_fix_total_coldkey_stake::<T>());
+        // Cant fail on insert.
+        StorageVersion::new(next_storage_version).put::<Pallet<T>>();
+        weight.saturating_accrue(T::DbWeight::get().writes(1));
+    }
+
+    // Return the migration weight.
+    weight
+}
+
 /// Performs migration to update the total issuance based on the sum of stakes and total balances.
 /// This migration is applicable only if the current storage version is 5, after which it updates the storage version to 6.
 ///

pallets/subtensor/src/swap.rs

@@ -851,7 +851,8 @@ impl<T: Config> Pallet<T> {
             log::info!("Transferring stake for hotkey {:?}: {}", hotkey, stake);
             if stake > 0 {
                 // Insert the stake for the hotkey and new coldkey
-                Stake::<T>::insert(hotkey, new_coldkey, stake);
+                let old_stake = Stake::<T>::get(hotkey, new_coldkey);
+                Stake::<T>::insert(hotkey, new_coldkey, stake.saturating_add(old_stake));
                 total_transferred_stake = total_transferred_stake.saturating_add(stake);
 
                 // Update the owner of the hotkey to the new coldkey
@@ -861,6 +862,52 @@ impl<T: Config> Pallet<T> {
                 weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 2));
             }
         }
+        log::info!(
+            "Starting transfer of delegated stakes for old coldkey: {:?}",
+            old_coldkey
+        );
+
+        for staking_hotkey in StakingHotkeys::<T>::get(old_coldkey) {
+            log::info!("Processing staking hotkey: {:?}", staking_hotkey);
+            if Stake::<T>::contains_key(staking_hotkey.clone(), old_coldkey) {
+                let hotkey = &staking_hotkey;
+                // Retrieve and remove the stake associated with the hotkey and old coldkey
+                let stake: u64 = Stake::<T>::get(hotkey, old_coldkey);
+                Stake::<T>::remove(hotkey, old_coldkey);
+                log::info!(
+                    "Transferring delegated stake for hotkey {:?}: {}",
+                    hotkey,
+                    stake
+                );
+                if stake > 0 {
+                    // Insert the stake for the hotkey and new coldkey
+                    let old_stake = Stake::<T>::get(hotkey, new_coldkey);
+                    Stake::<T>::insert(hotkey, new_coldkey, stake.saturating_add(old_stake));
+                    total_transferred_stake = total_transferred_stake.saturating_add(stake);
+                    log::info!(
+                        "Updated stake for hotkey {:?} under new coldkey {:?}: {}",
+                        hotkey,
+                        new_coldkey,
+                        stake.saturating_add(old_stake)
+                    );
+
+                    // Update the transaction weight
+                    weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 1));
+                }
+            } else {
+                log::info!(
+                    "No stake found for staking hotkey {:?} under old coldkey {:?}",
+                    staking_hotkey,
+                    old_coldkey
+                );
+                weight.saturating_accrue(T::DbWeight::get().reads(1));
+            }
+        }
+
+        log::info!(
+            "Completed transfer of delegated stakes for old coldkey: {:?}",
+            old_coldkey
+        );
 
         // Log the total transferred stake
         log::info!("Total transferred stake: {}", total_transferred_stake);
@@ -888,13 +935,30 @@ impl<T: Config> Pallet<T> {
         }
 
         // Update the list of owned hotkeys for both old and new coldkeys
+
+        let mut new_owned_hotkeys = OwnedHotkeys::<T>::get(new_coldkey);
+        for hotkey in old_owned_hotkeys {
+            if !new_owned_hotkeys.contains(&hotkey) {
+                new_owned_hotkeys.push(hotkey);
+            }
+        }
+
+        OwnedHotkeys::<T>::insert(new_coldkey, new_owned_hotkeys);
         OwnedHotkeys::<T>::remove(old_coldkey);
-        OwnedHotkeys::<T>::insert(new_coldkey, old_owned_hotkeys);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
 
         // Update the staking hotkeys for both old and new coldkeys
-        let staking_hotkeys: Vec<T::AccountId> = StakingHotkeys::<T>::take(old_coldkey);
-        StakingHotkeys::<T>::insert(new_coldkey, staking_hotkeys);
+        let staking_hotkeys: Vec<T::AccountId> = StakingHotkeys::<T>::get(old_coldkey);
+
+        let mut existing_staking_hotkeys = StakingHotkeys::<T>::get(new_coldkey);
+        for hotkey in staking_hotkeys {
+            if !existing_staking_hotkeys.contains(&hotkey) {
+                existing_staking_hotkeys.push(hotkey);
+            }
+        }
+
+        StakingHotkeys::<T>::remove(old_coldkey);
+        StakingHotkeys::<T>::insert(new_coldkey, existing_staking_hotkeys);
         weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
 
         // Log the total stake of old and new coldkeys after the swap
@@ -907,6 +971,7 @@ impl<T: Config> Pallet<T> {
             TotalColdkeyStake::<T>::get(new_coldkey)
         );
     }
+
     /// Swaps the total hotkey-coldkey stakes for the current interval from the old coldkey to the new coldkey.
     ///
     /// # Arguments

pallets/subtensor/tests/migration.rs

@@ -4,6 +4,7 @@ mod mock;
 use frame_support::assert_ok;
 use frame_system::Config;
 use mock::*;
+use pallet_subtensor::*;
 use sp_core::U256;
 
 #[test]
@@ -276,3 +277,73 @@ fn test_migration_delete_subnet_21() {
         assert!(!SubtensorModule::if_subnet_exist(21));
     })
 }
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0);
+        TotalColdkeyStake::<Test>::insert(coldkey, 0);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_value_already_in_total --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_value_already_in_total() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0);
+        TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_no_entry --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_no_entry() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(3), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 30000);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_no_entry_in_hotkeys --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_no_entry_in_hotkeys() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0);
+        TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 0);
+    })
+}
+
+// SKIP_WASM_BUILD=1 RUST_LOG=info cargo test --test migration -- test_migrate_fix_total_coldkey_stake_one_hotkey_stake_missing --exact --nocapture
+#[test]
+fn test_migrate_fix_total_coldkey_stake_one_hotkey_stake_missing() {
+    new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(0);
+        TotalColdkeyStake::<Test>::insert(coldkey, 100000000);
+        StakingHotkeys::<Test>::insert(coldkey, vec![U256::from(1), U256::from(2), U256::from(3)]);
+        Stake::<Test>::insert(U256::from(1), U256::from(0), 10000);
+        Stake::<Test>::insert(U256::from(2), U256::from(0), 10000);
+        pallet_subtensor::migration::do_migrate_fix_total_coldkey_stake::<Test>();
+        assert_eq!(TotalColdkeyStake::<Test>::get(coldkey), 20000);
+    })
+}