Handle balancer misbalance on injection

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -85,11 +85,6 @@ impl<T: Config> Pallet<T> {
                 let tao_to_swap_with: TaoBalance =
                     tou64!(excess_tao.get(netuid_i).unwrap_or(&asfloat!(0))).into();
 
-                // Inject tao and alpha into protocol liquidity. In theorry, it may not always
-                // be a success (returned values are 0s) in case of high liquidity disbalance
-                let (actual_injected_tao, actual_injected_alpha) =
-                    T::SwapInterface::adjust_protocol_liquidity(*netuid_i, tao_in_i, alpha_in_i);
-
                 // Clear per-block pool-side emission counters up front so a subnet
                 // disabled this block does not display stale values from an earlier block.
                 SubnetExcessTao::<T>::insert(*netuid_i, TaoBalance::ZERO);
@@ -130,38 +125,60 @@ impl<T: Config> Pallet<T> {
                     }
                 }
 
-                // Inject Alpha in.
-                SubnetAlphaInEmission::<T>::insert(*netuid_i, actual_injected_alpha);
-
-                // Mint alpha and resolve to alpha reserve
-                Self::resolve_to_alpha_in(Self::mint_alpha(*netuid_i, actual_injected_alpha));
-
-                // Inject TAO in.
-                if !actual_injected_tao.is_zero() {
-                    match Self::spend_tao(&subnet_account_id, remaining_credit, actual_injected_tao)
-                    {
+                // Materialize this block's TAO before updating balancer reservoir
+                // state. If spending fails, do not let the swap pallet consume
+                // reservoir state as if this block's TAO arrived.
+                let materialized_tao_delta = if tao_in_i.is_zero() {
+                    TaoBalance::ZERO
+                } else {
+                    match Self::spend_tao(&subnet_account_id, remaining_credit, tao_in_i) {
                         Ok(remainder) => {
                             remaining_credit = remainder;
-
-                            SubnetTaoInEmission::<T>::insert(*netuid_i, actual_injected_tao);
-                            SubnetTAO::<T>::mutate(*netuid_i, |total| {
-                                *total = total.saturating_add(actual_injected_tao);
-                            });
-                            TotalStake::<T>::mutate(|total| {
-                                *total = total.saturating_add(actual_injected_tao);
-                            });
-
-                            // Record emission injection as protocol inflow.
-                            Self::record_protocol_inflow(*netuid_i, actual_injected_tao);
+                            tao_in_i
                         }
                         Err(remainder) => {
                             remaining_credit = remainder;
                             let remaining_balance = remaining_credit.peek();
                             log::error!(
-                                "Failed to spend credit: injected_tao = {actual_injected_tao:?}, netuid_i = {netuid_i:?}, remaining_balance = {remaining_balance:?}"
+                                "Failed to spend credit: tao_to_materialize = {tao_in_i:?}, netuid_i = {netuid_i:?}, remaining_balance = {remaining_balance:?}"
                             );
+                            TaoBalance::ZERO
                         }
                     }
+                };
+
+                // Decide which current/reservoir liquidity can become price-active
+                // without pushing balancer weights out of range. Only already
+                // materialized current TAO is offered to the swap pallet.
+                let (
+                    price_active_tao,
+                    _tao_to_materialize,
+                    price_active_alpha,
+                    alpha_to_materialize,
+                ) = T::SwapInterface::adjust_protocol_liquidity(
+                    *netuid_i,
+                    materialized_tao_delta,
+                    alpha_in_i,
+                );
+
+                // Materialize this block's alpha emission, then add only the
+                // price-active portion to the pool reserve. The price-active
+                // portion may include alpha that was materialized in an earlier
+                // block and held in the reservoir.
+                let _ = Self::mint_alpha(*netuid_i, alpha_to_materialize);
+                SubnetAlphaInEmission::<T>::insert(*netuid_i, price_active_alpha);
+                Self::increase_provided_alpha_reserve(*netuid_i, price_active_alpha);
+
+                // Add only the price-active TAO to the pool reserve. This may
+                // include TAO materialized in an earlier block and held in the
+                // reservoir.
+                if !price_active_tao.is_zero() {
+                    SubnetTaoInEmission::<T>::insert(*netuid_i, price_active_tao);
+                    Self::increase_provided_tao_reserve(*netuid_i, price_active_tao);
+                    TotalStake::<T>::mutate(|total| {
+                        *total = total.saturating_add(price_active_tao);
+                    });
+                    Self::record_protocol_inflow(*netuid_i, price_active_tao);
                 }
             }
         }

pallets/subtensor/src/tests/coinbase.rs

@@ -3678,6 +3678,42 @@ fn test_coinbase_inject_and_maybe_swap_does_not_skew_reserves() {
     });
 }
 
+#[test]
+fn test_coinbase_failed_tao_materialization_does_not_activate_current_tao() {
+    new_test_ext(1).execute_with(|| {
+        let netuid = add_dynamic_network(&U256::from(1), &U256::from(2));
+        let initial_reserve = TaoBalance::from(1_000_000_u64);
+        let reservoir_tao = TaoBalance::from(100_u64);
+        let current_tao = TaoBalance::from(200_u64);
+        let current_alpha = AlphaBalance::from(100_u64);
+
+        mock::setup_reserves(netuid, initial_reserve, AlphaBalance::from(1_000_000_u64));
+        Swap::maybe_initialize_palswap(netuid, None);
+        pallet_subtensor_swap::BalancerTaoReservoir::<Test>::insert(netuid, reservoir_tao);
+
+        let tao_in = BTreeMap::from([(netuid, U96F32::saturating_from_num(current_tao))]);
+        let alpha_in = BTreeMap::from([(netuid, U96F32::saturating_from_num(current_alpha))]);
+        let excess_tao = BTreeMap::new();
+        let credit = SubtensorModule::mint_tao(TaoBalance::ZERO);
+
+        SubtensorModule::inject_and_maybe_swap(&[netuid], &tao_in, &alpha_in, &excess_tao, credit);
+
+        assert_eq!(
+            SubnetTAO::<Test>::get(netuid),
+            initial_reserve.saturating_add(reservoir_tao)
+        );
+        assert_eq!(SubnetTaoInEmission::<Test>::get(netuid), reservoir_tao);
+        assert_eq!(
+            SubnetProtocolFlow::<Test>::get(netuid),
+            reservoir_tao.to_u64() as i64
+        );
+        assert_eq!(
+            pallet_subtensor_swap::BalancerTaoReservoir::<Test>::get(netuid),
+            TaoBalance::ZERO
+        );
+    });
+}
+
 #[test]
 fn test_coinbase_drain_pending_increments_blockssincelaststep() {
     new_test_ext(1).execute_with(|| {

pallets/subtensor/src/tests/networks.rs

@@ -370,6 +370,8 @@ fn dissolve_clears_all_per_subnet_storages() {
         SubnetAlphaOutEmission::<Test>::insert(net, AlphaBalance::from(1));
         SubnetTaoInEmission::<Test>::insert(net, TaoBalance::from(1));
         SubnetVolume::<Test>::insert(net, 1u128);
+        pallet_subtensor_swap::BalancerTaoReservoir::<Test>::insert(net, TaoBalance::from(1));
+        pallet_subtensor_swap::BalancerAlphaReservoir::<Test>::insert(net, AlphaBalance::from(1));
 
         // Items now REMOVED (not zeroed) by dissolution
         SubnetAlphaIn::<Test>::insert(net, AlphaBalance::from(2));
@@ -521,6 +523,8 @@ fn dissolve_clears_all_per_subnet_storages() {
         assert!(!SubnetAlphaOutEmission::<Test>::contains_key(net));
         assert!(!SubnetTaoInEmission::<Test>::contains_key(net));
         assert!(!SubnetVolume::<Test>::contains_key(net));
+        assert!(!pallet_subtensor_swap::BalancerTaoReservoir::<Test>::contains_key(net));
+        assert!(!pallet_subtensor_swap::BalancerAlphaReservoir::<Test>::contains_key(net));
 
         // TAO Flow
         assert!(!SubnetTaoFlow::<Test>::contains_key(net));

pallets/swap/src/pallet/impls.rs

@@ -79,41 +79,112 @@ impl<T: Config> Pallet<T> {
         Ok(())
     }
 
-    /// Returns actually added Tao and Alpha, which may be zero in case
-    /// of a high disbalance
+    /// Adjusts balancer weights with minted TAO and alpha liquidity to
+    /// maintain price.
+    ///
+    /// If weights cannot be adjusted (get pushed out of range), the excess TAO
+    /// and/or Alpha are added to reservoirs and an attempt to use them will be made
+    /// later.
+    ///
+    /// Returns:
+    /// 1. price-active TAO delta to add to `SubnetTAO`
+    /// 2. TAO delta to materialize by the caller
+    /// 3. price-active Alpha delta to add to `SubnetAlphaIn`
+    /// 4. Alpha delta to materialize by the caller
+    ///
+    /// Amounts that would push weights out of range are materialized but left in
+    /// per-subnet reservoirs for a later balancer update.
+    ///
+    /// Reservoir amounts may be included in the balancer update, but they are not
+    /// returned for materialization because they were already materialized when
+    /// first stored.
     pub(super) fn adjust_protocol_liquidity(
         netuid: NetUid,
         tao_delta: TaoBalance,
         alpha_delta: AlphaBalance,
-    ) -> (TaoBalance, AlphaBalance) {
+    ) -> (TaoBalance, TaoBalance, AlphaBalance, AlphaBalance) {
         // Get reserves
         let alpha_reserve = T::AlphaReserve::reserve(netuid.into());
         let tao_reserve = T::TaoReserve::reserve(netuid.into());
-        let mut balancer = SwapBalancer::<T>::get(netuid);
+        let balancer = SwapBalancer::<T>::get(netuid);
+
+        let pending_tao = BalancerTaoReservoir::<T>::get(netuid).saturating_add(tao_delta);
+        let pending_alpha = BalancerAlphaReservoir::<T>::get(netuid).saturating_add(alpha_delta);
+
+        if let Some(new_balancer) = Self::try_update_balancer(
+            &balancer,
+            tao_reserve,
+            alpha_reserve,
+            pending_tao,
+            pending_alpha,
+        ) {
+            BalancerTaoReservoir::<T>::insert(netuid, TaoBalance::ZERO);
+            BalancerAlphaReservoir::<T>::insert(netuid, AlphaBalance::ZERO);
+            SwapBalancer::<T>::insert(netuid, new_balancer);
+            return (pending_tao, tao_delta, pending_alpha, alpha_delta);
+        }
 
-        // Update weights and log errors if they go out of range
-        if balancer
-            .update_weights_for_added_liquidity(
-                u64::from(tao_reserve),
-                u64::from(alpha_reserve),
-                u64::from(tao_delta),
-                u64::from(alpha_delta),
-            )
-            .is_err()
-        {
+        if let Some(new_balancer) = Self::try_update_balancer(
+            &balancer,
+            tao_reserve,
+            alpha_reserve,
+            TaoBalance::ZERO,
+            pending_alpha,
+        ) {
+            BalancerTaoReservoir::<T>::insert(netuid, pending_tao);
+            BalancerAlphaReservoir::<T>::insert(netuid, AlphaBalance::ZERO);
+            SwapBalancer::<T>::insert(netuid, new_balancer);
+            return (TaoBalance::ZERO, tao_delta, pending_alpha, alpha_delta);
+        }
+
+        if let Some(new_balancer) = Self::try_update_balancer(
+            &balancer,
+            tao_reserve,
+            alpha_reserve,
+            pending_tao,
+            AlphaBalance::ZERO,
+        ) {
+            BalancerTaoReservoir::<T>::insert(netuid, TaoBalance::ZERO);
+            BalancerAlphaReservoir::<T>::insert(netuid, pending_alpha);
+            SwapBalancer::<T>::insert(netuid, new_balancer);
+            return (pending_tao, tao_delta, AlphaBalance::ZERO, alpha_delta);
+        }
+
+        BalancerTaoReservoir::<T>::insert(netuid, pending_tao);
+        BalancerAlphaReservoir::<T>::insert(netuid, pending_alpha);
+        if pending_tao > TaoBalance::ZERO || pending_alpha > AlphaBalance::ZERO {
             log::warn!(
-                "Reserves are out of range for emission: netuid = {}, tao = {}, alpha = {}, tao_delta = {}, alpha_delta = {}",
+                "Reserves are out of range for emission: netuid = {}, tao = {}, alpha = {}, tao_delta = {}, alpha_delta = {}, tao_reservoir = {}, alpha_reservoir = {}",
                 netuid,
                 tao_reserve,
                 alpha_reserve,
                 tao_delta,
-                alpha_delta
+                alpha_delta,
+                pending_tao,
+                pending_alpha
             );
-            (TaoBalance::ZERO, AlphaBalance::ZERO)
-        } else {
-            SwapBalancer::<T>::insert(netuid, balancer);
-            (tao_delta, alpha_delta)
         }
+
+        (TaoBalance::ZERO, tao_delta, AlphaBalance::ZERO, alpha_delta)
+    }
+
+    fn try_update_balancer(
+        balancer: &Balancer,
+        tao_reserve: TaoBalance,
+        alpha_reserve: AlphaBalance,
+        tao_delta: TaoBalance,
+        alpha_delta: AlphaBalance,
+    ) -> Option<Balancer> {
+        let mut new_balancer = balancer.clone();
+        new_balancer
+            .update_weights_for_added_liquidity(
+                u64::from(tao_reserve),
+                u64::from(alpha_reserve),
+                u64::from(tao_delta),
+                u64::from(alpha_delta),
+            )
+            .ok()?;
+        Some(new_balancer)
     }
 
     /// Executes a token swap on the specified subnet.
@@ -276,6 +347,8 @@ impl<T: Config> Pallet<T> {
 
         FeeRate::<T>::remove(netuid);
         SwapBalancer::<T>::remove(netuid);
+        BalancerTaoReservoir::<T>::remove(netuid);
+        BalancerAlphaReservoir::<T>::remove(netuid);
 
         log::debug!(
             "clear_protocol_liquidity: netuid={netuid:?}, protocol_burned: τ={burned_tao:?}, α={burned_alpha:?}; state cleared"
@@ -390,7 +463,7 @@ impl<T: Config> SwapHandler for Pallet<T> {
         netuid: NetUid,
         tao_delta: TaoBalance,
         alpha_delta: AlphaBalance,
-    ) -> (TaoBalance, AlphaBalance) {
+    ) -> (TaoBalance, TaoBalance, AlphaBalance, AlphaBalance) {
         Self::adjust_protocol_liquidity(netuid, tao_delta, alpha_delta)
     }
 

pallets/swap/src/pallet/mod.rs

@@ -113,6 +113,15 @@ mod pallet {
     #[pallet::storage]
     pub type PalSwapInitialized<T> = StorageMap<_, Twox64Concat, NetUid, bool, ValueQuery>;
 
+    /// TAO protocol liquidity that could not be injected without exceeding balancer weight bounds.
+    #[pallet::storage]
+    pub type BalancerTaoReservoir<T> = StorageMap<_, Twox64Concat, NetUid, TaoBalance, ValueQuery>;
+
+    /// Alpha protocol liquidity that could not be injected without exceeding balancer weight bounds.
+    #[pallet::storage]
+    pub type BalancerAlphaReservoir<T> =
+        StorageMap<_, Twox64Concat, NetUid, AlphaBalance, ValueQuery>;
+
     /// --- Storage for migration run status
     #[pallet::storage]
     pub type HasMigrationRun<T: Config> =