address review feedback

moondance-labs · Nov 21, 2024 · 3152372 · 3152372
1 parent b47f66e
commit 3152372
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 67 deletions.
diff --git a/pallets/collator-assignment/src/benchmarking.rs b/pallets/collator-assignment/src/benchmarking.rs
@@ -68,7 +68,7 @@ mod benchmarks {
         let container_chains: Vec<_> = (0..y).map(|i| ParaId::from(2000 + i)).collect();
         let session_index = 0u32.into();
         T::ContainerChains::set_session_container_chains(session_index, &container_chains);
-        T::RemoveParaIdsWithNoCredits::make_valid_para_ids(&container_chains);
+        T::ParaIdAssignmentHooks::make_valid_para_ids(&container_chains);
         T::HostConfiguration::set_host_configuration(session_index);
 
         // Assign random collators to test worst case: when collators need to be checked against existing collators

diff --git a/pallets/collator-assignment/src/lib.rs b/pallets/collator-assignment/src/lib.rs
@@ -55,7 +55,7 @@ use {
     sp_std::{collections::btree_set::BTreeSet, fmt::Debug, prelude::*, vec},
     tp_traits::{
         CollatorAssignmentTip, GetContainerChainAuthor, GetHostConfiguration,
-        GetSessionContainerChains, ParaId, RemoveInvulnerables, RemoveParaIdsWithNoCredits,
+        GetSessionContainerChains, ParaId, ParaIdAssignmentHooks, RemoveInvulnerables,
         ShouldRotateAllCollators, Slot,
     },
 };
@@ -105,10 +105,7 @@ pub mod pallet {
         type ShouldRotateAllCollators: ShouldRotateAllCollators<Self::SessionIndex>;
         type GetRandomnessForNextBlock: GetRandomnessForNextBlock<BlockNumberFor<Self>>;
         type RemoveInvulnerables: RemoveInvulnerables<Self::AccountId>;
-        type RemoveParaIdsWithNoCredits: RemoveParaIdsWithNoCredits<
-            BalanceOf<Self>,
-            Self::AccountId,
-        >;
+        type ParaIdAssignmentHooks: ParaIdAssignmentHooks<BalanceOf<Self>, Self::AccountId>;
         type Currency: Currency<Self::AccountId>;
         type CollatorAssignmentTip: CollatorAssignmentTip<BalanceOf<Self>>;
         type ForceEmptyOrchestrator: Get<bool>;
@@ -311,16 +308,13 @@ pub mod pallet {
                 old_assigned.container_chains.keys().cloned().collect();
 
             // Remove the containerChains that do not have enough credits for block production
-            T::RemoveParaIdsWithNoCredits::pre_assignment_remove_para_ids_with_no_credits(
+            T::ParaIdAssignmentHooks::pre_assignment(
                 &mut container_chain_ids,
                 &old_assigned_para_ids,
             );
             // TODO: parathreads should be treated a bit differently, they don't need to have the same amount of credits
             // as parathreads because they will not be producing blocks on every slot.
-            T::RemoveParaIdsWithNoCredits::pre_assignment_remove_para_ids_with_no_credits(
-                &mut parathreads,
-                &old_assigned_para_ids,
-            );
+            T::ParaIdAssignmentHooks::pre_assignment(&mut parathreads, &old_assigned_para_ids);
 
             let mut shuffle_collators = None;
             // If the random_seed is all zeros, we don't shuffle the list of collators nor the list
@@ -469,7 +463,7 @@ pub mod pallet {
 
             // TODO: this probably is asking for a refactor
             // only apply the onCollatorAssignedHook if sufficient collators
-            T::RemoveParaIdsWithNoCredits::post_assignment_remove_para_ids_with_no_credits(
+            T::ParaIdAssignmentHooks::post_assignment(
                 &old_assigned_para_ids,
                 &mut new_assigned.container_chains,
                 &maybe_tip,

diff --git a/pallets/collator-assignment/src/mock.rs b/pallets/collator-assignment/src/mock.rs
@@ -33,8 +33,8 @@ use {
     },
     sp_std::collections::{btree_map::BTreeMap, btree_set::BTreeSet},
     tp_traits::{
-        CollatorAssignmentTip, ParaId, ParathreadParams, RemoveInvulnerables,
-        RemoveParaIdsWithNoCredits, SessionContainerChains,
+        CollatorAssignmentTip, ParaId, ParaIdAssignmentHooks, ParathreadParams,
+        RemoveInvulnerables, SessionContainerChains,
     },
     tracing_subscriber::{layer::SubscriberExt, FmtSubscriber},
 };
@@ -335,7 +335,7 @@ impl pallet_collator_assignment::Config for Test {
         RotateCollatorsEveryNSessions<MockCollatorRotationSessionPeriod>;
     type GetRandomnessForNextBlock = MockGetRandomnessForNextBlock;
     type RemoveInvulnerables = RemoveAccountIdsAbove100;
-    type RemoveParaIdsWithNoCredits = RemoveParaIdsAbove5000;
+    type ParaIdAssignmentHooks = MockParaIdAssignmentHooksImpl;
     type CollatorAssignmentTip = MockCollatorAssignmentTip;
     type ForceEmptyOrchestrator = ConstBool<false>;
     type Currency = ();
@@ -396,17 +396,14 @@ impl RemoveInvulnerables<u64> for RemoveAccountIdsAbove100 {
 }
 
 /// Any ParaId >= 5000 will be considered to not have enough credits
-pub struct RemoveParaIdsAbove5000;
+pub struct MockParaIdAssignmentHooksImpl;
 
-impl<AC> RemoveParaIdsWithNoCredits<u32, AC> for RemoveParaIdsAbove5000 {
-    fn pre_assignment_remove_para_ids_with_no_credits(
-        para_ids: &mut Vec<ParaId>,
-        _old_assigned: &BTreeSet<ParaId>,
-    ) {
+impl<AC> ParaIdAssignmentHooks<u32, AC> for MockParaIdAssignmentHooksImpl {
+    fn pre_assignment(para_ids: &mut Vec<ParaId>, _old_assigned: &BTreeSet<ParaId>) {
         para_ids.retain(|para_id| *para_id <= ParaId::from(5000));
     }
 
-    fn post_assignment_remove_para_ids_with_no_credits(
+    fn post_assignment(
         _current_assigned: &BTreeSet<ParaId>,
         new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
         _maybe_tip: &Option<u32>,

diff --git a/primitives/traits/src/lib.rs b/primitives/traits/src/lib.rs
@@ -272,14 +272,11 @@ impl<AccountId: Clone> RemoveInvulnerables<AccountId> for () {
 
 /// Helper trait for pallet_collator_assignment to be able to not assign collators to container chains with no credits
 /// in pallet_services_payment
-pub trait RemoveParaIdsWithNoCredits<B, AC> {
+pub trait ParaIdAssignmentHooks<B, AC> {
     /// Remove para ids with not enough credits. The resulting order will affect priority: the first para id in the list
     /// will be the first one to get collators.
-    fn pre_assignment_remove_para_ids_with_no_credits(
-        para_ids: &mut Vec<ParaId>,
-        old_assigned: &BTreeSet<ParaId>,
-    );
-    fn post_assignment_remove_para_ids_with_no_credits(
+    fn pre_assignment(para_ids: &mut Vec<ParaId>, old_assigned: &BTreeSet<ParaId>);
+    fn post_assignment(
         current_assigned: &BTreeSet<ParaId>,
         new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
         maybe_tip: &Option<B>,
@@ -290,14 +287,10 @@ pub trait RemoveParaIdsWithNoCredits<B, AC> {
     fn make_valid_para_ids(para_ids: &[ParaId]);
 }
 
-impl<B, AC> RemoveParaIdsWithNoCredits<B, AC> for () {
-    fn pre_assignment_remove_para_ids_with_no_credits(
-        _para_ids: &mut Vec<ParaId>,
-        _currently_assigned: &BTreeSet<ParaId>,
-    ) {
-    }
+impl<B, AC> ParaIdAssignmentHooks<B, AC> for () {
+    fn pre_assignment(_para_ids: &mut Vec<ParaId>, _currently_assigned: &BTreeSet<ParaId>) {}
 
-    fn post_assignment_remove_para_ids_with_no_credits(
+    fn post_assignment(
         _current_assigned: &BTreeSet<ParaId>,
         _new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
         _maybe_tip: &Option<B>,

diff --git a/runtime/dancebox/src/lib.rs b/runtime/dancebox/src/lib.rs
@@ -115,8 +115,8 @@ use {
     },
     tp_traits::{
         apply, derive_storage_traits, GetContainerChainAuthor, GetHostConfiguration,
-        GetSessionContainerChains, MaybeSelfChainBlockAuthor, RelayStorageRootProvider,
-        RemoveInvulnerables, RemoveParaIdsWithNoCredits, SlotFrequency,
+        GetSessionContainerChains, MaybeSelfChainBlockAuthor, ParaIdAssignmentHooks,
+        RelayStorageRootProvider, RemoveInvulnerables, SlotFrequency,
     },
     tp_xcm_core_buyer::BuyCoreCollatorProof,
     xcm_runtime_apis::{
@@ -812,9 +812,9 @@ impl RemoveInvulnerables<CollatorId> for RemoveInvulnerablesImpl {
     }
 }
 
-pub struct RemoveParaIdsWithNoCreditsImpl;
+pub struct ParaIdAssignmentHooksImpl;
 
-impl RemoveParaIdsWithNoCreditsImpl {
+impl ParaIdAssignmentHooksImpl {
     fn charge_para_ids_internal(
         blocks_per_session: tp_traits::BlockNumber,
         para_id: ParaId,
@@ -888,11 +888,8 @@ impl RemoveParaIdsWithNoCreditsImpl {
     }
 }
 
-impl<AC> RemoveParaIdsWithNoCredits<BalanceOf<Runtime>, AC> for RemoveParaIdsWithNoCreditsImpl {
-    fn pre_assignment_remove_para_ids_with_no_credits(
-        para_ids: &mut Vec<ParaId>,
-        currently_assigned: &BTreeSet<ParaId>,
-    ) {
+impl<AC> ParaIdAssignmentHooks<BalanceOf<Runtime>, AC> for ParaIdAssignmentHooksImpl {
+    fn pre_assignment(para_ids: &mut Vec<ParaId>, currently_assigned: &BTreeSet<ParaId>) {
         let blocks_per_session = Period::get();
         para_ids.retain(|para_id| {
             with_transaction(|| {
@@ -909,7 +906,7 @@ impl<AC> RemoveParaIdsWithNoCredits<BalanceOf<Runtime>, AC> for RemoveParaIdsWit
         });
     }
 
-    fn post_assignment_remove_para_ids_with_no_credits(
+    fn post_assignment(
         current_assigned: &BTreeSet<ParaId>,
         new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
         maybe_tip: &Option<BalanceOf<Runtime>>,
@@ -972,7 +969,7 @@ impl pallet_collator_assignment::Config for Runtime {
         RotateCollatorsEveryNSessions<ConfigurationCollatorRotationSessionPeriod>;
     type GetRandomnessForNextBlock = BabeGetRandomnessForNextBlock;
     type RemoveInvulnerables = RemoveInvulnerablesImpl;
-    type RemoveParaIdsWithNoCredits = RemoveParaIdsWithNoCreditsImpl;
+    type ParaIdAssignmentHooks = ParaIdAssignmentHooksImpl;
     type CollatorAssignmentTip = ServicesPayment;
     type Currency = Balances;
     type ForceEmptyOrchestrator = ConstBool<false>;

diff --git a/runtime/flashbox/src/lib.rs b/runtime/flashbox/src/lib.rs
@@ -103,8 +103,8 @@ use {
     sp_version::RuntimeVersion,
     tp_traits::{
         apply, derive_storage_traits, GetContainerChainAuthor, GetHostConfiguration,
-        GetSessionContainerChains, MaybeSelfChainBlockAuthor, RelayStorageRootProvider,
-        RemoveInvulnerables, RemoveParaIdsWithNoCredits, ShouldRotateAllCollators,
+        GetSessionContainerChains, MaybeSelfChainBlockAuthor, ParaIdAssignmentHooks,
+        RelayStorageRootProvider, RemoveInvulnerables, ShouldRotateAllCollators,
     },
 };
 pub use {
@@ -653,9 +653,9 @@ impl RemoveInvulnerables<CollatorId> for RemoveInvulnerablesImpl {
     }
 }
 
-pub struct RemoveParaIdsWithNoCreditsImpl;
+pub struct ParaIdAssignmentHooksImpl;
 
-impl RemoveParaIdsWithNoCreditsImpl {
+impl ParaIdAssignmentHooksImpl {
     fn charge_para_ids_internal(
         blocks_per_session: tp_traits::BlockNumber,
         para_id: ParaId,
@@ -729,11 +729,8 @@ impl RemoveParaIdsWithNoCreditsImpl {
     }
 }
 
-impl<AC> RemoveParaIdsWithNoCredits<BalanceOf<Runtime>, AC> for RemoveParaIdsWithNoCreditsImpl {
-    fn pre_assignment_remove_para_ids_with_no_credits(
-        para_ids: &mut Vec<ParaId>,
-        currently_assigned: &BTreeSet<ParaId>,
-    ) {
+impl<AC> ParaIdAssignmentHooks<BalanceOf<Runtime>, AC> for ParaIdAssignmentHooksImpl {
+    fn pre_assignment(para_ids: &mut Vec<ParaId>, currently_assigned: &BTreeSet<ParaId>) {
         let blocks_per_session = Period::get();
         para_ids.retain(|para_id| {
             with_transaction(|| {
@@ -750,7 +747,7 @@ impl<AC> RemoveParaIdsWithNoCredits<BalanceOf<Runtime>, AC> for RemoveParaIdsWit
         });
     }
 
-    fn post_assignment_remove_para_ids_with_no_credits(
+    fn post_assignment(
         current_assigned: &BTreeSet<ParaId>,
         new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
         maybe_tip: &Option<BalanceOf<Runtime>>,
@@ -820,7 +817,7 @@ impl pallet_collator_assignment::Config for Runtime {
     type ShouldRotateAllCollators = NeverRotateCollators;
     type GetRandomnessForNextBlock = ();
     type RemoveInvulnerables = RemoveInvulnerablesImpl;
-    type RemoveParaIdsWithNoCredits = RemoveParaIdsWithNoCreditsImpl;
+    type ParaIdAssignmentHooks = ParaIdAssignmentHooksImpl;
     type CollatorAssignmentTip = ServicesPayment;
     type Currency = Balances;
     type ForceEmptyOrchestrator = ConstBool<false>;

diff --git a/solo-chains/runtime/dancelight/src/lib.rs b/solo-chains/runtime/dancelight/src/lib.rs
@@ -92,7 +92,7 @@ use {
     },
     tp_traits::{
         apply, derive_storage_traits, EraIndex, GetHostConfiguration, GetSessionContainerChains,
-        RegistrarHandler, RemoveParaIdsWithNoCredits, Slot, SlotFrequency,
+        ParaIdAssignmentHooks, RegistrarHandler, Slot, SlotFrequency,
     },
 };
 
@@ -522,6 +522,8 @@ parameter_types! {
 #[cfg(feature = "runtime-benchmarks")]
 pub struct TreasuryBenchmarkHelper<T>(PhantomData<T>);
 
+#[cfg(feature = "runtime-benchmarks")]
+use frame_support::traits::Currency;
 use frame_support::traits::{ExistenceRequirement, OnUnbalanced, WithdrawReasons};
 use pallet_services_payment::BalanceOf;
 #[cfg(feature = "runtime-benchmarks")]
@@ -3015,9 +3017,9 @@ impl frame_support::traits::Randomness<Hash, BlockNumber> for BabeCurrentBlockRa
     }
 }
 
-pub struct RemoveParaIdsWithNoCreditsImpl;
+pub struct ParaIdAssignmentHooksImpl;
 
-impl RemoveParaIdsWithNoCreditsImpl {
+impl ParaIdAssignmentHooksImpl {
     fn charge_para_ids_internal(
         blocks_per_session: BlockNumber,
         para_id: ParaId,
@@ -3091,11 +3093,8 @@ impl RemoveParaIdsWithNoCreditsImpl {
     }
 }
 
-impl<AC> RemoveParaIdsWithNoCredits<BalanceOf<Runtime>, AC> for RemoveParaIdsWithNoCreditsImpl {
-    fn pre_assignment_remove_para_ids_with_no_credits(
-        para_ids: &mut Vec<ParaId>,
-        currently_assigned: &BTreeSet<ParaId>,
-    ) {
+impl<AC> ParaIdAssignmentHooks<BalanceOf<Runtime>, AC> for ParaIdAssignmentHooksImpl {
+    fn pre_assignment(para_ids: &mut Vec<ParaId>, currently_assigned: &BTreeSet<ParaId>) {
         let blocks_per_session = EpochDurationInBlocks::get();
         para_ids.retain(|para_id| {
             with_transaction(|| {
@@ -3112,7 +3111,7 @@ impl<AC> RemoveParaIdsWithNoCredits<BalanceOf<Runtime>, AC> for RemoveParaIdsWit
         });
     }
 
-    fn post_assignment_remove_para_ids_with_no_credits(
+    fn post_assignment(
         current_assigned: &BTreeSet<ParaId>,
         new_assigned: &mut BTreeMap<ParaId, Vec<AC>>,
         maybe_tip: &Option<BalanceOf<Runtime>>,
@@ -3226,7 +3225,7 @@ impl pallet_collator_assignment::Config for Runtime {
         RotateCollatorsEveryNSessions<ConfigurationCollatorRotationSessionPeriod>;
     type GetRandomnessForNextBlock = BabeGetRandomnessForNextBlock;
     type RemoveInvulnerables = ();
-    type RemoveParaIdsWithNoCredits = RemoveParaIdsWithNoCreditsImpl;
+    type ParaIdAssignmentHooks = ParaIdAssignmentHooksImpl;
     type CollatorAssignmentTip = ServicesPayment;
     type Currency = Balances;
     type ForceEmptyOrchestrator = ConstBool<true>;