Handle tao transfer errors in root claim

pallets/subtensor/src/macros/dispatches.rs

@@ -2177,7 +2177,7 @@ mod dispatches {
 
             Self::maybe_add_coldkey_index(&coldkey);
 
-            let weight = Self::do_root_claim(coldkey, Some(subnets));
+            let weight = Self::do_root_claim(coldkey, Some(subnets))?;
             Ok((Some(weight), Pays::Yes).into())
         }
 

pallets/subtensor/src/staking/claim_root.rs

@@ -1,6 +1,9 @@
 use super::*;
+use frame_support::dispatch::DispatchResult;
+use frame_support::storage::{TransactionOutcome, with_transaction};
 use frame_support::weights::Weight;
 use sp_core::Get;
+use sp_runtime::DispatchError;
 use sp_std::collections::btree_set::BTreeSet;
 use substrate_fixed::types::I96F32;
 use subtensor_swap_interface::SwapHandler;
@@ -130,7 +133,7 @@ impl<T: Config> Pallet<T> {
         netuid: NetUid,
         root_claim_type: RootClaimTypeEnum,
         ignore_minimum_condition: bool,
-    ) {
+    ) -> DispatchResult {
         // Subtract the root claimed.
         let owed: I96F32 = Self::get_root_owed_for_hotkey_coldkey_float(hotkey, coldkey, netuid);
 
@@ -140,7 +143,7 @@ impl<T: Config> Pallet<T> {
             log::debug!(
                 "root claim on subnet {netuid} is skipped: {owed:?} for h={hotkey:?},c={coldkey:?} "
             );
-            return; // no-op
+            return Ok(()); // no-op
         }
 
         // Convert owed to u64, mapping negative values to 0
@@ -154,7 +157,7 @@ impl<T: Config> Pallet<T> {
             log::debug!(
                 "root claim on subnet {netuid} is skipped: {owed:?} for h={hotkey:?},c={coldkey:?}"
             );
-            return; // no-op
+            return Ok(()); // no-op
         }
 
         let swap = match root_claim_type {
@@ -164,38 +167,48 @@ impl<T: Config> Pallet<T> {
         };
 
         if swap {
-            // Increase stake on root. Swap the alpha owed to TAO
-            let owed_tao = match Self::swap_alpha_for_tao(
-                netuid,
-                owed_u64.into(),
-                T::SwapInterface::min_price::<TaoBalance>(),
-                true,
-            ) {
-                Ok(owed_tao) => owed_tao,
-                Err(err) => {
-                    log::error!("Error swapping alpha for TAO: {err:?}");
-
-                    return;
-                }
-            };
-
-            // Record root sell as protocol outflow (reduces protocol cost).
-            let root_sell_tao: TaoBalance = owed_tao.amount_paid_out;
-            SubnetRootSellTao::<T>::mutate(netuid, |total| {
-                *total = total.saturating_add(root_sell_tao);
-            });
-            Self::record_protocol_outflow(netuid, root_sell_tao);
-
-            // Transfer unstaked TAO from subnet account to the root subnet account
-            // and increase root stake.
-            if let Some(root_subnet_account_id) = Self::get_subnet_account_id(NetUid::ROOT)
-                && Self::transfer_tao_from_subnet(
+            with_transaction(|| {
+                // Increase stake on root. Swap the alpha owed to TAO.
+                let owed_tao = match Self::swap_alpha_for_tao(
+                    netuid,
+                    owed_u64.into(),
+                    T::SwapInterface::min_price::<TaoBalance>(),
+                    true,
+                ) {
+                    Ok(owed_tao) => owed_tao,
+                    Err(err) => {
+                        log::error!("Error swapping alpha for TAO: {err:?}");
+
+                        return TransactionOutcome::Rollback(Err(err));
+                    }
+                };
+
+                let root_subnet_account_id = match Self::get_subnet_account_id(NetUid::ROOT) {
+                    Some(account_id) => account_id,
+                    None => {
+                        return TransactionOutcome::Rollback(Err(
+                            Error::<T>::RootNetworkDoesNotExist.into(),
+                        ));
+                    }
+                };
+
+                if let Err(err) = Self::transfer_tao_from_subnet(
                     netuid,
                     &root_subnet_account_id,
                     owed_tao.amount_paid_out.into(),
-                )
-                .is_ok()
-            {
+                ) {
+                    log::error!("Error transferring root claim TAO from subnet: {err:?}");
+
+                    return TransactionOutcome::Rollback(Err(err));
+                }
+
+                // Record root sell as protocol outflow (reduces protocol cost).
+                let root_sell_tao: TaoBalance = owed_tao.amount_paid_out;
+                SubnetRootSellTao::<T>::mutate(netuid, |total| {
+                    *total = total.saturating_add(root_sell_tao);
+                });
+                Self::record_protocol_outflow(netuid, root_sell_tao);
+
                 Self::increase_stake_for_hotkey_and_coldkey_on_subnet(
                     hotkey,
                     coldkey,
@@ -217,13 +230,15 @@ impl<T: Config> Pallet<T> {
                 TotalStake::<T>::mutate(|total| {
                     *total = total.saturating_add(owed_tao.amount_paid_out.into());
                 });
-            }
 
-            Self::add_stake_adjust_root_claimed_for_hotkey_and_coldkey(
-                hotkey,
-                coldkey,
-                owed_tao.amount_paid_out.into(),
-            );
+                Self::add_stake_adjust_root_claimed_for_hotkey_and_coldkey(
+                    hotkey,
+                    coldkey,
+                    owed_tao.amount_paid_out.into(),
+                );
+
+                TransactionOutcome::Commit(Ok(()))
+            })?;
         } else
         /* Keep */
         {
@@ -240,6 +255,8 @@ impl<T: Config> Pallet<T> {
         RootClaimed::<T>::mutate((netuid, hotkey, coldkey), |root_claimed| {
             *root_claimed = root_claimed.saturating_add(owed_u64.into());
         });
+
+        Ok(())
     }
 
     fn root_claim_on_subnet_weight(_root_claim_type: RootClaimTypeEnum) -> Weight {
@@ -251,7 +268,7 @@ impl<T: Config> Pallet<T> {
         hotkey: &T::AccountId,
         coldkey: &T::AccountId,
         subnets: Option<BTreeSet<NetUid>>,
-    ) -> Weight {
+    ) -> Result<Weight, DispatchError> {
         let mut weight = Weight::default();
 
         let root_claim_type = RootClaimType::<T>::get(coldkey);
@@ -271,11 +288,11 @@ impl<T: Config> Pallet<T> {
                 continue;
             }
 
-            Self::root_claim_on_subnet(hotkey, coldkey, *netuid, root_claim_type.clone(), false);
+            Self::root_claim_on_subnet(hotkey, coldkey, *netuid, root_claim_type.clone(), false)?;
             weight.saturating_accrue(Self::root_claim_on_subnet_weight(root_claim_type.clone()));
         }
 
-        weight
+        Ok(weight)
     }
 
     pub fn add_stake_adjust_root_claimed_for_hotkey_and_coldkey(
@@ -328,20 +345,33 @@ impl<T: Config> Pallet<T> {
         }
     }
 
-    pub fn do_root_claim(coldkey: T::AccountId, subnets: Option<BTreeSet<NetUid>>) -> Weight {
+    pub fn do_root_claim(
+        coldkey: T::AccountId,
+        subnets: Option<BTreeSet<NetUid>>,
+    ) -> Result<Weight, DispatchError> {
+        with_transaction(|| match Self::try_do_root_claim(coldkey, subnets) {
+            Ok(weight) => TransactionOutcome::Commit(Ok(weight)),
+            Err(err) => TransactionOutcome::Rollback(Err(err)),
+        })
+    }
+
+    fn try_do_root_claim(
+        coldkey: T::AccountId,
+        subnets: Option<BTreeSet<NetUid>>,
+    ) -> Result<Weight, DispatchError> {
         let mut weight = Weight::default();
 
         let hotkeys = StakingHotkeys::<T>::get(&coldkey);
         weight.saturating_accrue(T::DbWeight::get().reads(1));
 
-        hotkeys.iter().for_each(|hotkey| {
+        for hotkey in hotkeys.iter() {
             weight.saturating_accrue(T::DbWeight::get().reads(1));
-            weight.saturating_accrue(Self::root_claim_all(hotkey, &coldkey, subnets.clone()));
-        });
+            weight.saturating_accrue(Self::root_claim_all(hotkey, &coldkey, subnets.clone())?);
+        }
 
         Self::deposit_event(Event::RootClaimed { coldkey });
 
-        weight
+        Ok(weight)
     }
 
     fn block_hash_to_indices_weight(k: u64, _n: u64) -> Weight {
@@ -371,10 +401,11 @@ impl<T: Config> Pallet<T> {
         for i in coldkeys_to_claim.iter() {
             weight.saturating_accrue(T::DbWeight::get().reads(1));
             if let Ok(coldkey) = StakingColdkeysByIndex::<T>::try_get(i) {
-                weight.saturating_accrue(Self::do_root_claim(coldkey.clone(), None));
+                match Self::do_root_claim(coldkey.clone(), None) {
+                    Ok(claim_weight) => weight.saturating_accrue(claim_weight),
+                    Err(err) => log::error!("Error auto-claiming root dividends: {err:?}"),
+                }
             }
-
-            continue;
         }
 
         weight

pallets/subtensor/src/tests/claim_root.rs

@@ -1,18 +1,19 @@
-#![allow(clippy::expect_used)]
+#![allow(clippy::expect_used, clippy::unwrap_used)]
 
 use crate::RootAlphaDividendsPerSubnet;
 use crate::tests::mock::*;
 use crate::{
     DefaultMinRootClaimAmount, Error, MAX_NUM_ROOT_CLAIMS, MAX_ROOT_CLAIM_THRESHOLD, NetworksAdded,
     NumRootClaim, NumStakingColdkeys, PendingRootAlphaDivs, RootClaimable, RootClaimableThreshold,
-    StakingColdkeys, StakingColdkeysByIndex, SubnetAlphaIn, SubnetMechanism, SubnetMovingPrice,
-    SubnetTAO, SubnetTaoFlow, SubtokenEnabled, Tempo, pallet,
+    StakingColdkeys, StakingColdkeysByIndex, SubnetAlphaIn, SubnetAlphaOut, SubnetMechanism,
+    SubnetMovingPrice, SubnetProtocolFlow, SubnetRootSellTao, SubnetTAO, SubnetTaoFlow,
+    SubnetVolume, SubtokenEnabled, Tempo, TotalStake, pallet,
 };
 use crate::{RootClaimType, RootClaimTypeEnum, RootClaimed};
 use approx::assert_abs_diff_eq;
 use frame_support::dispatch::RawOrigin;
 use frame_support::pallet_prelude::Weight;
-use frame_support::traits::Get;
+use frame_support::traits::{Currency, Get};
 use frame_support::{assert_err, assert_noop, assert_ok};
 use sp_core::{H256, U256};
 use sp_runtime::DispatchError;
@@ -757,6 +758,100 @@ fn test_claim_root_with_drain_emissions_and_swap_claim_type() {
     });
 }
 
+#[test]
+fn test_claim_root_swap_failure_does_not_consume_claim() {
+    new_test_ext(1).execute_with(|| {
+        let owner_coldkey = U256::from(1001);
+        let other_coldkey = U256::from(10010);
+        let hotkey = U256::from(1002);
+        let coldkey = U256::from(1003);
+        let netuid = add_dynamic_network(&hotkey, &owner_coldkey);
+
+        SubtensorModule::set_tao_weight(u64::MAX);
+        SubnetTAO::<Test>::insert(netuid, TaoBalance::from(50_000_000_000_u64));
+        SubnetAlphaIn::<Test>::insert(netuid, AlphaBalance::from(100_000_000_000_u64));
+
+        mock_increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &coldkey,
+            NetUid::ROOT,
+            2_000_000_u64.into(),
+        );
+        mock_increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &other_coldkey,
+            NetUid::ROOT,
+            18_000_000_u64.into(),
+        );
+        mock_increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &owner_coldkey,
+            netuid,
+            10_000_000_u64.into(),
+        );
+
+        SubtensorModule::distribute_emission(
+            netuid,
+            AlphaBalance::ZERO,
+            AlphaBalance::ZERO,
+            10_000_000_u64.into(),
+            AlphaBalance::ZERO,
+        );
+
+        assert_ok!(SubtensorModule::set_root_claim_type(
+            RuntimeOrigin::signed(coldkey),
+            RootClaimTypeEnum::Swap
+        ));
+
+        let subnet_account = SubtensorModule::get_subnet_account_id(netuid).unwrap();
+        Balances::make_free_balance_be(&subnet_account, 0.into());
+
+        let root_claimed_before = RootClaimed::<Test>::get((netuid, &hotkey, &coldkey));
+        let root_stake_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &coldkey,
+            NetUid::ROOT,
+        );
+        let subnet_tao_before = SubnetTAO::<Test>::get(netuid);
+        let root_subnet_tao_before = SubnetTAO::<Test>::get(NetUid::ROOT);
+        let subnet_alpha_in_before = SubnetAlphaIn::<Test>::get(netuid);
+        let subnet_alpha_out_before = SubnetAlphaOut::<Test>::get(netuid);
+        let total_stake_before = TotalStake::<Test>::get();
+        let subnet_volume_before = SubnetVolume::<Test>::get(netuid);
+        let root_sell_before = SubnetRootSellTao::<Test>::get(netuid);
+        let protocol_flow_before = SubnetProtocolFlow::<Test>::get(netuid);
+
+        assert_noop!(
+            SubtensorModule::claim_root(RuntimeOrigin::signed(coldkey), BTreeSet::from([netuid])),
+            Error::<Test>::InsufficientBalance
+        );
+
+        assert_eq!(
+            RootClaimed::<Test>::get((netuid, &hotkey, &coldkey)),
+            root_claimed_before
+        );
+        assert_eq!(
+            SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                &coldkey,
+                NetUid::ROOT,
+            ),
+            root_stake_before
+        );
+        assert_eq!(SubnetTAO::<Test>::get(netuid), subnet_tao_before);
+        assert_eq!(SubnetTAO::<Test>::get(NetUid::ROOT), root_subnet_tao_before);
+        assert_eq!(SubnetAlphaIn::<Test>::get(netuid), subnet_alpha_in_before);
+        assert_eq!(SubnetAlphaOut::<Test>::get(netuid), subnet_alpha_out_before);
+        assert_eq!(TotalStake::<Test>::get(), total_stake_before);
+        assert_eq!(SubnetVolume::<Test>::get(netuid), subnet_volume_before);
+        assert_eq!(SubnetRootSellTao::<Test>::get(netuid), root_sell_before);
+        assert_eq!(
+            SubnetProtocolFlow::<Test>::get(netuid),
+            protocol_flow_before
+        );
+    });
+}
+
 #[test]
 fn test_claim_root_with_run_coinbase() {
     new_test_ext(1).execute_with(|| {