Custom dust collection type in subtensor

pallets/subtensor/src/coinbase/tao.rs

@@ -33,33 +33,21 @@ impl<T: Config> Pallet<T> {
         SubnetTAO::<T>::get(netuid)
     }
 
-    /// Internal function that transfers and updates subtensor pallet total issuance
-    /// in case of dust collection.
+    /// Internal function that transfers TAO and allows the origin account to be reaped.
+    ///
+    /// Dust collection is handled by the runtime's Balances `DustRemoval` implementation.
     fn transfer_allow_death_update_ti(
         origin_coldkey: &T::AccountId,
         destination_coldkey: &T::AccountId,
         amount: BalanceOf<T>,
     ) -> DispatchResult {
-        // If account balance remainder drops below ED, then account is killed, balance
-        // is lost, and we need to reduce total issuance in subtensor pallet. Measure
-        // balance TI before and after to detect the dust.
-        let balances_ti_before = <T as pallet::Config>::Currency::total_issuance();
-
         <T as pallet::Config>::Currency::transfer(
             origin_coldkey,
             destination_coldkey,
             amount,
             Preservation::Expendable,
         )?;
 
-        let balances_ti_after = <T as pallet::Config>::Currency::total_issuance();
-        if balances_ti_after < balances_ti_before {
-            let burned = balances_ti_before.saturating_sub(balances_ti_after);
-            TotalIssuance::<T>::mutate(|total| {
-                *total = total.saturating_sub(burned);
-            });
-        }
-
         Ok(())
     }
 

pallets/subtensor/src/lib.rs

@@ -69,6 +69,21 @@ pub const MAX_SUBNET_CLAIMS: usize = 5;
 
 pub const MAX_ROOT_CLAIM_THRESHOLD: u64 = 10_000_000;
 
+pub struct SubtensorDustRemoval<T>(PhantomData<T>);
+impl<T> frame_support::traits::OnUnbalanced<pallet_balances::CreditOf<T, ()>>
+    for SubtensorDustRemoval<T>
+where
+    T: Config + pallet_balances::Config,
+    <T as pallet_balances::Config>::Balance: Into<TaoBalance> + Copy,
+{
+    fn on_nonzero_unbalanced(dust: pallet_balances::CreditOf<T, ()>) {
+        let amount: TaoBalance = frame_support::traits::Imbalance::peek(&dust).into();
+        TotalIssuance::<T>::mutate(|total| {
+            *total = total.saturating_sub(amount);
+        });
+    }
+}
+
 #[allow(deprecated)]
 #[deny(missing_docs)]
 #[import_section(errors::errors)]

pallets/subtensor/src/migrations/migrate_fix_total_issuance_evm_fees.rs

@@ -3,24 +3,31 @@ use frame_support::traits::fungible::Inspect;
 use frame_support::weights::Weight;
 
 pub fn migrate_fix_total_issuance_evm_fees<T: Config>() -> Weight {
-    let migration_name = b"migrate_fix_total_issuance_evm_fees".to_vec();
-    let mut weight = T::DbWeight::get().reads(1);
-
-    if HasMigrationRun::<T>::get(&migration_name) {
-        log::info!(
-            "Migration '{:?}' has already run. Skipping.",
-            String::from_utf8_lossy(&migration_name)
-        );
+    let migration_names: [&[u8]; 2] = [
+        // Fix testnet TotalIssuance after the earlier EVM fees issue caused the
+        // Subtensor pallet's accounting to diverge from the balances pallet.
+        b"migrate_fix_total_issuance_evm_fees",
+        // Fix Subtensor TotalIssuance after dust collection caused accounting drift.
+        b"migrate_fix_total_issuance_after_dust_collection",
+    ];
+    let mut weight = T::DbWeight::get().reads(migration_names.len() as u64);
+
+    let Some(migration_name) = migration_names
+        .iter()
+        .map(|name| name.to_vec())
+        .find(|name| !HasMigrationRun::<T>::get(name))
+    else {
+        log::info!("All total issuance fix migrations have already run. Skipping.");
         return weight;
-    }
+    };
 
     log::info!(
         "Running migration '{}'",
         String::from_utf8_lossy(&migration_name)
     );
 
-    // Fix testnet TotalIssuance after the earlier EVM fees issue caused the
-    // Subtensor pallet's accounting to diverge from the balances pallet.
+    // All migration instances reset Subtensor TotalIssuance to the authoritative
+    // Balances pallet total issuance.
     let balances_total_issuance = <T as Config>::Currency::total_issuance();
     let subtensor_total_issuance_before = TotalIssuance::<T>::get();
     TotalIssuance::<T>::put(balances_total_issuance);

pallets/subtensor/src/tests/migration.rs

@@ -4362,6 +4362,7 @@ fn test_migrate_subnet_balances() {
 fn test_migrate_fix_total_issuance_evm_fees() {
     new_test_ext(1).execute_with(|| {
         const MIGRATION_NAME: &[u8] = b"migrate_fix_total_issuance_evm_fees";
+        const DUST_MIGRATION_NAME: &[u8] = b"migrate_fix_total_issuance_after_dust_collection";
 
         let account = U256::from(42);
         let balances_total_issuance = TaoBalance::from(123_456_789_u64);
@@ -4382,16 +4383,29 @@ fn test_migrate_fix_total_issuance_evm_fees() {
         assert!(!weight.is_zero(), "weight must be non-zero");
         assert_eq!(TotalIssuance::<Test>::get(), balances_total_issuance);
         assert!(HasMigrationRun::<Test>::get(MIGRATION_NAME.to_vec()));
+        assert!(!HasMigrationRun::<Test>::get(
+            DUST_MIGRATION_NAME.to_vec()
+        ));
 
         let second_wrong_value = TaoBalance::from(555_u64);
         TotalIssuance::<Test>::put(second_wrong_value);
 
         crate::migrations::migrate_fix_total_issuance_evm_fees::migrate_fix_total_issuance_evm_fees::<Test>();
 
+        assert_eq!(TotalIssuance::<Test>::get(), balances_total_issuance);
+        assert!(HasMigrationRun::<Test>::get(
+            DUST_MIGRATION_NAME.to_vec()
+        ));
+
+        let third_wrong_value = TaoBalance::from(777_u64);
+        TotalIssuance::<Test>::put(third_wrong_value);
+
+        crate::migrations::migrate_fix_total_issuance_evm_fees::migrate_fix_total_issuance_evm_fees::<Test>();
+
         assert_eq!(
             TotalIssuance::<Test>::get(),
-            second_wrong_value,
-            "migration must not run more than once"
+            third_wrong_value,
+            "migration must not run after all known migration keys have run"
         );
     });
 }

pallets/subtensor/src/tests/mock.rs

@@ -88,7 +88,7 @@ pub type BlockNumber = u64;
 impl pallet_balances::Config for Test {
     type Balance = Balance;
     type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
+    type DustRemoval = crate::SubtensorDustRemoval<Test>;
     type ExistentialDeposit = ExistentialDeposit;
     type AccountStore = System;
     type MaxLocks = ();

pallets/subtensor/src/tests/mock_high_ed.rs

@@ -70,7 +70,7 @@ pub type BlockNumber = u64;
 impl pallet_balances::Config for Test {
     type Balance = Balance;
     type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
+    type DustRemoval = crate::SubtensorDustRemoval<Test>;
     type ExistentialDeposit = ExistentialDeposit;
     type AccountStore = System;
     type MaxLocks = ();

runtime/src/lib.rs

@@ -519,7 +519,7 @@ impl pallet_balances::Config for Runtime {
     type Balance = Balance;
     // The ubiquitous event type.
     type RuntimeEvent = RuntimeEvent;
-    type DustRemoval = ();
+    type DustRemoval = pallet_subtensor::SubtensorDustRemoval<Runtime>;
     type ExistentialDeposit = ExistentialDeposit;
     type AccountStore = System;
     type WeightInfo = pallet_balances::weights::SubstrateWeight<Runtime>;

runtime/tests/balances_dust.rs

@@ -0,0 +1,60 @@
+#![allow(clippy::unwrap_used)]
+
+use frame_support::traits::fungible::{Inspect, Mutate};
+use frame_support::traits::tokens::Preservation;
+use node_subtensor_runtime::{Balances, BuildStorage, Runtime, RuntimeGenesisConfig};
+use sp_core::crypto::AccountId32;
+use subtensor_runtime_common::TaoBalance;
+
+fn new_test_ext() -> sp_io::TestExternalities {
+    let mut ext: sp_io::TestExternalities = RuntimeGenesisConfig {
+        ..Default::default()
+    }
+    .build_storage()
+    .unwrap()
+    .into();
+    ext.execute_with(|| frame_system::Pallet::<Runtime>::set_block_number(1));
+    ext
+}
+
+fn add_balance_to_account(account: &AccountId32, tao: TaoBalance) {
+    let credit = pallet_subtensor::Pallet::<Runtime>::mint_tao(tao);
+    let _ = pallet_subtensor::Pallet::<Runtime>::spend_tao(account, credit, tao).unwrap();
+}
+
+#[test]
+fn balances_dust_removal_updates_subtensor_total_issuance() {
+    new_test_ext().execute_with(|| {
+        let origin = AccountId32::new([1u8; 32]);
+        let destination = AccountId32::new([2u8; 32]);
+        let existential_deposit = TaoBalance::from(500u64);
+        let dust = TaoBalance::from(1u64);
+        let transfer_amount = existential_deposit;
+
+        add_balance_to_account(&origin, transfer_amount + dust);
+
+        let balances_issuance_before = Balances::total_issuance();
+        let subtensor_issuance_before = pallet_subtensor::Pallet::<Runtime>::get_total_issuance();
+        assert_eq!(balances_issuance_before, subtensor_issuance_before);
+
+        <Balances as Mutate<AccountId32>>::transfer(
+            &origin,
+            &destination,
+            transfer_amount,
+            Preservation::Expendable,
+        )
+        .unwrap();
+
+        assert_eq!(Balances::total_balance(&origin), 0u64.into());
+        assert_eq!(Balances::total_balance(&destination), transfer_amount);
+        assert_eq!(balances_issuance_before - Balances::total_issuance(), dust);
+        assert_eq!(
+            subtensor_issuance_before - pallet_subtensor::Pallet::<Runtime>::get_total_issuance(),
+            dust
+        );
+        assert_eq!(
+            Balances::total_issuance(),
+            pallet_subtensor::Pallet::<Runtime>::get_total_issuance()
+        );
+    });
+}