Fix root basket claim weight accounting

pallets/subtensor/src/benchmarks/benchmarks.rs

@@ -1962,102 +1962,67 @@ mod pallet_benchmarks {
     }
 
     #[benchmark]
-    fn claim_root() {
+    fn claim_root(h: Linear<1, { crate::MAX_ROOT_CLAIM_WORK }>) {
         let coldkey: T::AccountId = whitelisted_caller();
-        let hotkey: T::AccountId = account("A", 0, 1);
-
+        let owner_coldkey: T::AccountId = account("claim_owner_cold", 0, 0);
+        let owner_hotkey: T::AccountId = account("claim_owner_hot", 0, 1);
         let netuid = Subtensor::<T>::get_next_netuid();
 
         let lock_cost = Subtensor::<T>::get_network_lock_cost();
-        add_balance_to_coldkey_account::<T>(&coldkey, lock_cost.into());
+        add_balance_to_coldkey_account::<T>(&owner_coldkey, lock_cost.into());
 
         assert_ok!(Subtensor::<T>::register_network(
-            RawOrigin::Signed(coldkey.clone()).into(),
-            hotkey.clone()
+            RawOrigin::Signed(owner_coldkey).into(),
+            owner_hotkey
         ));
-
         SubtokenEnabled::<T>::insert(netuid, true);
-
         Subtensor::<T>::set_network_registration_allowed(netuid, true);
-
         NetworkRegistrationAllowed::<T>::insert(netuid, true);
         FirstEmissionBlockNumber::<T>::insert(netuid, 0);
-
         SubnetMechanism::<T>::insert(netuid, 1);
         SubnetworkN::<T>::insert(netuid, 1);
         Subtensor::<T>::set_tao_weight(u64::MAX);
-
-        let root_stake = 100_000_000u64;
-        Subtensor::<T>::increase_stake_for_hotkey_and_coldkey_on_subnet(
-            &hotkey,
-            &coldkey,
-            NetUid::ROOT,
-            root_stake.into(),
-        );
-
-        let initial_total_hotkey_alpha = 100_000_000u64;
-        Subtensor::<T>::increase_stake_for_hotkey_and_coldkey_on_subnet(
-            &hotkey,
-            &coldkey,
+        set_reserves::<T>(
             netuid,
-            initial_total_hotkey_alpha.into(),
+            TaoBalance::from(100_000_000_000_000_u64),
+            AlphaBalance::from(100_000_000_000_000_u64),
         );
+        RootClaimableThreshold::<T>::insert(NetUid::ROOT, I96F32::from_num(0));
 
-        // Point the validator's basket weight vector at the subnet so the distributed root
-        // dividend is deposited into its fund (instead of being recycled for lack of weights).
-        if let Ok(root_uid) = Uids::<T>::try_get(NetUid::ROOT, &hotkey) {
-            Weights::<T>::insert(
-                NetUidStorageIndex::ROOT,
-                root_uid,
-                vec![(u16::from(netuid), 1u16)],
+        let escrow = Subtensor::<T>::get_beta_escrow_account_id();
+        let holding_alpha = AlphaBalance::from(1_000_000_u64);
+        for i in 0..h {
+            let hotkey: T::AccountId = account("claim_hot", i, 1);
+            Subtensor::<T>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                &coldkey,
+                NetUid::ROOT,
+                AlphaBalance::from(1_u64),
+            );
+            Subtensor::<T>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                &escrow,
+                netuid,
+                holding_alpha,
             );
+            BasketShares::<T>::insert(&hotkey, 1_u64);
+            BasketRate::<T>::insert(&hotkey, I96F32::from_num(1));
         }
 
-        let pending_root_alpha = 10_000_000u64;
-        Subtensor::<T>::distribute_emission(
-            netuid,
-            AlphaBalance::ZERO,
-            pending_root_alpha.into(),
-            pending_root_alpha.into(),
-            AlphaBalance::ZERO,
-        );
-
-        let initial_stake = Subtensor::<T>::get_stake_for_hotkey_and_coldkey_on_subnet(
-            &hotkey,
-            &coldkey,
-            NetUid::ROOT,
-        );
-
         #[extrinsic_call]
         _(RawOrigin::Signed(coldkey.clone()));
 
-        let new_stake = Subtensor::<T>::get_stake_for_hotkey_and_coldkey_on_subnet(
-            &hotkey,
-            &coldkey,
-            NetUid::ROOT,
-        );
-
-        // The claim must actually pay out (strict: a no-op claim is a broken benchmark).
-        assert!(new_stake > initial_stake);
+        // Every work unit must execute the active holding-redemption path.
+        let first_hotkey: T::AccountId = account("claim_hot", 0, 1);
+        let last_hotkey: T::AccountId = account("claim_hot", h.saturating_sub(1), 1);
+        assert_eq!(BasketShares::<T>::get(first_hotkey), 0);
+        assert_eq!(BasketShares::<T>::get(last_hotkey), 0);
     }
 
     #[benchmark]
     fn sudo_set_root_claim_threshold() {
-        let coldkey: T::AccountId = whitelisted_caller();
-        let hotkey: T::AccountId = account("A", 0, 1);
-
-        let netuid = Subtensor::<T>::get_next_netuid();
-
-        let lock_cost = Subtensor::<T>::get_network_lock_cost();
-        add_balance_to_coldkey_account::<T>(&coldkey, lock_cost.into());
-
-        assert_ok!(Subtensor::<T>::register_network(
-            RawOrigin::Signed(coldkey.clone()).into(),
-            hotkey.clone()
-        ));
-
         #[extrinsic_call]
-        _(RawOrigin::Root, netuid, 100);
+        _(RawOrigin::Root, NetUid::ROOT, 100);
     }
 
     #[benchmark]

pallets/subtensor/src/lib.rs

@@ -65,6 +65,10 @@ pub const ALPHA_MAP_BATCH_SIZE: usize = 30;
 
 pub const MAX_ROOT_CLAIM_THRESHOLD: u64 = 10_000_000;
 
+/// Maximum number of validator hotkeys or raw basket-position rows one `claim_root` call may
+/// inspect.
+pub const MAX_ROOT_CLAIM_WORK: u32 = 256;
+
 pub struct SubtensorDustRemoval<T>(PhantomData<T>);
 impl<T> frame_support::traits::OnUnbalanced<pallet_balances::CreditOf<T, ()>>
     for SubtensorDustRemoval<T>

pallets/subtensor/src/macros/dispatches.rs

@@ -13,6 +13,7 @@ mod dispatches {
 
     use crate::MAX_CRV3_COMMIT_SIZE_BYTES;
     use crate::MAX_ROOT_CLAIM_THRESHOLD;
+    use crate::MAX_ROOT_CLAIM_WORK;
     /// Dispatchable functions allow users to interact with the pallet and invoke state changes.
     /// These functions materialize as "extrinsics", which are often compared to transactions.
     /// Dispatchable functions must be annotated with a weight and must return a DispatchResult.
@@ -1907,16 +1908,20 @@ mod dispatches {
         /// * `RootClaimed`: On the successfully claiming the root emissions for a coldkey.
         ///
         /// # Errors
-        /// * `InvalidSubnetNumber`: The subnet set is empty or exceeds the maximum number of claims.
-        ///
+        /// * `TooManyRootClaimHotkeys`: The coldkey exceeds the benchmarked hotkey bound.
+        /// * `TooManyRootClaimHoldings`: The coldkey exceeds the benchmarked basket-position bound.
         #[pallet::call_index(121)]
-        #[pallet::weight(<T as crate::pallet::Config>::WeightInfo::claim_root())]
+        #[pallet::weight(
+            <T as crate::pallet::Config>::WeightInfo::claim_root(MAX_ROOT_CLAIM_WORK)
+        )]
         pub fn claim_root(origin: OriginFor<T>) -> DispatchResultWithPostInfo {
             let coldkey: T::AccountId = ensure_signed(origin)?;
 
+            let (hotkeys, work) = Self::bounded_root_claim_work(&coldkey)?;
+            Self::do_root_claim(coldkey.clone(), hotkeys)?;
             Self::maybe_add_coldkey_index(&coldkey);
 
-            let weight = Self::do_root_claim(coldkey)?;
+            let weight = <T as crate::pallet::Config>::WeightInfo::claim_root(work);
             Ok((Some(weight), Pays::Yes).into())
         }
 

pallets/subtensor/src/macros/errors.rs

@@ -350,5 +350,9 @@ mod errors {
         /// (measured from the last stake add/remove) and cannot be unstaked yet. Prevents
         /// epoch-boundary just-in-time dividend sniping.
         RootStakeLocked,
+        /// A root claim would inspect more validator hotkeys than its benchmarked bound.
+        TooManyRootClaimHotkeys,
+        /// A root claim would inspect more basket-position rows than its benchmarked bound.
+        TooManyRootClaimHoldings,
     }
 }

pallets/subtensor/src/staking/claim_root.rs

@@ -403,10 +403,7 @@ impl<T: Config> Pallet<T> {
             Self::hotkey_account_exists(&hotkey),
             Error::<T>::HotKeyAccountNotExists
         );
-        ensure!(
-            tao >= DefaultMinStake::<T>::get(),
-            Error::<T>::AmountTooLow
-        );
+        ensure!(tao >= DefaultMinStake::<T>::get(), Error::<T>::AmountTooLow);
         ensure!(
             Self::can_remove_balance_from_coldkey_account(&coldkey, tao.into()),
             Error::<T>::NotEnoughBalanceToStake
@@ -417,8 +414,8 @@ impl<T: Config> Pallet<T> {
 
         // Each weight slot can add at most one new holding, so pre-deploy holdings plus the
         // slot count bounds the holdings the two NAV valuations will sweep.
-        let num_holdings = (Self::get_basket_holdings(&hotkey).len() as u64)
-            .saturating_add(valid.len() as u64);
+        let num_holdings =
+            (Self::get_basket_holdings(&hotkey).len() as u64).saturating_add(valid.len() as u64);
 
         with_transaction(
             || match Self::try_stake_into_basket(&coldkey, &hotkey, tao, &valid) {
@@ -663,42 +660,67 @@ impl<T: Config> Pallet<T> {
         })
     }
 
-    fn root_claim_weight(num_holdings: u64) -> Weight {
-        // Per-holding: escrow stake read/write + swap + protocol-flow bookkeeping.
-        Weight::from_parts(20_000_000, 3000)
-            .saturating_add(T::DbWeight::get().reads(4_u64))
-            .saturating_add(T::DbWeight::get().writes(3_u64))
-            .saturating_mul(num_holdings.max(1))
-            .saturating_add(T::DbWeight::get().reads_writes(4_u64, 3_u64))
+    /// Returns the bounded hotkey set and benchmark work parameter for a root claim.
+    ///
+    /// The benchmark uses one hotkey with one basket position per work unit. Bounding both
+    /// independent dimensions by the same parameter therefore covers callers with many empty
+    /// hotkeys, callers with many holdings on one hotkey, and mixtures of the two.
+    pub(crate) fn bounded_root_claim_work(
+        coldkey: &T::AccountId,
+    ) -> Result<(Vec<T::AccountId>, u32), DispatchError> {
+        let hotkey_count = StakingHotkeys::<T>::decode_len(coldkey).unwrap_or_default();
+        ensure!(
+            hotkey_count <= crate::MAX_ROOT_CLAIM_WORK as usize,
+            Error::<T>::TooManyRootClaimHotkeys
+        );
+        let hotkey_count = hotkey_count as u32;
+
+        let hotkeys = StakingHotkeys::<T>::get(coldkey);
+        let escrow = Self::get_beta_escrow_account_id();
+        let mut holding_rows = 0_u32;
+
+        for hotkey in &hotkeys {
+            for _ in Alpha::<T>::iter_prefix((hotkey, &escrow))
+                .map(|_| ())
+                .chain(AlphaV2::<T>::iter_prefix((hotkey, &escrow)).map(|_| ()))
+            {
+                holding_rows = holding_rows.saturating_add(1);
+                ensure!(
+                    holding_rows <= crate::MAX_ROOT_CLAIM_WORK,
+                    Error::<T>::TooManyRootClaimHoldings
+                );
+            }
+        }
+
+        Ok((hotkeys, hotkey_count.max(holding_rows).max(1)))
     }
 
-    pub fn do_root_claim(coldkey: T::AccountId) -> Result<Weight, DispatchError> {
-        with_transaction(|| match Self::try_do_root_claim(coldkey) {
-            Ok(weight) => TransactionOutcome::Commit(Ok(weight)),
+    pub fn do_root_claim(
+        coldkey: T::AccountId,
+        hotkeys: Vec<T::AccountId>,
+    ) -> Result<(), DispatchError> {
+        with_transaction(|| match Self::try_do_root_claim(coldkey, &hotkeys) {
+            Ok(()) => TransactionOutcome::Commit(Ok(())),
             Err(err) => TransactionOutcome::Rollback(Err(err)),
         })
     }
 
-    fn try_do_root_claim(coldkey: T::AccountId) -> Result<Weight, DispatchError> {
-        let mut weight = Weight::default();
-
-        let hotkeys = StakingHotkeys::<T>::get(&coldkey);
-        weight.saturating_accrue(T::DbWeight::get().reads(1));
-
+    fn try_do_root_claim(
+        coldkey: T::AccountId,
+        hotkeys: &[T::AccountId],
+    ) -> Result<(), DispatchError> {
         let mut total_tao: u64 = 0;
-        for hotkey in hotkeys.iter() {
-            let num_holdings = Self::get_basket_holdings(hotkey).len() as u64;
+        for hotkey in hotkeys {
             let realized = Self::root_claim_for_hotkey(hotkey, &coldkey, false)?;
             total_tao = total_tao.saturating_add(realized);
-            weight.saturating_accrue(Self::root_claim_weight(num_holdings));
         }
 
         Self::deposit_event(Event::RootClaimed {
             coldkey,
             tao: total_tao.into(),
         });
 
-        Ok(weight)
+        Ok(())
     }
 
     pub fn maybe_add_coldkey_index(coldkey: &T::AccountId) {