Fix balancer performance

chain-extensions/src/tests.rs

@@ -1297,15 +1297,26 @@ fn add_stake_recycle_rollback_on_recycle_failure() {
 
         let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
-        // Set up very low reserves so recycle will fail with InsufficientLiquidity
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::insert_lock_state(
+            &coldkey,
+            netuid,
+            &hotkey,
+            pallet_subtensor::staking::lock::LockState {
+                locked_mass: AlphaBalance::from(u64::MAX / 4),
+                conviction: U64F64::saturating_from_num(0),
+                last_update: pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64(),
+            },
+        );
+
+        // Leave enough input-side liquidity for add_stake to pass the 1000x swap input cap.
+        // The lock above makes the recycle leg fail, exercising atomic rollback.
         mock::setup_reserves(
             netuid,
-            TaoBalance::from(1_000_u64),
+            TaoBalance::from(tao_amount_raw / 1000 + 1),
             AlphaBalance::from(1_000_u64),
         );
 
-        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
-
         add_balance_to_coldkey_account(
             &coldkey,
             TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),
@@ -1368,15 +1379,26 @@ fn add_stake_burn_rollback_on_burn_failure() {
 
         let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
-        // Set up very low reserves so burn will fail with InsufficientLiquidity
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::insert_lock_state(
+            &coldkey,
+            netuid,
+            &hotkey,
+            pallet_subtensor::staking::lock::LockState {
+                locked_mass: AlphaBalance::from(u64::MAX / 4),
+                conviction: U64F64::saturating_from_num(0),
+                last_update: pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64(),
+            },
+        );
+
+        // Leave enough input-side liquidity for add_stake to pass the 1000x swap input cap.
+        // The lock above makes the burn leg fail, exercising atomic rollback.
         mock::setup_reserves(
             netuid,
-            TaoBalance::from(1_000_u64),
+            TaoBalance::from(tao_amount_raw / 1000 + 1),
             AlphaBalance::from(1_000_u64),
         );
 
-        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
-
         add_balance_to_coldkey_account(
             &coldkey,
             TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),

pallets/subtensor/src/staking/add_stake.rs

@@ -48,6 +48,8 @@ impl<T: Config> Pallet<T> {
             "do_add_stake( origin:{coldkey:?} hotkey:{hotkey:?}, netuid:{netuid:?}, stake_to_be_added:{stake_to_be_added:?} )"
         );
 
+        Self::ensure_add_stake_input_within_swap_limit(netuid, stake_to_be_added)?;
+
         // 2. Validate user input
         Self::validate_add_stake(
             &coldkey,
@@ -124,6 +126,8 @@ impl<T: Config> Pallet<T> {
             "do_add_stake( origin:{coldkey:?} hotkey:{hotkey:?}, netuid:{netuid:?}, stake_to_be_added:{stake_to_be_added:?} )"
         );
 
+        Self::ensure_add_stake_input_within_swap_limit(netuid, stake_to_be_added)?;
+
         // 2. Calculate the maximum amount that can be executed with price limit
         let max_amount: TaoBalance = Self::get_max_amount_add(netuid, limit_price)?.into();
         let mut possible_stake = stake_to_be_added;
@@ -175,11 +179,27 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        // Use reverting swap to estimate max limit amount
-        let order = GetAlphaForTao::<T>::with_amount(u64::MAX);
+        // Use the largest supported input instead of probing the swap path with u64::MAX.
+        let max_supported_input = SubnetTAO::<T>::get(netuid).saturating_mul(1_000.into());
+        let order = GetAlphaForTao::<T>::with_amount(max_supported_input);
         let result = T::SwapInterface::swap(netuid.into(), order, limit_price, false, true)
             .map(|r| r.amount_paid_in.saturating_add(r.fee_paid))?;
 
         Ok(result.into())
     }
+
+    fn ensure_add_stake_input_within_swap_limit(
+        netuid: NetUid,
+        amount: TaoBalance,
+    ) -> Result<(), Error<T>> {
+        if !netuid.is_root() && SubnetMechanism::<T>::get(netuid) == 1 {
+            let max_supported_input = SubnetTAO::<T>::get(netuid).saturating_mul(1_000.into());
+            ensure!(
+                amount <= max_supported_input,
+                Error::<T>::InsufficientLiquidity
+            );
+        }
+
+        Ok(())
+    }
 }

pallets/subtensor/src/staking/recycle_alpha.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::{Error, system::ensure_signed};
+use frame_support::storage::{TransactionOutcome, with_transaction};
 use subtensor_runtime_common::{AlphaBalance, NetUid};
 
 impl<T: Config> Pallet<T> {
@@ -132,22 +133,38 @@ impl<T: Config> Pallet<T> {
         amount: TaoBalance,
         limit: Option<TaoBalance>,
     ) -> DispatchResult {
-        let alpha = if let Some(limit) = limit {
-            Self::do_add_stake_limit(origin.clone(), hotkey.clone(), netuid, amount, limit, false)?
-        } else {
-            Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?
-        };
-
-        Self::do_burn_alpha(origin, hotkey.clone(), alpha, netuid)?;
-
-        Self::deposit_event(Event::AddStakeBurn {
-            netuid,
-            hotkey,
-            amount,
-            alpha,
-        });
-
-        Ok(())
+        with_transaction(|| {
+            let result = (|| {
+                let alpha = if let Some(limit) = limit {
+                    Self::do_add_stake_limit(
+                        origin.clone(),
+                        hotkey.clone(),
+                        netuid,
+                        amount,
+                        limit,
+                        false,
+                    )?
+                } else {
+                    Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?
+                };
+
+                Self::do_burn_alpha(origin, hotkey.clone(), alpha, netuid)?;
+
+                Self::deposit_event(Event::AddStakeBurn {
+                    netuid,
+                    hotkey,
+                    amount,
+                    alpha,
+                });
+
+                Ok(())
+            })();
+
+            match result {
+                Ok(()) => TransactionOutcome::Commit(Ok(())),
+                Err(err) => TransactionOutcome::Rollback(Err(err)),
+            }
+        })
     }
 
     /// Atomically stakes TAO and recycles the resulting alpha.
@@ -160,8 +177,17 @@ impl<T: Config> Pallet<T> {
         netuid: NetUid,
         amount: TaoBalance,
     ) -> Result<AlphaBalance, DispatchError> {
-        let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
-        Self::do_recycle_alpha(origin, hotkey, alpha, netuid)
+        with_transaction(|| {
+            let result = (|| {
+                let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
+                Self::do_recycle_alpha(origin, hotkey, alpha, netuid)
+            })();
+
+            match result {
+                Ok(alpha) => TransactionOutcome::Commit(Ok(alpha)),
+                Err(err) => TransactionOutcome::Rollback(Err(err)),
+            }
+        })
     }
 
     /// Atomically stakes TAO and burns the resulting alpha. Permissionless
@@ -173,7 +199,16 @@ impl<T: Config> Pallet<T> {
         netuid: NetUid,
         amount: TaoBalance,
     ) -> Result<AlphaBalance, DispatchError> {
-        let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
-        Self::do_burn_alpha(origin, hotkey, alpha, netuid)
+        with_transaction(|| {
+            let result = (|| {
+                let alpha = Self::do_add_stake(origin.clone(), hotkey.clone(), netuid, amount)?;
+                Self::do_burn_alpha(origin, hotkey, alpha, netuid)
+            })();
+
+            match result {
+                Ok(alpha) => TransactionOutcome::Commit(Ok(alpha)),
+                Err(err) => TransactionOutcome::Rollback(Err(err)),
+            }
+        })
     }
 }

pallets/subtensor/src/staking/remove_stake.rs

@@ -55,6 +55,7 @@ impl<T: Config> Pallet<T> {
         let alpha_available =
             Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid);
         let alpha_unstaked = alpha_unstaked.min(alpha_available);
+        Self::ensure_remove_stake_input_within_swap_limit(netuid, alpha_unstaked)?;
 
         // 2. Validate the user input
         Self::validate_remove_stake(
@@ -336,6 +337,8 @@ impl<T: Config> Pallet<T> {
             "do_remove_stake( origin:{coldkey:?} hotkey:{hotkey:?}, netuid: {netuid:?}, alpha_unstaked:{alpha_unstaked:?} )"
         );
 
+        Self::ensure_remove_stake_input_within_swap_limit(netuid, alpha_unstaked)?;
+
         // 2. Calculate the maximum amount that can be executed with price limit
         let max_amount = Self::get_max_amount_remove(netuid, limit_price)?;
         let mut possible_alpha = alpha_unstaked;
@@ -394,14 +397,30 @@ impl<T: Config> Pallet<T> {
             }
         }
 
-        // Use reverting swap to estimate max limit amount
-        let order = GetTaoForAlpha::<T>::with_amount(u64::MAX);
+        // Use the largest supported input instead of probing the swap path with u64::MAX.
+        let max_supported_input = SubnetAlphaIn::<T>::get(netuid).saturating_mul(1_000.into());
+        let order = GetTaoForAlpha::<T>::with_amount(max_supported_input);
         let result = T::SwapInterface::swap(netuid.into(), order, limit_price.into(), false, true)
             .map(|r| r.amount_paid_in.saturating_add(r.fee_paid))?;
 
         Ok(result)
     }
 
+    fn ensure_remove_stake_input_within_swap_limit(
+        netuid: NetUid,
+        amount: AlphaBalance,
+    ) -> Result<(), Error<T>> {
+        if !netuid.is_root() && SubnetMechanism::<T>::get(netuid) == 1 {
+            let max_supported_input = SubnetAlphaIn::<T>::get(netuid).saturating_mul(1_000.into());
+            ensure!(
+                amount <= max_supported_input,
+                Error::<T>::InsufficientLiquidity
+            );
+        }
+
+        Ok(())
+    }
+
     pub fn do_remove_stake_full_limit(
         origin: OriginFor<T>,
         hotkey: T::AccountId,