diff --git a/pallets/subtensor/src/staking/claim_root.rs b/pallets/subtensor/src/staking/claim_root.rs index 24a26d154c..27a46539bc 100644 --- a/pallets/subtensor/src/staking/claim_root.rs +++ b/pallets/subtensor/src/staking/claim_root.rs @@ -1,6 +1,7 @@ use super::*; use frame_support::weights::Weight; use sp_core::Get; +use sp_std::collections::btree_map::BTreeMap; use sp_std::collections::btree_set::BTreeSet; use substrate_fixed::types::I96F32; use subtensor_swap_interface::SwapHandler; @@ -389,14 +390,21 @@ impl Pallet { } /// Claim all root dividends for subnet and remove all associated data. + /// + /// Uses `translate` to iterate entries one at a time without collecting all + /// keys into memory, preventing unbounded memory usage when the number of + /// hotkeys is large. pub fn finalize_all_subnet_root_dividends(netuid: NetUid) { - let hotkeys = RootClaimable::::iter_keys().collect::>(); - - for hotkey in hotkeys.iter() { - RootClaimable::::mutate(hotkey, |claimable| { + RootClaimable::::translate::, _>( + |_hotkey, mut claimable| { claimable.remove(&netuid); - }); - } + if claimable.is_empty() { + None + } else { + Some(claimable) + } + }, + ); let _ = RootClaimed::::clear_prefix((netuid,), u32::MAX, None); } diff --git a/pallets/subtensor/src/tests/claim_root.rs b/pallets/subtensor/src/tests/claim_root.rs index 717f3a5d28..d8d8aee1b9 100644 --- a/pallets/subtensor/src/tests/claim_root.rs +++ b/pallets/subtensor/src/tests/claim_root.rs @@ -2057,3 +2057,104 @@ fn test_claim_root_with_moved_stake() { assert_abs_diff_eq!(bob_stake_diff2, estimated_stake as u64, epsilon = 100u64,); }); } + +/// Regression test for issue #2411: finalize_all_subnet_root_dividends must +/// clean up the target netuid across all hotkeys without collecting all keys +/// into memory, and must preserve claimable entries for other subnets. +#[test] +fn test_finalize_all_subnet_root_dividends_multiple_hotkeys() { + new_test_ext(1).execute_with(|| { + let owner_coldkey = U256::from(5001); + let hotkey1 = U256::from(5002); + let hotkey2 = U256::from(5003); + let hotkey3 = U256::from(5004); + + let netuid_a = add_dynamic_network(&hotkey1, &owner_coldkey); + let netuid_b = add_dynamic_network(&hotkey2, &owner_coldkey); + + // Manually populate RootClaimable for multiple hotkeys across two subnets. + let rate = I96F32::from(42); + for hotkey in [hotkey1, hotkey2, hotkey3] { + RootClaimable::::mutate(&hotkey, |claimable| { + claimable.insert(netuid_a, rate); + claimable.insert(netuid_b, rate); + }); + } + + // Populate some RootClaimed entries for netuid_a. + RootClaimed::::insert((netuid_a, &hotkey1, &owner_coldkey), 100u128); + RootClaimed::::insert((netuid_a, &hotkey2, &owner_coldkey), 200u128); + + // Finalize root dividends for netuid_a only. + SubtensorModule::finalize_all_subnet_root_dividends(netuid_a); + + // netuid_a entries must be removed from all hotkeys. + for hotkey in [hotkey1, hotkey2, hotkey3] { + assert!( + !RootClaimable::::get(hotkey).contains_key(&netuid_a), + "netuid_a should be removed from hotkey {hotkey:?}" + ); + } + + // netuid_b entries must be preserved for all hotkeys. + for hotkey in [hotkey1, hotkey2, hotkey3] { + assert_eq!( + RootClaimable::::get(hotkey).get(&netuid_b), + Some(&rate), + "netuid_b should be preserved for hotkey {hotkey:?}" + ); + } + + // RootClaimed entries for netuid_a must be cleared. + assert!(!RootClaimed::::contains_key(( + netuid_a, + &hotkey1, + &owner_coldkey, + ))); + assert!(!RootClaimed::::contains_key(( + netuid_a, + &hotkey2, + &owner_coldkey, + ))); + }); +} + +/// Test that hotkeys whose BTreeMap becomes empty after cleanup are fully +/// removed from storage (translate returns None). +#[test] +fn test_finalize_all_subnet_root_dividends_removes_empty_entries() { + new_test_ext(1).execute_with(|| { + let owner_coldkey = U256::from(6001); + let hotkey_single = U256::from(6002); + let hotkey_multi = U256::from(6003); + + let netuid = add_dynamic_network(&hotkey_single, &owner_coldkey); + let netuid_other = add_dynamic_network(&hotkey_multi, &owner_coldkey); + + let rate = I96F32::from(10); + + // hotkey_single only has the target netuid. + RootClaimable::::mutate(&hotkey_single, |c| { + c.insert(netuid, rate); + }); + + // hotkey_multi has the target netuid plus another. + RootClaimable::::mutate(&hotkey_multi, |c| { + c.insert(netuid, rate); + c.insert(netuid_other, rate); + }); + + SubtensorModule::finalize_all_subnet_root_dividends(netuid); + + // hotkey_single's entry should be entirely gone (empty map removed). + assert!( + RootClaimable::::get(hotkey_single).is_empty(), + "empty map should be cleaned up" + ); + + // hotkey_multi should still have netuid_other. + let remaining = RootClaimable::::get(hotkey_multi); + assert_eq!(remaining.len(), 1); + assert!(remaining.contains_key(&netuid_other)); + }); +}