Add migration to populate StakingHotkeys, unstake delegations in do_unstake_all_and_transfer_to_new_coldkey

Author: gztensor

pallets/subtensor/src/lib.rs

@@ -380,6 +380,9 @@ pub mod pallet {
         ValueQuery,
         DefaultAccountTake<T>,
     >;
+    #[pallet::storage] // --- DMAP ( cold ) --> Vec<hot> | Maps coldkey to hotkeys that stake to it
+    pub type StakingHotkeys<T: Config> =
+        StorageMap<_, Blake2_128Concat, T::AccountId, Vec<T::AccountId>, ValueQuery>;
     /// -- ITEM (switches liquid alpha on)
     #[pallet::type_value]
     pub fn DefaultLiquidAlpha<T: Config>() -> bool {
@@ -1225,6 +1228,13 @@ pub mod pallet {
 
                     Stake::<T>::insert(hotkey.clone(), coldkey.clone(), stake);
 
+                    // Update StakingHotkeys map
+                    let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
+                    if !staking_hotkeys.contains(hotkey) {
+                        staking_hotkeys.push(hotkey.clone());
+                        StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
+                    }
+
                     next_uid = next_uid.checked_add(1).expect(
                         "should not have total number of hotkey accounts larger than u16::MAX",
                     );
@@ -1337,7 +1347,9 @@ pub mod pallet {
                 // Doesn't check storage version. TODO: Remove after upgrade
                 .saturating_add(migration::migration5_total_issuance::<T>(false))
                 // Populate OwnedHotkeys map for coldkey swap. Doesn't update storage vesion.
-                .saturating_add(migration::migrate_populate_owned::<T>());
+                .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>());
 
             weight
         }

pallets/subtensor/src/migration.rs

@@ -539,3 +539,66 @@ pub fn migrate_populate_owned<T: Config>() -> Weight {
         Weight::zero()
     }
 }
+
+/// Populate the StakingHotkeys map from Stake map
+pub fn migrate_populate_staking_hotkeys<T: Config>() -> Weight {
+    // Setup migration weight
+    let mut weight = T::DbWeight::get().reads(1);
+    let migration_name = "Populate StakingHotkeys map";
+
+    // Check if this migration is needed (if StakingHotkeys map is empty)
+    let migrate = StakingHotkeys::<T>::iter().next().is_none();
+
+    // Only runs if the migration is needed
+    if migrate {
+        info!(target: LOG_TARGET_1, ">>> Starting Migration: {}", migration_name);
+
+        let mut longest_hotkey_vector: usize = 0;
+        let mut longest_coldkey: Option<T::AccountId> = None;
+        let mut keys_touched: u64 = 0;
+        let mut storage_reads: u64 = 0;
+        let mut storage_writes: u64 = 0;
+
+        // Iterate through all Owner entries
+        Stake::<T>::iter().for_each(|(hotkey, coldkey, _stake)| {
+            storage_reads = storage_reads.saturating_add(1); // Read from Owner storage
+            let mut hotkeys = StakingHotkeys::<T>::get(&coldkey);
+            storage_reads = storage_reads.saturating_add(1); // Read from StakingHotkeys storage
+
+            // Add the hotkey if it's not already in the vector
+            if !hotkeys.contains(&hotkey) {
+                hotkeys.push(hotkey);
+                keys_touched = keys_touched.saturating_add(1);
+
+                // Update longest hotkey vector info
+                if longest_hotkey_vector < hotkeys.len() {
+                    longest_hotkey_vector = hotkeys.len();
+                    longest_coldkey = Some(coldkey.clone());
+                }
+
+                // Update the StakingHotkeys storage
+                StakingHotkeys::<T>::insert(&coldkey, hotkeys);
+                storage_writes = storage_writes.saturating_add(1); // Write to StakingHotkeys storage
+            }
+
+            // Accrue weight for reads and writes
+            weight = weight.saturating_add(T::DbWeight::get().reads_writes(2, 1));
+        });
+
+        // Log migration results
+        info!(
+            target: LOG_TARGET_1,
+            "Migration {} finished. Keys touched: {}, Longest hotkey vector: {}, Storage reads: {}, Storage writes: {}",
+            migration_name, keys_touched, longest_hotkey_vector, storage_reads, storage_writes
+        );
+        if let Some(c) = longest_coldkey {
+            info!(target: LOG_TARGET_1, "Longest hotkey vector is controlled by: {:?}", c);
+        }
+
+        weight
+    } else {
+        info!(target: LOG_TARGET_1, "Migration {} already done!", migration_name);
+        Weight::zero()
+    }
+}
+

pallets/subtensor/src/staking.rs

@@ -569,6 +569,13 @@ impl<T: Config> Pallet<T> {
                 hotkeys.push(hotkey.clone());
                 OwnedHotkeys::<T>::insert(coldkey, hotkeys);
             }
+
+            // Update StakingHotkeys map
+            let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
+            if !staking_hotkeys.contains(hotkey) {
+                staking_hotkeys.push(hotkey.clone());
+                StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
+            }
         }
     }
 
@@ -648,6 +655,13 @@ impl<T: Config> Pallet<T> {
             Stake::<T>::get(hotkey, coldkey).saturating_add(increment),
         );
         TotalStake::<T>::put(TotalStake::<T>::get().saturating_add(increment));
+
+        // Update StakingHotkeys map
+        let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
+        if !staking_hotkeys.contains(hotkey) {
+            staking_hotkeys.push(hotkey.clone());
+            StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
+        }
     }
 
     // Decreases the stake on the cold - hot pairing by the decrement while decreasing other counters.
@@ -668,6 +682,8 @@ impl<T: Config> Pallet<T> {
             Stake::<T>::get(hotkey, coldkey).saturating_sub(decrement),
         );
         TotalStake::<T>::put(TotalStake::<T>::get().saturating_sub(decrement));
+
+        // TODO: Tech debt: Remove StakingHotkeys entry if stake goes to 0
     }
 
     /// Empties the stake associated with a given coldkey-hotkey account pairing.
@@ -693,6 +709,11 @@ impl<T: Config> Pallet<T> {
         TotalStake::<T>::mutate(|stake| *stake = stake.saturating_sub(current_stake));
         TotalIssuance::<T>::mutate(|issuance| *issuance = issuance.saturating_sub(current_stake));
 
+        // Update StakingHotkeys map
+        let mut staking_hotkeys = StakingHotkeys::<T>::get(coldkey);
+        staking_hotkeys.retain(|h| h != hotkey);
+        StakingHotkeys::<T>::insert(coldkey, staking_hotkeys);
+
         current_stake
     }
 
@@ -897,6 +918,25 @@ impl<T: Config> Pallet<T> {
             }
         }
 
+        // Unstake all delegate stake make by this coldkey to non-owned hotkeys
+        let staking_hotkeys = StakingHotkeys::<T>::get(&current_coldkey);
+
+        // iterate over all staking hotkeys.
+        for hotkey in staking_hotkeys {
+            // Get the current stake
+            let current_stake: u64 =
+                Self::get_stake_for_coldkey_and_hotkey(&current_coldkey, &hotkey);
+
+            // Unstake all balance if there's any stake
+            if current_stake > 0 {
+                Self::do_remove_stake(
+                    RawOrigin::Signed(current_coldkey.clone()).into(),
+                    hotkey.clone(),
+                    current_stake,
+                )?;
+            }
+        }
+
         let total_balance = Self::get_coldkey_balance(&current_coldkey);
         log::info!("Total Bank Balance: {:?}", total_balance);
 

pallets/subtensor/src/swap.rs

@@ -272,12 +272,22 @@ impl<T: Config> Pallet<T> {
         for (coldkey, stake_amount) in stakes {
             Stake::<T>::insert(new_hotkey, &coldkey, stake_amount);
             writes = writes.saturating_add(1u64); // One write for insert
+
+            // Update StakingHotkeys map
+            let mut staking_hotkeys = StakingHotkeys::<T>::get(&coldkey);
+            if !staking_hotkeys.contains(new_hotkey) {
+                staking_hotkeys.push(new_hotkey.clone());
+                StakingHotkeys::<T>::insert(coldkey.clone(), staking_hotkeys);
+                writes = writes.saturating_add(1u64); // One write for insert
+            }
         }
 
         // Clear the prefix for the old hotkey after transferring all stakes
         let _ = Stake::<T>::clear_prefix(old_hotkey, stake_count, None);
         writes = writes.saturating_add(1); // One write for insert; // One write for clear_prefix
 
+        // TODO: Remove all entries for old hotkey from StakingHotkeys map
+
         weight.saturating_accrue(T::DbWeight::get().writes(writes));
     }
 
@@ -521,7 +531,12 @@ impl<T: Config> Pallet<T> {
             let stake = Stake::<T>::get(&hotkey, old_coldkey);
             Stake::<T>::remove(&hotkey, old_coldkey);
             Stake::<T>::insert(&hotkey, new_coldkey, stake);
-            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
+
+            // Update StakingHotkeys map
+            let staking_hotkeys = StakingHotkeys::<T>::get(old_coldkey);
+            StakingHotkeys::<T>::insert(new_coldkey.clone(), staking_hotkeys);
+
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(2, 3));
         }
     }
 

pallets/subtensor/tests/swap.rs

@@ -625,7 +625,7 @@ fn test_swap_stake_weight_update() {
         SubtensorModule::swap_stake(&old_hotkey, &new_hotkey, &mut weight);
 
         // Verify the weight update
-        let expected_weight = <Test as frame_system::Config>::DbWeight::get().writes(2);
+        let expected_weight = <Test as frame_system::Config>::DbWeight::get().writes(3);
         assert_eq!(weight, expected_weight);
     });
 }
@@ -1213,7 +1213,7 @@ fn test_swap_stake_for_coldkey() {
         assert_eq!(TotalIssuance::<Test>::get(), stake_amount1 + stake_amount2);
 
         // Verify weight update
-        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(3, 4);
+        let expected_weight = <Test as frame_system::Config>::DbWeight::get().reads_writes(5, 6);
         assert_eq!(weight, expected_weight);
     });
 }