diff --git a/pallets/subtensor/src/staking/remove_stake.rs b/pallets/subtensor/src/staking/remove_stake.rs index 0b36dbf224..8ca9ccc2b4 100644 --- a/pallets/subtensor/src/staking/remove_stake.rs +++ b/pallets/subtensor/src/staking/remove_stake.rs @@ -666,7 +666,7 @@ impl Pallet { let mut hotkeys_in_subnet: Vec = Vec::new(); let mut coldkeys = BTreeSet::::new(); - let mut last_hot = None; + let mut last_completed_key = last_key.clone(); let iter = match last_key { Some(key) => TotalHotkeyAlpha::::iter_from(key), @@ -725,7 +725,8 @@ impl Pallet { // 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 Pallet { for (cold, value) in coldkey_value_vec { stakers.push((hot.clone(), cold, value)); } - last_hot = Some(hot.clone()); + last_completed_key = Some(TotalHotkeyAlpha::::hashed_key_for(&hot, netuid)); } } @@ -825,10 +826,7 @@ impl Pallet { // 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::::hashed_key_for(&hot, netuid)), - ) + (read_all, last_completed_key) } pub fn destroy_alpha_in_out_stakes_clean_alpha( diff --git a/pallets/subtensor/src/tests/destroy_alpha_tests.rs b/pallets/subtensor/src/tests/destroy_alpha_tests.rs index d7186f1f16..aef09d009e 100644 --- a/pallets/subtensor/src/tests/destroy_alpha_tests.rs +++ b/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(), + ::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::::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 = ::DbWeight::get() + .reads(3) + .saturating_add(::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(), + ::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::::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 = ::DbWeight::get() + .reads(14) + .saturating_add(::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 = ::DbWeight::get() + .reads(3) + .saturating_add(::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" + ); + }); +}