Handle swap errors when taking fees in alpha

pallets/transaction-fee/src/lib.rs

@@ -3,6 +3,7 @@
 // FRAME
 use frame_support::{
     pallet_prelude::*,
+    storage::{TransactionOutcome, with_transaction},
     traits::{
         Imbalance, IsSubType, OnUnbalanced,
         fungible::{
@@ -16,8 +17,9 @@ use frame_support::{
 
 // Runtime
 use sp_runtime::{
-    Perbill, Saturating,
+    DispatchError, Perbill, Saturating,
     traits::{DispatchInfoOf, PostDispatchInfoOf},
+    transaction_validity::{InvalidTransaction, TransactionValidityError},
 };
 
 // Pallets
@@ -67,7 +69,7 @@ pub trait AlphaFeeHandler<T: frame_system::Config> {
         coldkey: &AccountIdOf<T>,
         alpha_vec: &[(AccountIdOf<T>, NetUid)],
         tao_amount: TaoBalance,
-    ) -> (AlphaBalance, TaoBalance, NetUid);
+    ) -> Result<(AlphaBalance, TaoBalance, NetUid), TransactionValidityError>;
     fn get_all_netuids_for_coldkey_and_hotkey(
         coldkey: &AccountIdOf<T>,
         hotkey: &AccountIdOf<T>,
@@ -155,9 +157,9 @@ where
         coldkey: &AccountIdOf<T>,
         alpha_vec: &[(AccountIdOf<T>, NetUid)],
         tao_amount: TaoBalance,
-    ) -> (AlphaBalance, TaoBalance, NetUid) {
+    ) -> Result<(AlphaBalance, TaoBalance, NetUid), TransactionValidityError> {
         if alpha_vec.len() != 1 {
-            return (0.into(), 0.into(), NetUid::ROOT);
+            return Ok((0.into(), 0.into(), NetUid::ROOT));
         }
 
         if let Some((hotkey, netuid)) = alpha_vec.first() {
@@ -176,26 +178,33 @@ where
 
             // Sell alpha_fee and burn received tao (ignore unstake_from_subnet return).
             if let Some(author) = T::author() {
-                let swap_result = pallet_subtensor::Pallet::<T>::unstake_from_subnet(
-                    hotkey,
-                    coldkey,
-                    &author,
-                    *netuid,
-                    alpha_fee,
-                    0.into(),
-                    true,
-                );
-                if let Ok(tao_amount) = swap_result {
-                    (alpha_fee, tao_amount, *netuid)
-                } else {
-                    (0.into(), 0.into(), NetUid::ROOT)
-                }
+                with_transaction(
+                    || -> TransactionOutcome<Result<TaoBalance, DispatchError>> {
+                        match pallet_subtensor::Pallet::<T>::unstake_from_subnet(
+                            hotkey,
+                            coldkey,
+                            &author,
+                            *netuid,
+                            alpha_fee,
+                            0.into(),
+                            true,
+                        ) {
+                            Ok(tao_amount) => TransactionOutcome::Commit(Ok(tao_amount)),
+                            Err(err) => TransactionOutcome::Rollback(Err(err)),
+                        }
+                    },
+                )
+                .map(|tao_amount| (alpha_fee, tao_amount, *netuid))
+                .map_err(|err| {
+                    log::error!("Error withdrawing transaction fee in alpha: {err:?}");
+                    InvalidTransaction::Payment.into()
+                })
             } else {
                 // Fallback: no author => no fees (do nothing)
-                (0.into(), 0.into(), NetUid::ROOT)
+                Ok((0.into(), 0.into(), NetUid::ROOT))
             }
         } else {
-            (0.into(), 0.into(), NetUid::ROOT)
+            Ok((0.into(), 0.into(), NetUid::ROOT))
         }
     }
 
@@ -343,7 +352,7 @@ where
                 if !alpha_vec.is_empty() {
                     let fee_u64: u64 = fee.saturated_into::<u64>();
                     let (alpha_fee, tao_amount, netuid) =
-                        OU::withdraw_in_alpha(who, &alpha_vec, fee_u64.into());
+                        OU::withdraw_in_alpha(who, &alpha_vec, fee_u64.into())?;
                     return Ok(Some(WithdrawnFee::Alpha((alpha_fee, tao_amount, netuid))));
                 }
                 Err(InvalidTransaction::Payment.into())

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

@@ -3,6 +3,7 @@ use crate::{AlphaFeeHandler, SubtensorTxFeeHandler, TransactionFeeHandler, Trans
 use approx::assert_abs_diff_eq;
 use frame_support::dispatch::GetDispatchInfo;
 use frame_support::pallet_prelude::Zero;
+use frame_support::traits::Currency;
 use frame_support::{assert_err, assert_ok};
 use pallet_subtensor_swap::AlphaSqrtPrice;
 use sp_runtime::{
@@ -187,7 +188,7 @@ fn test_rejects_multi_subnet_alpha_fee_deduction() {
                 &alpha_vec,
                 1.into(),
             ),
-            (0.into(), 0.into(), NetUid::ROOT)
+            Ok((0.into(), 0.into(), NetUid::ROOT))
         );
 
         let alpha_after_0 = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
@@ -326,6 +327,85 @@ fn test_remove_stake_fees_alpha() {
     });
 }
 
+#[test]
+fn test_alpha_fee_withdraw_failure_aborts_and_rolls_back() {
+    new_test_ext().execute_with(|| {
+        let stake_amount = TAO;
+        let unstake_amount = AlphaBalance::from(TAO / 50);
+        let sn = setup_subnets(1, 1);
+        let netuid = sn.subnets[0].netuid;
+        let hotkey = sn.hotkeys[0];
+
+        setup_stake(netuid, &sn.coldkey, &hotkey, stake_amount);
+
+        let current_balance = Balances::free_balance(sn.coldkey);
+        remove_balance_from_coldkey_account(
+            &sn.coldkey,
+            current_balance - ExistentialDeposit::get(),
+        );
+
+        // Force the alpha-fee unstake to fail after AMM bookkeeping by draining
+        // the subnet account used by transfer_tao_from_subnet.
+        let subnet_account = SubtensorModule::get_subnet_account_id(netuid).unwrap();
+        Balances::make_free_balance_be(&subnet_account, 0.into());
+
+        let block_builder = U256::from(MOCK_BLOCK_BUILDER);
+        let block_builder_balance_before = Balances::free_balance(block_builder);
+        let signer_balance_before = Balances::free_balance(sn.coldkey);
+        let alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &sn.coldkey,
+            netuid,
+        );
+        let subnet_alpha_in_before = SubnetAlphaIn::<Test>::get(netuid);
+        let subnet_alpha_out_before = SubnetAlphaOut::<Test>::get(netuid);
+        let subnet_tao_before = SubnetTAO::<Test>::get(netuid);
+        let total_stake_before = TotalStake::<Test>::get();
+        let subnet_volume_before = SubnetVolume::<Test>::get(netuid);
+
+        let call = RuntimeCall::SubtensorModule(pallet_subtensor::Call::remove_stake {
+            hotkey,
+            netuid,
+            amount_unstaked: unstake_amount,
+        });
+        let info = call.get_dispatch_info();
+        let ext = pallet_transaction_payment::ChargeTransactionPayment::<Test>::from(0.into());
+
+        let result =
+            ext.dispatch_transaction(RuntimeOrigin::signed(sn.coldkey).into(), call, &info, 0, 0);
+
+        assert_eq!(
+            result.unwrap_err(),
+            TransactionValidityError::Invalid(InvalidTransaction::Payment)
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                &sn.coldkey,
+                netuid,
+            ),
+            alpha_before
+        );
+        assert_eq!(SubnetAlphaIn::<Test>::get(netuid), subnet_alpha_in_before);
+        assert_eq!(SubnetAlphaOut::<Test>::get(netuid), subnet_alpha_out_before);
+        assert_eq!(SubnetTAO::<Test>::get(netuid), subnet_tao_before);
+        assert_eq!(TotalStake::<Test>::get(), total_stake_before);
+        assert_eq!(SubnetVolume::<Test>::get(netuid), subnet_volume_before);
+        assert_eq!(Balances::free_balance(sn.coldkey), signer_balance_before);
+        assert_eq!(
+            Balances::free_balance(block_builder),
+            block_builder_balance_before
+        );
+
+        assert!(!System::events().iter().any(|event_record| {
+            matches!(
+                &event_record.event,
+                RuntimeEvent::SubtensorModule(SubtensorEvent::TransactionFeePaidWithAlpha { .. })
+            )
+        }));
+    });
+}
+
 // Test that unstaking on root with no free balance results in charging fees from
 // staked amount
 //