feat/s2-deregistrations (#204)

* feat: move storage DefaultMinValidatorStake to pallet Config (#201)

* feat: add validator blacklist (#202)

* feat: add validator blacklist

* feat: remove  check

* big wip

* adding missing file

* feat: fixing linear consensus

* feat: refactoring deregistration logic

* fix: storage version

* feat: add weights to yuma/linear

* feat: updating benchmars

* feat: updating non-running migration

---------

Co-authored-by: João Victor <[email protected]>
Co-authored-by: devwckd <[email protected]>

pallets/governance/src/benchmarking.rs

@@ -271,7 +271,7 @@ benchmarks! {
         let caller: T::AccountId = account("Alice", 0, 1);
         let application_key: T::AccountId = account("Bob", 0, 2);
         Curator::<T>::set(caller.clone());
-    }: add_to_whitelist(RawOrigin::Signed(caller), application_key, 1)
+    }: add_to_whitelist(RawOrigin::Signed(caller), application_key)
 
     // 12
     remove_from_whitelist {
@@ -282,7 +282,7 @@ benchmarks! {
         Curator::<T>::set(caller.clone());
         // Now add it to whitelist
         GovernanceMod::<T>::add_to_whitelist(RawOrigin::Signed(caller.clone()).into(),
-    application_key.clone(), 1)?; }: remove_from_whitelist(RawOrigin::Signed(caller),
+    application_key.clone())?; }: remove_from_whitelist(RawOrigin::Signed(caller),
     application_key)
 
 }

pallets/subnet_emission/src/benchmarking.rs

@@ -0,0 +1,50 @@
+#![cfg(feature = "runtime-benchmarks")]
+
+use crate::{Pallet as SubnetEmissionMod, *};
+use frame_benchmarking::{account, benchmarks};
+use frame_system::RawOrigin;
+pub use pallet::*;
+use pallet_subspace::{Pallet as SubspaceMod, SubnetBurn};
+use sp_std::vec::Vec;
+
+fn register_mock<T: Config>(
+    key: T::AccountId,
+    module_key: T::AccountId,
+    name: Vec<u8>,
+) -> Result<(), &'static str> {
+    let address = "test".as_bytes().to_vec();
+    let network = "testnet".as_bytes().to_vec();
+
+    let enough_stake = 10000000000000u64;
+    SubspaceMod::<T>::add_balance_to_account(
+        &key,
+        SubspaceMod::<T>::u64_to_balance(SubnetBurn::<T>::get() + enough_stake).unwrap(),
+    );
+    let metadata = Some("metadata".as_bytes().to_vec());
+    SubspaceMod::<T>::register(
+        RawOrigin::Signed(key.clone()).into(),
+        network,
+        name,
+        address,
+        module_key.clone(),
+        metadata,
+    )?;
+    SubspaceMod::<T>::increase_stake(&key, &module_key, enough_stake);
+    Ok(())
+}
+
+benchmarks! {
+    process_subnets {
+        let caller: T::AccountId = account("Alice", 0, 1);
+        // Add Alice's funds to submit the proposal
+        SubspaceMod::<T>::add_balance_to_account(
+            &caller,
+            SubspaceMod::<T>::u64_to_balance(1_000_000_000_000_000).unwrap()
+        );
+
+        register_mock::<T>(caller.clone(), caller.clone(),
+    "test".as_bytes().to_vec())?;
+
+        reigster_mock()
+    }: process_subnets()
+}

pallets/subnet_emission/src/distribute_emission.rs

@@ -1,7 +1,7 @@
 use super::*;
 use crate::subnet_consensus::{linear::LinearEpoch, treasury::TreasuryEpoch, yuma::YumaEpoch};
 
-use frame_support::storage::with_storage_layer;
+use frame_support::{storage::with_storage_layer, traits::Get, weights::Weight};
 use pallet_subnet_emission_api::SubnetConsensus;
 use pallet_subspace::N;
 
@@ -14,19 +14,25 @@ use pallet_subspace::N;
 ///
 /// This function iterates through all subnets, updates their pending emissions,
 /// and runs an epoch if it's time for that subnet.
-fn process_subnets<T: Config>(block_number: u64, subnets_emission_distribution: PricedSubnets) {
-    for netuid in N::<T>::iter_keys() {
+fn process_subnets<T: Config>(
+    block_number: u64,
+    subnets_emission_distribution: PricedSubnets,
+) -> Weight {
+    let total_weight = N::<T>::iter_keys().fold(Weight::zero(), |acc_weight, netuid| {
         update_pending_emission::<T>(
             netuid,
             subnets_emission_distribution.get(&netuid).unwrap_or(&0),
         );
+        let mut weight = acc_weight.saturating_add(T::DbWeight::get().writes(1));
 
-        if pallet_subspace::Pallet::<T>::blocks_until_next_epoch(netuid, block_number) > 0 {
-            continue;
+        if pallet_subspace::Pallet::<T>::blocks_until_next_epoch(netuid, block_number) == 0 {
+            weight = weight.saturating_add(run_epoch::<T>(netuid));
         }
 
-        run_epoch::<T>(netuid);
-    }
+        weight
+    });
+
+    total_weight
 }
 
 /// Updates the pending emission for a given subnet.
@@ -55,21 +61,30 @@ fn update_pending_emission<T: Config>(netuid: u16, new_queued_emission: &u64) {
 /// This function clears the set weight rate limiter, retrieves the pending emission,
 /// and if there's emission to distribute, runs the consensus algorithm. If successful,
 /// it finalizes the epoch. If an error occurs during consensus, it logs the error
-fn run_epoch<T: Config>(netuid: u16) {
+fn run_epoch<T: Config>(netuid: u16) -> Weight {
     log::trace!("running epoch for subnet {netuid}");
 
+    let mut weight = T::DbWeight::get().reads(1);
+
     let emission_to_drain = PendingEmission::<T>::get(netuid);
     if emission_to_drain > 0 {
         match run_consensus_algorithm::<T>(netuid, emission_to_drain) {
-            Ok(()) => finalize_epoch::<T>(netuid),
+            Ok(consensus_weight) => {
+                weight = weight.saturating_add(consensus_weight);
+                finalize_epoch::<T>(netuid);
+                weight
+            }
             Err(e) => {
                 log::error!(
                     "Error running consensus algorithm for subnet {}: {:?}",
                     netuid,
                     e
                 );
+                Weight::zero()
             }
         }
+    } else {
+        weight
     }
 }
 
@@ -93,14 +108,14 @@ fn run_epoch<T: Config>(netuid: u16) {
 fn run_consensus_algorithm<T: Config>(
     netuid: u16,
     emission_to_drain: u64,
-) -> Result<(), &'static str> {
+) -> Result<Weight, &'static str> {
     with_storage_layer(|| {
         let Some(consensus_type) = SubnetConsensusType::<T>::get(netuid) else {
-            return Ok(());
+            return Ok(T::DbWeight::get().reads(1));
         };
 
         match consensus_type {
-            SubnetConsensus::Root => Ok(()),
+            SubnetConsensus::Root => Ok(T::DbWeight::get().reads(1)),
             SubnetConsensus::Treasury => run_treasury_consensus::<T>(netuid, emission_to_drain),
             SubnetConsensus::Linear => run_linear_consensus::<T>(netuid, emission_to_drain),
             SubnetConsensus::Yuma => run_yuma_consensus::<T>(netuid, emission_to_drain),
@@ -122,14 +137,14 @@ fn run_consensus_algorithm<T: Config>(
 fn run_linear_consensus<T: Config>(
     netuid: u16,
     emission_to_drain: u64,
-) -> Result<(), &'static str> {
+) -> Result<Weight, &'static str> {
     LinearEpoch::<T>::new(netuid, emission_to_drain)
         .run()
-        .map(|_| ())
+        .map(|(_, weight)| weight)
         .map_err(|err| {
             log::error!(
                 "Failed to run linear consensus algorithm: {err:?}, skipping this block. \
-            {emission_to_drain} tokens will be emitted on the next epoch."
+                {emission_to_drain} tokens will be emitted on the next epoch."
             );
             "linear failed"
         })
@@ -147,14 +162,20 @@ fn run_linear_consensus<T: Config>(
 /// A Result indicating success or failure of the Yuma consensus algorithm.
 ///
 /// This function creates and runs a new YumaEpoch, logging any errors that occur.
-fn run_yuma_consensus<T: Config>(netuid: u16, emission_to_drain: u64) -> Result<(), &'static str> {
-    YumaEpoch::<T>::new(netuid, emission_to_drain).run().map(|_| ()).map_err(|err| {
-        log::error!(
-            "Failed to run yuma consensus algorithm: {err:?}, skipping this block. \
+fn run_yuma_consensus<T: Config>(
+    netuid: u16,
+    emission_to_drain: u64,
+) -> Result<Weight, &'static str> {
+    YumaEpoch::<T>::new(netuid, emission_to_drain)
+        .run()
+        .map(|(_, weight)| weight)
+        .map_err(|err| {
+            log::error!(
+                "Failed to run yuma consensus algorithm: {err:?}, skipping this block. \
             {emission_to_drain} tokens will be emitted on the next epoch."
-        );
-        "yuma failed"
-    })
+            );
+            "yuma failed"
+        })
 }
 
 /// Runs the treasury consensus algorithm for a given network and emission amount.
@@ -171,14 +192,14 @@ fn run_yuma_consensus<T: Config>(netuid: u16, emission_to_drain: u64) -> Result<
 fn run_treasury_consensus<T: Config>(
     netuid: u16,
     emission_to_drain: u64,
-) -> Result<(), &'static str> {
+) -> Result<Weight, &'static str> {
     TreasuryEpoch::<T>::new(netuid, emission_to_drain)
         .run()
-        .map(|_| ())
+        .map(|_| T::DbWeight::get().reads_writes(1, 1))
         .map_err(|err| {
             log::error!(
                 "Failed to run treasury consensus algorithm: {err:?}, skipping this block. \
-            {emission_to_drain} tokens will be emitted on the next epoch."
+                {emission_to_drain} tokens will be emitted on the next epoch."
             );
             "treasury failed"
         })
@@ -213,11 +234,11 @@ impl<T: Config> Pallet<T> {
     ///
     /// This function calculates the emission distribution across subnets and
     /// processes each subnet accordingly.
-    pub fn process_emission_distribution(block_number: u64, emission_per_block: u64) {
+    pub fn process_emission_distribution(block_number: u64, emission_per_block: u64) -> Weight {
         log::debug!("stepping block {block_number:?}");
 
         let subnets_emission_distribution = Self::get_subnet_pricing(emission_per_block);
-        process_subnets::<T>(block_number, subnets_emission_distribution);
+        process_subnets::<T>(block_number, subnets_emission_distribution)
     }
 
     // ---------------------------------

pallets/subnet_emission/src/lib.rs

@@ -10,16 +10,13 @@ use sp_std::collections::btree_map::BTreeMap;
 
 pub mod distribute_emission;
 pub mod migrations;
+pub mod post_runtime;
 pub mod subnet_pricing {
     pub mod demo;
     pub mod root;
 }
 
-pub mod subnet_consensus {
-    pub mod linear;
-    pub mod treasury;
-    pub mod yuma;
-}
+pub mod subnet_consensus;
 
 // TODO:
 // move some import outside of the macro
@@ -112,6 +109,11 @@ pub mod pallet {
             }
             Weight::zero()
         }
+
+        fn on_idle(_n: BlockNumberFor<T>, remaining: Weight) -> Weight {
+            log::info!("on_idle: {remaining:?}");
+            Self::deregister_excess_modules(remaining)
+        }
     }
 
     #[pallet::event]

pallets/subnet_emission/src/post_runtime.rs

@@ -0,0 +1,72 @@
+use super::*;
+use frame_support::{storage::with_storage_layer, traits::Get, weights::Weight};
+use pallet_subnet_emission_api::SubnetConsensus;
+use pallet_subspace::{MaxAllowedUids, Pallet as PalletS, N};
+
+impl<T: Config> Pallet<T> {
+    pub(crate) fn deregister_excess_modules(mut remaining: Weight) -> Weight {
+        let netuid = Self::get_consensus_netuid(SubnetConsensus::Linear).unwrap_or(2);
+        const MAX_MODULES_PER_ITERATION: usize = 5;
+        const MAX_UIDS: u16 = 524;
+
+        let db_weight = T::DbWeight::get();
+        let mut weight = db_weight.reads(2);
+        let find_id_weight = db_weight.reads(1);
+        let deregister_weight = Weight::from_parts(300_495_000, 21587)
+            .saturating_add(T::DbWeight::get().reads(34_u64))
+            .saturating_add(T::DbWeight::get().writes(48_u64));
+
+        if !remaining
+            .all_gte(weight.saturating_add(find_id_weight).saturating_add(deregister_weight))
+        {
+            log::info!("not enough weight remaining: {remaining:?}");
+            return Weight::zero();
+        }
+
+        remaining = remaining.saturating_sub(weight);
+
+        let mut module_count = N::<T>::get(netuid);
+        while module_count > MAX_UIDS {
+            for _ in 0..MAX_MODULES_PER_ITERATION {
+                if !remaining.all_gte(find_id_weight.saturating_add(deregister_weight)) {
+                    log::info!("not enough weight remaining: {remaining:?}");
+                    return weight;
+                }
+
+                if let Some(uid) = PalletS::<T>::get_lowest_uid(netuid, false) {
+                    log::info!("deregistering module with uid {uid}");
+
+                    weight = weight.saturating_add(find_id_weight);
+                    remaining = remaining.saturating_sub(find_id_weight);
+
+                    let result =
+                        with_storage_layer(|| PalletS::<T>::remove_module(netuid, uid, true));
+                    if result.is_ok() {
+                        weight = weight.saturating_add(deregister_weight);
+                        remaining = remaining.saturating_sub(deregister_weight);
+                        module_count = module_count.saturating_sub(1);
+                    } else {
+                        log::error!(
+                            "failed to deregister module {uid} due to: {:?}",
+                            result.unwrap_err()
+                        );
+                    }
+                } else {
+                    // No more modules to deregister
+                    break;
+                }
+
+                if module_count <= MAX_UIDS {
+                    break;
+                }
+            }
+
+            if module_count <= MAX_UIDS {
+                break;
+            }
+        }
+
+        MaxAllowedUids::<T>::set(netuid, MAX_UIDS);
+        weight
+    }
+}