Merge pull request #1196 from opentensor/fix/swap-pending-childkeys-on-hk-swap

Fix/swap pending childkeys on hk swap

Author: sam0x17

pallets/subtensor/src/swap/swap_hotkey.rs

@@ -364,6 +364,7 @@ impl<T: Config> Pallet<T> {
             ChildKeys::<T>::remove(old_hotkey, netuid);
             // Insert the same child entries for the new hotkey
             ChildKeys::<T>::insert(new_hotkey, netuid, my_children.clone());
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
             for (_, child_key_i) in my_children {
                 // For each child, update their parent list
                 let mut child_parents: Vec<(u64, T::AccountId)> =
@@ -376,6 +377,7 @@ impl<T: Config> Pallet<T> {
                 }
                 // Update the child's parent list
                 ParentKeys::<T>::insert(child_key_i, netuid, child_parents);
+                weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
             }
         }
 
@@ -388,6 +390,7 @@ impl<T: Config> Pallet<T> {
             ParentKeys::<T>::remove(old_hotkey, netuid);
             // Insert the same parent entries for the new hotkey
             ParentKeys::<T>::insert(new_hotkey, netuid, parents.clone());
+            weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
             for (_, parent_key_i) in parents {
                 // For each parent, update their children list
                 let mut parent_children: Vec<(u64, T::AccountId)> =
@@ -400,11 +403,39 @@ impl<T: Config> Pallet<T> {
                 }
                 // Update the parent's children list
                 ChildKeys::<T>::insert(parent_key_i, netuid, parent_children);
+                weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 1));
             }
         }
 
-        // 14. Swap Stake Delta for all coldkeys.
-        // DEPRECATED
+        // 14. Swap PendingChildKeys.
+        // PendingChildKeys( netuid, parent ) --> Vec<(proportion,child), cool_down_block>
+        for netuid in Self::get_all_subnet_netuids() {
+            weight.saturating_accrue(T::DbWeight::get().reads(1));
+            if PendingChildKeys::<T>::contains_key(netuid, old_hotkey) {
+                let (children, cool_down_block) = PendingChildKeys::<T>::get(netuid, old_hotkey);
+                PendingChildKeys::<T>::remove(netuid, old_hotkey);
+                PendingChildKeys::<T>::insert(netuid, new_hotkey, (children, cool_down_block));
+                weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
+            }
+
+            // Also check for others with our hotkey as a child
+            for (hotkey, (children, cool_down_block)) in PendingChildKeys::<T>::iter_prefix(netuid)
+            {
+                weight.saturating_accrue(T::DbWeight::get().reads(1));
+
+                if let Some(potential_idx) =
+                    children.iter().position(|(_, child)| *child == *old_hotkey)
+                {
+                    let mut new_children = children.clone();
+                    let entry_to_remove = new_children.remove(potential_idx);
+                    new_children.push((entry_to_remove.0, new_hotkey.clone())); // Keep the proportion.
+
+                    PendingChildKeys::<T>::remove(netuid, hotkey.clone());
+                    PendingChildKeys::<T>::insert(netuid, hotkey, (new_children, cool_down_block));
+                    weight.saturating_accrue(T::DbWeight::get().reads_writes(1, 2));
+                }
+            }
+        }
 
         // Return successful after swapping all the relevant terms.
         Ok(())

pallets/subtensor/src/tests/swap_hotkey.rs

@@ -1257,12 +1257,18 @@ fn test_swap_parent_hotkey_childkey_maps() {
         let parent_old = U256::from(1);
         let coldkey = U256::from(2);
         let child = U256::from(3);
+        let child_other = U256::from(4);
         let parent_new = U256::from(4);
         add_network(netuid, 1, 0);
         SubtensorModule::create_account_if_non_existent(&coldkey, &parent_old);
 
         // Set child and verify state maps
         mock_set_children(&coldkey, &parent_old, netuid, &[(u64::MAX, child)]);
+        // Wait rate limit
+        step_rate_limit(&TransactionType::SetChildren, netuid);
+        // Schedule some pending child keys.
+        mock_schedule_children(&coldkey, &parent_old, netuid, &[(u64::MAX, child_other)]);
+
         assert_eq!(
             ParentKeys::<Test>::get(child, netuid),
             vec![(u64::MAX, parent_old)]
@@ -1271,6 +1277,8 @@ fn test_swap_parent_hotkey_childkey_maps() {
             ChildKeys::<Test>::get(parent_old, netuid),
             vec![(u64::MAX, child)]
         );
+        let existing_pending_child_keys = PendingChildKeys::<Test>::get(netuid, parent_old);
+        assert_eq!(existing_pending_child_keys.0, vec![(u64::MAX, child_other)]);
 
         // Swap
         let mut weight = Weight::zero();
@@ -1290,6 +1298,10 @@ fn test_swap_parent_hotkey_childkey_maps() {
             ChildKeys::<Test>::get(parent_new, netuid),
             vec![(u64::MAX, child)]
         );
+        assert_eq!(
+            PendingChildKeys::<Test>::get(netuid, parent_new),
+            existing_pending_child_keys // Entry under new hotkey.
+        );
     })
 }
 
@@ -1307,6 +1319,10 @@ fn test_swap_child_hotkey_childkey_maps() {
 
         // Set child and verify state maps
         mock_set_children(&coldkey, &parent, netuid, &[(u64::MAX, child_old)]);
+        // Wait rate limit
+        step_rate_limit(&TransactionType::SetChildren, netuid);
+        // Schedule some pending child keys.
+        mock_schedule_children(&coldkey, &parent, netuid, &[(u64::MAX, child_old)]);
 
         assert_eq!(
             ParentKeys::<Test>::get(child_old, netuid),
@@ -1316,6 +1332,8 @@ fn test_swap_child_hotkey_childkey_maps() {
             ChildKeys::<Test>::get(parent, netuid),
             vec![(u64::MAX, child_old)]
         );
+        let existing_pending_child_keys = PendingChildKeys::<Test>::get(netuid, parent);
+        assert_eq!(existing_pending_child_keys.0, vec![(u64::MAX, child_old)]);
 
         // Swap
         let mut weight = Weight::zero();
@@ -1335,5 +1353,9 @@ fn test_swap_child_hotkey_childkey_maps() {
             ChildKeys::<Test>::get(parent, netuid),
             vec![(u64::MAX, child_new)]
         );
+        assert_eq!(
+            PendingChildKeys::<Test>::get(netuid, parent),
+            (vec![(u64::MAX, child_new)], existing_pending_child_keys.1) // Same cooldown block.
+        );
     })
 }