fix benchmark remove stake

pallets/subtensor/src/benchmarks.rs

@@ -913,6 +913,61 @@ mod pallet_benchmarks {
         );
     }
 
+    #[benchmark]
+    fn remove_stake() {
+        let netuid = NetUid::from(1);
+        let tempo: u16 = 1;
+        let seed: u32 = 1;
+
+        Subtensor::<T>::increase_total_stake(1_000_000_000_000_u64.into());
+
+        Subtensor::<T>::init_new_network(netuid, tempo);
+        Subtensor::<T>::set_network_registration_allowed(netuid, true);
+        SubtokenEnabled::<T>::insert(netuid, true);
+
+        Subtensor::<T>::set_max_allowed_uids(netuid, 4096);
+        assert_eq!(Subtensor::<T>::get_max_allowed_uids(netuid), 4096);
+
+        let coldkey: T::AccountId = account("Test", 0, seed);
+        let hotkey: T::AccountId = account("Alice", 0, seed);
+        Subtensor::<T>::set_burn(netuid, benchmark_registration_burn());
+
+        let tao_reserve = TaoBalance::from(1_000_000_000_000_u64);
+        let alpha_in = AlphaBalance::from(100_000_000_000_000_u64);
+        set_reserves::<T>(netuid, tao_reserve, alpha_in);
+
+        let wallet_bal = 1000000u32.into();
+        Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), wallet_bal);
+
+        assert_ok!(Subtensor::<T>::burned_register(
+            RawOrigin::Signed(coldkey.clone()).into(),
+            netuid,
+            hotkey.clone()
+        ));
+
+        let staked_amt = TaoBalance::from(100_000_000_000_u64);
+        Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), staked_amt);
+
+        assert_ok!(Subtensor::<T>::add_stake(
+            RawOrigin::Signed(coldkey.clone()).into(),
+            hotkey.clone(),
+            netuid,
+            staked_amt
+        ));
+
+        let amount_unstaked = AlphaBalance::from(30_000_000_000_u64);
+
+        StakingOperationRateLimiter::<T>::remove((hotkey.clone(), coldkey.clone(), netuid));
+
+        #[extrinsic_call]
+        _(
+            RawOrigin::Signed(coldkey.clone()),
+            hotkey.clone(),
+            netuid,
+            amount_unstaked,
+        );
+    }
+
     #[benchmark]
     fn remove_stake_limit() {
         let netuid = NetUid::from(1);

pallets/subtensor/src/macros/dispatches.rs

@@ -736,9 +736,7 @@ mod dispatches {
         /// 	- Thrown if there is not enough stake on the hotkey to withdwraw this amount.
         ///
         #[pallet::call_index(3)]
-        #[pallet::weight((Weight::from_parts(196_800_000, 0)
-		.saturating_add(T::DbWeight::get().reads(19))
-		.saturating_add(T::DbWeight::get().writes(10)), DispatchClass::Normal, Pays::Yes))]
+        #[pallet::weight(<T as crate::pallet::Config>::WeightInfo::remove_stake())]
         pub fn remove_stake(
             origin: OriginFor<T>,
             hotkey: T::AccountId,

pallets/subtensor/src/weights.rs

@@ -89,6 +89,7 @@ pub trait WeightInfo {
 	fn add_stake_burn() -> Weight;
 	fn set_pending_childkey_cooldown() -> Weight;
     fn set_auto_parent_delegation_enabled() -> Weight;
+	fn remove_stake() -> Weight;
 }
 
 /// Weights for `pallet_subtensor` using the Substrate node and recommended hardware.
@@ -1133,6 +1134,17 @@ impl<T: frame_system::Config> WeightInfo for SubstrateWeight<T> {
 			.saturating_add(T::DbWeight::get().reads(30_u64))
 			.saturating_add(T::DbWeight::get().writes(13_u64))
 	}
+
+	fn remove_stake() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `2211`
+		//  Estimated: `10626`
+		// Minimum execution time: 387_646_000 picoseconds.
+		Weight::from_parts(403_169_000, 10626)
+			.saturating_add(T::DbWeight::get().reads(30_u64))
+			.saturating_add(T::DbWeight::get().writes(13_u64))
+	}
+
 	/// Storage: `SubtensorModule::Alpha` (r:2 w:0)
 	/// Proof: `SubtensorModule::Alpha` (`max_values`: None, `max_size`: None, mode: `Measured`)
 	/// Storage: `SubtensorModule::AlphaV2` (r:2 w:2)
@@ -3251,6 +3263,17 @@ impl WeightInfo for () {
 			.saturating_add(RocksDbWeight::get().reads(30_u64))
 			.saturating_add(RocksDbWeight::get().writes(13_u64))
 	}
+
+	fn remove_stake() -> Weight {
+		// Proof Size summary in bytes:
+		//  Measured:  `2211`
+		//  Estimated: `10626`
+		// Minimum execution time: 387_646_000 picoseconds.
+		Weight::from_parts(403_169_000, 10626)
+			.saturating_add(RocksDbWeight::get().reads(30_u64))
+			.saturating_add(RocksDbWeight::get().writes(13_u64))
+	}
+
 	/// Storage: `SubtensorModule::Alpha` (r:2 w:0)
 	/// Proof: `SubtensorModule::Alpha` (`max_values`: None, `max_size`: None, mode: `Measured`)
 	/// Storage: `SubtensorModule::AlphaV2` (r:2 w:2)

pallets/transaction-fee/src/tests/mod.rs

@@ -602,7 +602,7 @@ fn test_remove_stake_edge_alpha() {
         );
 
         // For-set Alpha balance to low, but enough to pay tx fees at the current Alpha price
-        let new_current_stake = AlphaBalance::from(1_000_000);
+        let new_current_stake = AlphaBalance::from(2_000_000);
         SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &sn.hotkeys[0],
             &sn.coldkey,