Fix staking counter

pallets/subtensor/src/macros/errors.rs

@@ -91,9 +91,7 @@ mod errors {
         /// A transactor exceeded the rate limit for setting or swapping hotkey.
         HotKeySetTxRateLimitExceeded,
         /// A transactor exceeded the rate limit for staking.
-        StakeRateLimitExceeded,
-        /// A transactor exceeded the rate limit for unstaking.
-        UnstakeRateLimitExceeded,
+        StakingRateLimitExceeded,
         /// Registration is disabled.
         SubNetRegistrationDisabled,
         /// The number of registration attempts exceeded the allowed number in the interval.

pallets/subtensor/src/staking/add_stake.rs

@@ -62,19 +62,6 @@ impl<T: Config> Pallet<T> {
             Error::<T>::HotKeyNotDelegateAndSignerNotOwnHotKey
         );
 
-        // Ensure we don't exceed stake rate limit
-        let stakes_this_interval =
-            Self::get_stakes_this_interval_for_coldkey_hotkey(&coldkey, &hotkey);
-        ensure!(
-            stakes_this_interval < Self::get_target_stakes_per_interval(),
-            Error::<T>::StakeRateLimitExceeded
-        );
-
-        // Track this addition in the stake delta.
-        StakeDeltaSinceLastEmissionDrain::<T>::mutate(&hotkey, &coldkey, |stake_delta| {
-            *stake_delta = stake_delta.saturating_add_unsigned(stake_to_be_added as u128);
-        });
-
         // If coldkey is not owner of the hotkey, it's a nomination stake.
         if !Self::coldkey_owns_hotkey(&coldkey, &hotkey) {
             let total_stake_after_add =
@@ -86,24 +73,24 @@ impl<T: Config> Pallet<T> {
             );
         }
 
+        Self::try_increase_staking_counter(&coldkey, &hotkey)?;
+
         // Ensure the remove operation from the coldkey is a success.
         let actual_amount_to_stake =
             Self::remove_balance_from_coldkey_account(&coldkey, stake_to_be_added)?;
 
         // If we reach here, add the balance to the hotkey.
         Self::increase_stake_on_coldkey_hotkey_account(&coldkey, &hotkey, actual_amount_to_stake);
 
+        // Track this addition in the stake delta.
+        StakeDeltaSinceLastEmissionDrain::<T>::mutate(&hotkey, &coldkey, |stake_delta| {
+            *stake_delta = stake_delta.saturating_add_unsigned(stake_to_be_added as u128);
+        });
+
         // Set last block for rate limiting
-        let block: u64 = Self::get_current_block_as_u64();
+        let block = Self::get_current_block_as_u64();
         Self::set_last_tx_block(&coldkey, block);
 
-        // Emit the staking event.
-        Self::set_stakes_this_interval_for_coldkey_hotkey(
-            &coldkey,
-            &hotkey,
-            stakes_this_interval.saturating_add(1),
-            block,
-        );
         log::debug!(
             "StakeAdded( hotkey:{:?}, stake_to_be_added:{:?} )",
             hotkey,

pallets/subtensor/src/staking/helpers.rs

@@ -59,42 +59,6 @@ impl<T: Config> Pallet<T> {
         Stake::<T>::get(hotkey, coldkey)
     }
 
-    // Retrieves the total stakes for a given hotkey (account ID) for the current staking interval.
-    pub fn get_stakes_this_interval_for_coldkey_hotkey(
-        coldkey: &T::AccountId,
-        hotkey: &T::AccountId,
-    ) -> u64 {
-        // Retrieve the configured stake interval duration from storage.
-        let stake_interval = StakeInterval::<T>::get();
-
-        // Obtain the current block number as an unsigned 64-bit integer.
-        let current_block = Self::get_current_block_as_u64();
-
-        // Fetch the total stakes and the last block number when stakes were made for the hotkey.
-        let (stakes, block_last_staked_at) =
-            TotalHotkeyColdkeyStakesThisInterval::<T>::get(coldkey, hotkey);
-
-        // Calculate the block number after which the stakes for the hotkey should be reset.
-        let block_to_reset_after = block_last_staked_at.saturating_add(stake_interval);
-
-        // If the current block number is beyond the reset point,
-        // it indicates the end of the staking interval for the hotkey.
-        if block_to_reset_after <= current_block {
-            // Reset the stakes for this hotkey for the current interval.
-            Self::set_stakes_this_interval_for_coldkey_hotkey(
-                coldkey,
-                hotkey,
-                0,
-                block_last_staked_at,
-            );
-            // Return 0 as the stake amount since we've just reset the stakes.
-            return 0;
-        }
-
-        // If the staking interval has not yet ended, return the current stake amount.
-        stakes
-    }
-
     pub fn get_target_stakes_per_interval() -> u64 {
         TargetStakesPerInterval::<T>::get()
     }

pallets/subtensor/src/staking/remove_stake.rs

@@ -65,13 +65,7 @@ impl<T: Config> Pallet<T> {
             Error::<T>::NotEnoughStakeToWithdraw
         );
 
-        // Ensure we don't exceed stake rate limit
-        let unstakes_this_interval =
-            Self::get_stakes_this_interval_for_coldkey_hotkey(&coldkey, &hotkey);
-        ensure!(
-            unstakes_this_interval < Self::get_target_stakes_per_interval(),
-            Error::<T>::UnstakeRateLimitExceeded
-        );
+        Self::try_increase_staking_counter(&coldkey, &hotkey)?;
 
         // We remove the balance from the hotkey.
         Self::decrease_stake_on_coldkey_hotkey_account(&coldkey, &hotkey, stake_to_be_removed);
@@ -98,16 +92,10 @@ impl<T: Config> Pallet<T> {
         }
 
         // Set last block for rate limiting
-        let block: u64 = Self::get_current_block_as_u64();
+        let block = Self::get_current_block_as_u64();
         Self::set_last_tx_block(&coldkey, block);
 
         // Emit the unstaking event.
-        Self::set_stakes_this_interval_for_coldkey_hotkey(
-            &coldkey,
-            &hotkey,
-            unstakes_this_interval.saturating_add(1),
-            block,
-        );
         log::debug!(
             "StakeRemoved( hotkey:{:?}, stake_to_be_removed:{:?} )",
             hotkey,