fix: must_use, try_decrease, alpha-fees

pallets/subtensor/src/staking/helpers.rs

@@ -207,12 +207,17 @@ impl<T: Config> Pallet<T> {
                     // Add the stake to the coldkey account.
                     Self::add_balance_to_coldkey_account(coldkey, cleared_stake.into());
                 } else {
-                    // Just clear small alpha
+                    // We should only get here when
+                    // - alpha < min_alpha_stake
+                    // - the unstake fails (for whatever reason)
+
+                    // Unstake and clear the amount, burning the alpha.
                     let alpha =
                         Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, netuid);
-                    Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+                    let alpha_unstaked = Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
                         hotkey, coldkey, netuid, alpha,
                     );
+                    Self::burn_subnet_alpha(netuid, alpha_unstaked);
                 }
             }
         }

pallets/subtensor/src/staking/stake_utils.rs

@@ -6,6 +6,8 @@ use substrate_fixed::types::{I64F64, I96F32, U64F64, U96F32};
 use subtensor_runtime_common::{AlphaCurrency, Currency, NetUid, TaoCurrency};
 use subtensor_swap_interface::{Order, SwapHandler, SwapResult};
 
+use frame_support::storage::{TransactionOutcome, transactional};
+
 impl<T: Config> Pallet<T> {
     /// Retrieves the total alpha issuance for a given subnet.
     ///
@@ -553,6 +555,7 @@ impl<T: Config> Pallet<T> {
     /// * `netuid` - The unique identifier of the subnet.
     /// * `amount` - The amount of alpha to be added.
     ///
+    #[must_use]
     pub fn decrease_stake_for_hotkey_and_coldkey_on_subnet(
         hotkey: &T::AccountId,
         coldkey: &T::AccountId,
@@ -576,6 +579,43 @@ impl<T: Config> Pallet<T> {
         actual_alpha.neg().max(0).unsigned_abs().into()
     }
 
+    pub fn try_decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        hotkey: &T::AccountId,
+        coldkey: &T::AccountId,
+        netuid: NetUid,
+        amount: AlphaCurrency,
+    ) -> bool {
+        let alpha_share_pool = Self::get_alpha_share_pool(hotkey.clone(), netuid);
+
+        if let Ok(value) = alpha_share_pool.try_get_value(coldkey) {
+            if value < amount.to_u64() {
+                return false;
+            }
+
+            // Try to decrease the stake, if we unstake too much, fail
+            // Rollback both operations
+            let result = transactional::with_transaction(
+                || -> TransactionOutcome<Result<(), sp_runtime::DispatchError>> {
+                    let result = Self::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+                        hotkey, coldkey, netuid, amount,
+                    );
+                    if result > amount {
+                        return TransactionOutcome::Rollback(Err(
+                            Error::<T>::NotEnoughStakeToWithdraw.into(),
+                        ));
+                    }
+
+                    TransactionOutcome::Rollback(Ok(()))
+                },
+            )
+            .is_ok();
+
+            return result;
+        }
+
+        false
+    }
+
     /// Swaps TAO for the alpha token on the subnet.
     ///
     /// Updates TaoIn, AlphaIn, and AlphaOut

pallets/subtensor/src/tests/children.rs

@@ -2334,7 +2334,7 @@ fn test_do_set_child_cooldown_period() {
         );
 
         wait_and_set_pending_children(netuid);
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &parent,
             &coldkey,
             netuid,
@@ -2475,7 +2475,7 @@ fn test_revoke_child_no_min_stake_check() {
         assert_eq!(children_before, vec![]);
 
         wait_and_set_pending_children(netuid);
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &parent,
             &coldkey,
             NetUid::ROOT,
@@ -2545,7 +2545,7 @@ fn test_do_set_child_registration_disabled() {
         ));
 
         wait_and_set_pending_children(netuid);
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &parent,
             &coldkey,
             netuid,

pallets/subtensor/src/tests/staking.rs

@@ -1240,7 +1240,7 @@ fn test_remove_stake_from_hotkey_account() {
         );
 
         // Remove stake
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey_id,
             &coldkey_id,
             netuid,
@@ -1293,7 +1293,7 @@ fn test_remove_stake_from_hotkey_account_registered_in_various_networks() {
         );
 
         // Remove all stake
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &hotkey_id,
             &coldkey_id,
             netuid,

pallets/subtensor/src/tests/staking2.rs

@@ -247,7 +247,7 @@ fn test_share_based_staking() {
 
         // Test Case 7: Stake Removal
         // Verify correct stake removal from both accounts
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &primary_hotkey,
             &primary_coldkey,
             netuid,
@@ -273,7 +273,7 @@ fn test_share_based_staking() {
             "Stake removal should decrease balance by exact amount"
         );
 
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &primary_hotkey,
             &secondary_coldkey,
             netuid,
@@ -346,7 +346,7 @@ fn test_share_based_staking() {
         log::info!(
             "Attempting to remove excessive stake: {available_stake} + 1000 = {excessive_amount}"
         );
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &primary_hotkey,
             &primary_coldkey,
             netuid,
@@ -383,7 +383,7 @@ fn test_share_based_staking() {
 
         // Test removing stake from non-existent coldkey
         let non_existent_coldkey = U256::from(5);
-        SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+        let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
             &primary_hotkey,
             &non_existent_coldkey,
             netuid,
@@ -435,7 +435,7 @@ fn test_share_based_staking_denominator_precision() {
                     .to_num::<u64>()
                     .into(),
             );
-            SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+            let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
                 &hotkey1,
                 &coldkey1,
                 netuid,
@@ -490,7 +490,7 @@ fn test_share_based_staking_stake_unstake_inject() {
                 netuid,
                 stake_amount,
             );
-            SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+            let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
                 &hotkey1,
                 &coldkey1,
                 netuid,
@@ -502,7 +502,7 @@ fn test_share_based_staking_stake_unstake_inject() {
                 netuid,
                 stake_amount,
             );
-            SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+            let _ = SubtensorModule::decrease_stake_for_hotkey_and_coldkey_on_subnet(
                 &hotkey1,
                 &coldkey2,
                 netuid,

pallets/transaction-fee/src/lib.rs

@@ -67,7 +67,7 @@ pub trait AlphaFeeHandler<T: frame_system::Config> {
         coldkey: &AccountIdOf<T>,
         alpha_vec: &[(AccountIdOf<T>, NetUid)],
         tao_amount: u64,
-    );
+    ) -> Result<(), DispatchError>;
     fn get_all_netuids_for_coldkey_and_hotkey(
         coldkey: &AccountIdOf<T>,
         hotkey: &AccountIdOf<T>,
@@ -151,22 +151,37 @@ where
                 ),
             );
             let alpha_price = pallet_subtensor_swap::Pallet::<T>::current_alpha_price(*netuid);
-            alpha_price.saturating_mul(alpha_balance) >= tao_per_entry
+            if alpha_price.saturating_mul(alpha_balance) < tao_per_entry {
+                return false;
+            }
+
+            let alpha_fee = U96F32::saturating_from_num(tao_per_entry)
+                .checked_div(alpha_price)
+                .unwrap_or(alpha_balance)
+                .min(alpha_balance)
+                .saturating_to_num::<u64>();
+
+            pallet_subtensor::Pallet::<T>::try_decrease_stake_for_hotkey_and_coldkey_on_subnet(
+                hotkey,
+                coldkey,
+                *netuid,
+                alpha_fee.into(),
+            )
         })
     }
 
     fn withdraw_in_alpha(
         coldkey: &AccountIdOf<T>,
         alpha_vec: &[(AccountIdOf<T>, NetUid)],
         tao_amount: u64,
-    ) {
+    ) -> Result<(), DispatchError> {
         if alpha_vec.is_empty() {
-            return;
+            return Ok(());
         }
 
         let tao_per_entry = tao_amount.checked_div(alpha_vec.len() as u64).unwrap_or(0);
 
-        alpha_vec.iter().for_each(|(hotkey, netuid)| {
+        for (hotkey, netuid) in alpha_vec.iter() {
             // Divide tao_amount evenly among all alpha entries
             let alpha_balance = U96F32::saturating_from_num(
                 pallet_subtensor::Pallet::<T>::get_stake_for_hotkey_and_coldkey_on_subnet(
@@ -180,13 +195,20 @@ where
                 .min(alpha_balance)
                 .saturating_to_num::<u64>();
 
-            pallet_subtensor::Pallet::<T>::decrease_stake_for_hotkey_and_coldkey_on_subnet(
-                hotkey,
-                coldkey,
-                *netuid,
-                alpha_fee.into(),
+            let alpha_removed =
+                pallet_subtensor::Pallet::<T>::decrease_stake_for_hotkey_and_coldkey_on_subnet(
+                    hotkey,
+                    coldkey,
+                    *netuid,
+                    alpha_fee.into(),
+                );
+            ensure!(
+                alpha_removed == alpha_fee.into(),
+                "Alpha is not greater than alpha fee"
             );
-        });
+        }
+
+        Ok(())
     }
 
     fn get_all_netuids_for_coldkey_and_hotkey(
@@ -308,7 +330,7 @@ where
 
     fn withdraw_fee(
         who: &AccountIdOf<T>,
-        _call: &CallOf<T>,
+        call: &CallOf<T>,
         _dispatch_info: &DispatchInfoOf<CallOf<T>>,
         fee: Self::Balance,
         _tip: Self::Balance,
@@ -327,20 +349,22 @@ where
         ) {
             Ok(imbalance) => Ok(Some(WithdrawnFee::Tao(imbalance))),
             Err(_) => {
-                // let alpha_vec = Self::fees_in_alpha::<T>(who, call);
-                // if !alpha_vec.is_empty() {
-                //     let fee_u64: u64 = fee.into();
-                //     OU::withdraw_in_alpha(who, &alpha_vec, fee_u64);
-                //     return Ok(Some(WithdrawnFee::Alpha));
-                // }
+                let alpha_vec = Self::fees_in_alpha::<T>(who, call);
+                if !alpha_vec.is_empty() {
+                    let fee_u64: u64 = fee.into();
+                    OU::withdraw_in_alpha(who, &alpha_vec, fee_u64).map_err(|_| {
+                        TransactionValidityError::Invalid(InvalidTransaction::Payment)
+                    })?;
+                    return Ok(Some(WithdrawnFee::Alpha));
+                }
                 Err(InvalidTransaction::Payment.into())
             }
         }
     }
 
     fn can_withdraw_fee(
         who: &AccountIdOf<T>,
-        _call: &CallOf<T>,
+        call: &CallOf<T>,
         _dispatch_info: &DispatchInfoOf<CallOf<T>>,
         fee: Self::Balance,
         _tip: Self::Balance,
@@ -353,14 +377,14 @@ where
         match F::can_withdraw(who, fee) {
             WithdrawConsequence::Success => Ok(()),
             _ => {
-                // // Fallback to fees in Alpha if possible
-                // let alpha_vec = Self::fees_in_alpha::<T>(who, call);
-                // if !alpha_vec.is_empty() {
-                //     let fee_u64: u64 = fee.into();
-                //     if OU::can_withdraw_in_alpha(who, &alpha_vec, fee_u64) {
-                //         return Ok(());
-                //     }
-                // }
+                // Fallback to fees in Alpha if possible
+                let alpha_vec = Self::fees_in_alpha::<T>(who, call);
+                if !alpha_vec.is_empty() {
+                    let fee_u64: u64 = fee.into();
+                    if OU::can_withdraw_in_alpha(who, &alpha_vec, fee_u64) {
+                        return Ok(());
+                    }
+                }
                 Err(InvalidTransaction::Payment.into())
             }
         }