Merge pull request #1386 from opentensor/fix/staking-fees

sam0x17 · web-flow · commit 1b40d1de351f · 2025-03-10T21:13:35.000-04:00
Fix/staking fees
diff --git a/pallets/subtensor/src/staking/add_stake.rs b/pallets/subtensor/src/staking/add_stake.rs
@@ -1,5 +1,5 @@
 use super::*;
-use sp_core::Get;
+use substrate_fixed::types::I96F32;
 
 impl<T: Config> Pallet<T> {
     /// ---- The implementation for the extrinsic add_stake: Adds stake to a hotkey account.
@@ -58,13 +58,19 @@ impl<T: Config> Pallet<T> {
         )?;
 
         // 3. Ensure the remove operation from the coldkey is a success.
-        let tao_staked: u64 =
-            Self::remove_balance_from_coldkey_account(&coldkey, stake_to_be_added)?;
+        let tao_staked: I96F32 =
+            Self::remove_balance_from_coldkey_account(&coldkey, stake_to_be_added)?.into();
 
         // 4. Swap the stake into alpha on the subnet and increase counters.
         // Emit the staking event.
         let fee = DefaultStakingFee::<T>::get();
-        Self::stake_into_subnet(&hotkey, &coldkey, netuid, tao_staked, fee);
+        Self::stake_into_subnet(
+            &hotkey,
+            &coldkey,
+            netuid,
+            tao_staked.saturating_to_num::<u64>(),
+            fee,
+        );
 
         // Ok and return.
         Ok(())
@@ -148,12 +154,19 @@ impl<T: Config> Pallet<T> {
         }
 
         // 5. Ensure the remove operation from the coldkey is a success.
-        let tao_staked: u64 = Self::remove_balance_from_coldkey_account(&coldkey, possible_stake)?;
+        let tao_staked: I96F32 =
+            Self::remove_balance_from_coldkey_account(&coldkey, possible_stake)?.into();
 
         // 6. Swap the stake into alpha on the subnet and increase counters.
         // Emit the staking event.
         let fee = DefaultStakingFee::<T>::get();
-        Self::stake_into_subnet(&hotkey, &coldkey, netuid, tao_staked, fee);
+        Self::stake_into_subnet(
+            &hotkey,
+            &coldkey,
+            netuid,
+            tao_staked.saturating_to_num::<u64>(),
+            fee,
+        );
 
         // Ok and return.
         Ok(())
diff --git a/pallets/subtensor/src/staking/move_stake.rs b/pallets/subtensor/src/staking/move_stake.rs
@@ -1,7 +1,7 @@
 use super::*;
 use safe_math::*;
 use sp_core::Get;
-use substrate_fixed::types::U64F64;
+use substrate_fixed::types::{I96F32, U64F64};
 
 impl<T: Config> Pallet<T> {
     /// Moves stake from one hotkey to another across subnets.
@@ -338,7 +338,13 @@ impl<T: Config> Pallet<T> {
         };
 
         // Unstake from the origin subnet, returning TAO (or a 1:1 equivalent).
-        let fee = DefaultStakingFee::<T>::get().safe_div(2);
+        let fee = Self::calculate_staking_fee(
+            origin_netuid,
+            origin_hotkey,
+            I96F32::saturating_from_num(alpha_amount),
+        )
+        .safe_div(2);
+
         let tao_unstaked = Self::unstake_from_subnet(
             origin_hotkey,
             origin_coldkey,
diff --git a/pallets/subtensor/src/staking/remove_stake.rs b/pallets/subtensor/src/staking/remove_stake.rs
@@ -1,5 +1,5 @@
 use super::*;
-use sp_core::Get;
+use substrate_fixed::types::I96F32;
 
 impl<T: Config> Pallet<T> {
     /// ---- The implementation for the extrinsic remove_stake: Removes stake from a hotkey account and adds it onto a coldkey.
@@ -58,7 +58,11 @@ impl<T: Config> Pallet<T> {
         )?;
 
         // 3. Swap the alpba to tao and update counters for this subnet.
-        let fee = DefaultStakingFee::<T>::get();
+        let fee = Self::calculate_staking_fee(
+            netuid,
+            &hotkey,
+            I96F32::saturating_from_num(alpha_unstaked),
+        );
         let tao_unstaked: u64 =
             Self::unstake_from_subnet(&hotkey, &coldkey, netuid, alpha_unstaked, fee);
 
@@ -109,8 +113,6 @@ impl<T: Config> Pallet<T> {
         origin: T::RuntimeOrigin,
         hotkey: T::AccountId,
     ) -> dispatch::DispatchResult {
-        let fee = DefaultStakingFee::<T>::get();
-
         // 1. We check the transaction is signed by the caller and retrieve the T::AccountId coldkey information.
         let coldkey = ensure_signed(origin)?;
         log::info!("do_unstake_all( origin:{:?} hotkey:{:?} )", coldkey, hotkey);
@@ -126,20 +128,26 @@ impl<T: Config> Pallet<T> {
         log::debug!("All subnet netuids: {:?}", netuids);
 
         // 4. Iterate through all subnets and remove stake.
-        for netuid in netuids.iter() {
+        for netuid in netuids.into_iter() {
             // Ensure that the hotkey has enough stake to withdraw.
             let alpha_unstaked =
-                Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, *netuid);
+                Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid);
+            let fee = Self::calculate_staking_fee(
+                netuid,
+                &hotkey,
+                I96F32::saturating_from_num(alpha_unstaked),
+            );
+
             if alpha_unstaked > 0 {
                 // Swap the alpha to tao and update counters for this subnet.
                 let tao_unstaked: u64 =
-                    Self::unstake_from_subnet(&hotkey, &coldkey, *netuid, alpha_unstaked, fee);
+                    Self::unstake_from_subnet(&hotkey, &coldkey, netuid, alpha_unstaked, fee);
 
                 // Add the balance to the coldkey. If the above fails we will not credit this coldkey.
                 Self::add_balance_to_coldkey_account(&coldkey, tao_unstaked);
 
                 // If the stake is below the minimum, we clear the nomination from storage.
-                Self::clear_small_nomination_if_required(&hotkey, &coldkey, *netuid);
+                Self::clear_small_nomination_if_required(&hotkey, &coldkey, netuid);
             }
         }
 
@@ -177,8 +185,6 @@ impl<T: Config> Pallet<T> {
         origin: T::RuntimeOrigin,
         hotkey: T::AccountId,
     ) -> dispatch::DispatchResult {
-        let fee = DefaultStakingFee::<T>::get();
-
         // 1. We check the transaction is signed by the caller and retrieve the T::AccountId coldkey information.
         let coldkey = ensure_signed(origin)?;
         log::info!("do_unstake_all( origin:{:?} hotkey:{:?} )", coldkey, hotkey);
@@ -195,22 +201,28 @@ impl<T: Config> Pallet<T> {
 
         // 4. Iterate through all subnets and remove stake.
         let mut total_tao_unstaked: u64 = 0;
-        for netuid in netuids.iter() {
+        for netuid in netuids.into_iter() {
             // If not Root network.
-            if *netuid != Self::get_root_netuid() {
+            if netuid != Self::get_root_netuid() {
                 // Ensure that the hotkey has enough stake to withdraw.
                 let alpha_unstaked =
-                    Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, *netuid);
+                    Self::get_stake_for_hotkey_and_coldkey_on_subnet(&hotkey, &coldkey, netuid);
+                let fee = Self::calculate_staking_fee(
+                    netuid,
+                    &hotkey,
+                    I96F32::saturating_from_num(alpha_unstaked),
+                );
+
                 if alpha_unstaked > 0 {
                     // Swap the alpha to tao and update counters for this subnet.
-                    let tao_unstaked: u64 =
-                        Self::unstake_from_subnet(&hotkey, &coldkey, *netuid, alpha_unstaked, fee);
+                    let tao_unstaked =
+                        Self::unstake_from_subnet(&hotkey, &coldkey, netuid, alpha_unstaked, fee);
 
                     // Increment total
                     total_tao_unstaked = total_tao_unstaked.saturating_add(tao_unstaked);
 
                     // If the stake is below the minimum, we clear the nomination from storage.
-                    Self::clear_small_nomination_if_required(&hotkey, &coldkey, *netuid);
+                    Self::clear_small_nomination_if_required(&hotkey, &coldkey, netuid);
                 }
             }
         }
@@ -302,8 +314,12 @@ impl<T: Config> Pallet<T> {
         )?;
 
         // 4. Swap the alpha to tao and update counters for this subnet.
-        let fee = DefaultStakingFee::<T>::get();
-        let tao_unstaked: u64 =
+        let fee = Self::calculate_staking_fee(
+            netuid,
+            &hotkey,
+            I96F32::saturating_from_num(alpha_unstaked),
+        );
+        let tao_unstaked =
             Self::unstake_from_subnet(&hotkey, &coldkey, netuid, possible_alpha, fee);
 
         // 5. We add the balance to the coldkey. If the above fails we will not credit this coldkey.
diff --git a/pallets/subtensor/src/staking/stake_utils.rs b/pallets/subtensor/src/staking/stake_utils.rs
@@ -806,7 +806,7 @@ impl<T: Config> Pallet<T> {
     /// Stakes TAO into a subnet for a given hotkey and coldkey pair.
     ///
     /// We update the pools associated with a subnet as well as update hotkey alpha shares.
-    pub fn stake_into_subnet(
+    pub(crate) fn stake_into_subnet(
         hotkey: &T::AccountId,
         coldkey: &T::AccountId,
         netuid: u16,
@@ -1066,6 +1066,28 @@ impl<T: Config> Pallet<T> {
 
         Ok(())
     }
+
+    pub(crate) fn calculate_staking_fee(
+        netuid: u16,
+        hotkey: &T::AccountId,
+        alpha_estimate: I96F32,
+    ) -> u64 {
+        if (netuid == Self::get_root_netuid()) || (SubnetMechanism::<T>::get(netuid)) == 0 {
+            DefaultStakingFee::<T>::get()
+        } else {
+            let fee = alpha_estimate
+                .saturating_mul(
+                    I96F32::saturating_from_num(AlphaDividendsPerSubnet::<T>::get(netuid, &hotkey))
+                        .safe_div(I96F32::saturating_from_num(TotalHotkeyAlpha::<T>::get(
+                            &hotkey, netuid,
+                        ))),
+                )
+                .saturating_mul(Self::get_alpha_price(netuid)) // fee needs to be in TAO
+                .saturating_to_num::<u64>();
+
+            fee.max(DefaultStakingFee::<T>::get())
+        }
+    }
 }
 
 ///////////////////////////////////////////
diff --git a/pallets/subtensor/src/tests/staking.rs b/pallets/subtensor/src/tests/staking.rs
@@ -2219,11 +2219,71 @@ fn test_remove_stake_fee_goes_to_subnet_tao() {
     });
 }
 
+// cargo test --package pallet-subtensor --lib -- tests::staking::test_remove_stake_fee_realistic_values --exact --show-output --nocapture
 #[test]
-fn test_stake_below_min_validate() {
-    // Testing the signed extension validate function
-    // correctly filters the `add_stake` transaction.
+fn test_remove_stake_fee_realistic_values() {
+    new_test_ext(1).execute_with(|| {
+        let subnet_owner_coldkey = U256::from(1001);
+        let subnet_owner_hotkey = U256::from(1002);
+        let hotkey = U256::from(2);
+        let coldkey = U256::from(3);
+        let alpha_to_unstake = 111_180_000_000;
+        let alpha_divs = 2_816_190;
+
+        let netuid = add_dynamic_network(&subnet_owner_hotkey, &subnet_owner_coldkey);
+        SubtensorModule::create_account_if_non_existent(&coldkey, &hotkey);
+
+        // Mock a realistic scenario:
+        //   Subnet 1 has 3896 TAO and 128_011 Alpha in reserves, which
+        //   makes its price ~0.03.
+        //   A hotkey has 111 Alpha stake and is unstaking all Alpha.
+        //   Alpha dividends of this hotkey are ~0.0028
+        //   This makes fee be equal ~0.0028 Alpha ~= 84000 rao
+        let tao_reserve: U96F32 = U96F32::from_num(3_896_056_559_708_u64);
+        let alpha_in: U96F32 = U96F32::from_num(128_011_331_299_964_u64);
+        SubnetTAO::<Test>::insert(netuid, tao_reserve.to_num::<u64>());
+        SubnetAlphaIn::<Test>::insert(netuid, alpha_in.to_num::<u64>());
+        AlphaDividendsPerSubnet::<Test>::insert(netuid, hotkey, alpha_divs);
+        let current_price = SubtensorModule::get_alpha_price(netuid).to_num::<f64>();
+
+        // Add stake first time to init TotalHotkeyAlpha
+        SubtensorModule::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &coldkey,
+            netuid,
+            alpha_to_unstake,
+        );
+
+        // Estimate fees
+        let expected_fee: f64 = current_price * alpha_divs as f64;
+
+        // Remove stake to measure fee
+        let balance_before = SubtensorModule::get_coldkey_balance(&coldkey);
+        let expected_tao_no_fee =
+            SubtensorModule::sim_swap_alpha_for_tao(netuid, alpha_to_unstake).unwrap();
+
+        assert_ok!(SubtensorModule::remove_stake(
+            RuntimeOrigin::signed(coldkey),
+            hotkey,
+            netuid,
+            alpha_to_unstake
+        ));
 
+        // Calculate expected fee
+        let balance_after = SubtensorModule::get_coldkey_balance(&coldkey);
+        let actual_fee = expected_tao_no_fee as f64 - (balance_after - balance_before) as f64;
+        log::info!("Actual fee: {:?}", actual_fee);
+
+        assert_abs_diff_eq!(
+            actual_fee as u64,
+            expected_fee as u64,
+            epsilon = expected_fee as u64 / 1000
+        );
+    });
+}
+
+#[test]
+fn test_stake_below_min_validate() {
     new_test_ext(0).execute_with(|| {
         let subnet_owner_coldkey = U256::from(1001);
         let subnet_owner_hotkey = U256::from(1002);
diff --git a/pallets/subtensor/src/tests/swap_coldkey.rs b/pallets/subtensor/src/tests/swap_coldkey.rs
@@ -14,6 +14,7 @@ use frame_support::traits::schedule::DispatchTime;
 use frame_support::traits::schedule::v3::Named as ScheduleNamed;
 use sp_core::{Get, H256, U256};
 use sp_runtime::DispatchError;
+use substrate_fixed::types::I96F32;
 
 // // SKIP_WASM_BUILD=1 RUST_LOG=debug cargo test --test swap_coldkey -- test_swap_total_hotkey_coldkey_stakes_this_interval --exact --nocapture
 // #[test]
@@ -537,7 +538,9 @@ fn test_swap_concurrent_modifications() {
         let netuid: u16 = 1;
         let initial_stake = 1_000_000_000_000;
         let additional_stake = 500_000_000_000;
-        let fee = DefaultStakingFee::<Test>::get();
+        let initial_stake_alpha =
+            I96F32::from(initial_stake).saturating_mul(SubtensorModule::get_alpha_price(netuid));
+        let fee = SubtensorModule::calculate_staking_fee(netuid, &hotkey, initial_stake_alpha);
 
         // Setup initial state
         add_network(netuid, 1, 1);
@@ -588,15 +591,14 @@ fn test_swap_concurrent_modifications() {
             &mut weight
         ));
 
-        let eps = 500; // RAO
         assert_abs_diff_eq!(
             SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(
                 &hotkey,
                 &new_coldkey,
                 netuid
             ),
             stake_before_swap + additional_stake - fee,
-            epsilon = eps
+            epsilon = (stake_before_swap + additional_stake - fee) / 1000
         );
         assert!(!Alpha::<Test>::contains_key((hotkey, old_coldkey, netuid)));
     });
