preserve dissolution stake settlement cursor

pallets/subtensor/src/staking/remove_stake.rs

@@ -666,7 +666,7 @@ impl<T: Config> Pallet<T> {
 
         let mut hotkeys_in_subnet: Vec<T::AccountId> = Vec::new();
         let mut coldkeys = BTreeSet::<T::AccountId>::new();
-        let mut last_hot = None;
+        let mut last_completed_key = last_key.clone();
 
         let iter = match last_key {
             Some(key) => TotalHotkeyAlpha::<T>::iter_from(key),
@@ -725,7 +725,8 @@ impl<T: Config> Pallet<T> {
                     // reserve the weight for the add_balance_to_coldkey_account function call later
                     if !weight_meter.can_consume(need_to_consume_weight) {
                         inner_read_all = false;
-                        last_hot = Some(hot.clone());
+                        // This hotkey has not been fully processed or queued for payout yet.
+                        // Keep the cursor on the last completed hotkey so the next pass retries it.
                         break;
                     }
                     weight_meter.consume(need_to_consume_weight);
@@ -741,7 +742,7 @@ impl<T: Config> Pallet<T> {
                 for (cold, value) in coldkey_value_vec {
                     stakers.push((hot.clone(), cold, value));
                 }
-                last_hot = Some(hot.clone());
+                last_completed_key = Some(TotalHotkeyAlpha::<T>::hashed_key_for(&hot, netuid));
             }
         }
 
@@ -825,10 +826,7 @@ impl<T: Config> Pallet<T> {
         // ignore the weight for handling the final operation, we must set the correct status for the next run
         status.subnet_distributed_tao = Some(distributed_tao_value_u128);
 
-        (
-            read_all,
-            last_hot.map(|hot| TotalHotkeyAlpha::<T>::hashed_key_for(&hot, netuid)),
-        )
+        (read_all, last_completed_key)
     }
 
     pub fn destroy_alpha_in_out_stakes_clean_alpha(

pallets/subtensor/src/tests/destroy_alpha_tests.rs

@@ -444,3 +444,183 @@ fn test_destroy_alpha_in_out_stakes_settle_stakes_multi_block_total_issuance() {
         );
     });
 }
+
+#[test]
+fn test_destroy_alpha_in_out_stakes_settle_stakes_retries_unpaid_hotkey() {
+    new_test_ext(0).execute_with(|| {
+        let cold_base = U256::from(7000);
+        let hot_base = U256::from(8000);
+        let netuid = add_dynamic_network(&hot_base, &cold_base);
+
+        setup_reserves(netuid, 1_000_000u64.into(), 10_000_000u64.into());
+
+        for index in 1..=3 {
+            let cold = U256::from(cold_base + index);
+            let hot = U256::from(hot_base + index);
+            assert_ok!(SubtensorModule::create_account_if_non_existent(&cold, &hot));
+            add_balance_to_coldkey_account(&cold, 1_000u64.into());
+            assert_ok!(SubtensorModule::stake_into_subnet(
+                &hot,
+                &cold,
+                netuid,
+                1_000u64.into(),
+                <Test as Config>::SwapInterface::max_price(),
+                false,
+            ));
+        }
+
+        let mut status = dissolve_cleanup_status(netuid);
+        let mut meter = WeightMeter::with_limit(Weight::from_parts(u64::MAX, u64::MAX));
+        assert!(run_resumable_netuid_cleanup_with_status(
+            netuid,
+            &mut meter,
+            &mut status,
+            |netuid, weight_meter, last_key, status| {
+                SubtensorModule::destroy_alpha_in_out_stakes_get_total_alpha_value(
+                    netuid,
+                    weight_meter,
+                    last_key,
+                    status,
+                )
+            },
+        ));
+        status.subnet_distributed_tao = Some(0);
+
+        let first_hot = TotalHotkeyAlpha::<Test>::iter()
+            .find_map(|(hot, this_netuid, _)| (this_netuid == netuid).then_some(hot))
+            .expect("staked hotkey should exist");
+        let first_cold = SubtensorModule::alpha_iter_single_prefix(&first_hot)
+            .find_map(|(cold, this_netuid, _)| (this_netuid == netuid).then_some(cold))
+            .expect("staked coldkey should exist");
+        let first_balance_before = SubtensorModule::get_coldkey_balance(&first_cold);
+
+        let tight = <Test as frame_system::Config>::DbWeight::get()
+            .reads(3)
+            .saturating_add(<Test as frame_system::Config>::DbWeight::get().writes(1));
+        let mut tight_meter = WeightMeter::with_limit(tight);
+        let (done, new_key) = SubtensorModule::destroy_alpha_in_out_stakes_settle_stakes(
+            netuid,
+            &mut tight_meter,
+            None,
+            &mut status,
+        );
+
+        assert!(!done);
+        assert_eq!(
+            new_key, None,
+            "cursor must not advance past a hotkey that was read but not paid"
+        );
+        assert_eq!(
+            SubtensorModule::get_coldkey_balance(&first_cold),
+            first_balance_before,
+            "tight pass should not pay without reserving full weight"
+        );
+
+        let mut full_meter = WeightMeter::with_limit(Weight::from_parts(u64::MAX, u64::MAX));
+        let (done, _) = SubtensorModule::destroy_alpha_in_out_stakes_settle_stakes(
+            netuid,
+            &mut full_meter,
+            new_key,
+            &mut status,
+        );
+
+        assert!(done);
+        assert!(
+            SubtensorModule::get_coldkey_balance(&first_cold) > first_balance_before,
+            "retry should pay the hotkey that could not fit in the previous pass"
+        );
+    });
+}
+
+#[test]
+fn test_destroy_alpha_in_out_stakes_settle_stakes_keeps_previous_cursor() {
+    new_test_ext(0).execute_with(|| {
+        let cold_base = U256::from(9000);
+        let hot_base = U256::from(10000);
+        let netuid = add_dynamic_network(&hot_base, &cold_base);
+
+        setup_reserves(netuid, 1_000_000u64.into(), 10_000_000u64.into());
+
+        for index in 1..=3 {
+            let cold = U256::from(cold_base + index);
+            let hot = U256::from(hot_base + index);
+            assert_ok!(SubtensorModule::create_account_if_non_existent(&cold, &hot));
+            add_balance_to_coldkey_account(&cold, 1_000u64.into());
+            assert_ok!(SubtensorModule::stake_into_subnet(
+                &hot,
+                &cold,
+                netuid,
+                1_000u64.into(),
+                <Test as Config>::SwapInterface::max_price(),
+                false,
+            ));
+        }
+
+        let mut status = dissolve_cleanup_status(netuid);
+        let mut meter = WeightMeter::with_limit(Weight::from_parts(u64::MAX, u64::MAX));
+        assert!(run_resumable_netuid_cleanup_with_status(
+            netuid,
+            &mut meter,
+            &mut status,
+            |netuid, weight_meter, last_key, status| {
+                SubtensorModule::destroy_alpha_in_out_stakes_get_total_alpha_value(
+                    netuid,
+                    weight_meter,
+                    last_key,
+                    status,
+                )
+            },
+        ));
+        status.subnet_distributed_tao = Some(0);
+
+        let first_hot = TotalHotkeyAlpha::<Test>::iter()
+            .find_map(|(hot, this_netuid, _)| (this_netuid == netuid).then_some(hot))
+            .expect("staked hotkey should exist");
+        let first_cold = SubtensorModule::alpha_iter_single_prefix(&first_hot)
+            .find_map(|(cold, this_netuid, _)| (this_netuid == netuid).then_some(cold))
+            .expect("staked coldkey should exist");
+        let first_balance_before = SubtensorModule::get_coldkey_balance(&first_cold);
+
+        let one_hotkey = <Test as frame_system::Config>::DbWeight::get()
+            .reads(14)
+            .saturating_add(<Test as frame_system::Config>::DbWeight::get().writes(4));
+        let mut first_meter = WeightMeter::with_limit(one_hotkey);
+        let (done, previous_key) = SubtensorModule::destroy_alpha_in_out_stakes_settle_stakes(
+            netuid,
+            &mut first_meter,
+            None,
+            &mut status,
+        );
+        let previous_key = previous_key.expect("first pass should complete one hotkey");
+
+        assert!(!done);
+        assert!(
+            SubtensorModule::get_coldkey_balance(&first_cold) > first_balance_before,
+            "first pass should pay the completed hotkey"
+        );
+        let first_balance_after = SubtensorModule::get_coldkey_balance(&first_cold);
+
+        let tight = <Test as frame_system::Config>::DbWeight::get()
+            .reads(3)
+            .saturating_add(<Test as frame_system::Config>::DbWeight::get().writes(1));
+        let mut tight_meter = WeightMeter::with_limit(tight);
+        let (done, retry_key) = SubtensorModule::destroy_alpha_in_out_stakes_settle_stakes(
+            netuid,
+            &mut tight_meter,
+            Some(previous_key.clone()),
+            &mut status,
+        );
+
+        assert!(!done);
+        assert_eq!(
+            retry_key,
+            Some(previous_key),
+            "cursor should stay at the previous completed hotkey when no new hotkey fits"
+        );
+        assert_eq!(
+            SubtensorModule::get_coldkey_balance(&first_cold),
+            first_balance_after,
+            "tight retry should not rewind and pay the completed hotkey again"
+        );
+    });
+}