diff --git a/pallets/subtensor/src/staking/claim_root.rs b/pallets/subtensor/src/staking/claim_root.rs index 24a26d154c..9512d8c1a8 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,19 @@ 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| { - claimable.remove(&netuid); - }); - } + 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..220528e454 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)); + }); +} diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs index 919468c1ef..4e5dbc8376 100644 --- a/runtime/src/lib.rs +++ b/runtime/src/lib.rs @@ -241,7 +241,7 @@ pub const VERSION: RuntimeVersion = RuntimeVersion { // `spec_version`, and `authoring_version` are the same between Wasm and native. // This value is set to 100 to notify Polkadot-JS App (https://polkadot.js.org/apps) to use // the compatible custom types. - spec_version: 377, + spec_version: 378, impl_version: 1, apis: RUNTIME_API_VERSIONS, transaction_version: 1,