Give priority to invulnerables as orchestrator collators

pallets/collator-assignment/src/lib.rs

@@ -55,8 +55,7 @@ use {
     sp_std::{fmt::Debug, prelude::*, vec},
     tp_collator_assignment::AssignedCollators,
     tp_traits::{
-        GetContainerChainAuthor, GetHostConfiguration, GetSessionContainerChains, ParaId,
-        ShouldRotateAllCollators, Slot,
+        GetContainerChainAuthor, GetHostConfiguration, GetSessionContainerChains, ParaId, Slot, RemoveInvulnerables, ShouldRotateAllCollators
     },
 };
 
@@ -95,6 +94,7 @@ pub mod pallet {
         type ContainerChains: GetSessionContainerChains<Self::SessionIndex>;
         type ShouldRotateAllCollators: ShouldRotateAllCollators<Self::SessionIndex>;
         type GetRandomnessForNextBlock: GetRandomnessForNextBlock<BlockNumberFor<Self>>;
+        type RemoveInvulnerables: RemoveInvulnerables<Self::AccountId>;
         /// The weight information of this pallet.
         type WeightInfo: WeightInfo;
     }
@@ -313,6 +313,13 @@ pub mod pallet {
             }
 
             // Fill orchestrator chain collators up to min_num_orchestrator_chain
+            // Give priority to invulnerables
+            let num_missing_orchestrator_collators = min_num_orchestrator_chain - new_assigned.orchestrator_chain.len();
+            let invulnerables_for_orchestrator = T::RemoveInvulnerables::remove_invulnerables(&mut new_collators, num_missing_orchestrator_collators);
+            new_assigned
+                .fill_orchestrator_chain_collators(min_num_orchestrator_chain, &mut invulnerables_for_orchestrator.into_iter());
+            // If there are no enough invulnerables, or if the invulnerables are currently assigned to other chains,
+            // fill orchestrator chain with regular collators
             let mut new_collators = new_collators.into_iter();
             new_assigned
                 .fill_orchestrator_chain_collators(min_num_orchestrator_chain, &mut new_collators);

pallets/collator-assignment/src/mock.rs

@@ -30,7 +30,7 @@ use {
         traits::{BlakeTwo256, IdentityLookup},
         BuildStorage,
     },
-    tp_traits::ParaId,
+    tp_traits::{ParaId, RemoveInvulnerables},
 };
 
 type Block = frame_system::mocking::MockBlock<Test>;
@@ -187,6 +187,7 @@ impl pallet_collator_assignment::Config for Test {
     type ContainerChains = ContainerChainsGetter;
     type ShouldRotateAllCollators = RotateCollatorsEveryNSessions<CollatorRotationSessionPeriod>;
     type GetRandomnessForNextBlock = MockGetRandomnessForNextBlock;
+    type RemoveInvulnerables = RemoveAccountIdsAbove100;
     type WeightInfo = ();
 }
 
@@ -222,3 +223,20 @@ pub fn run_to_block(n: u64) {
         CollatorAssignment::on_finalize(x);
     }
 }
+
+/// Any AccountId >= 100 will be considered an invulnerable
+pub struct RemoveAccountIdsAbove100;
+
+impl RemoveInvulnerables<u64> for RemoveAccountIdsAbove100 {
+    fn remove_invulnerables(collators: &mut Vec<u64>, num_invulnerables: usize) -> Vec<u64> {
+        let mut invulnerables = vec![];
+        collators.retain(|x| if invulnerables.len() < num_invulnerables && x >= 100 {
+            invulnerables.push(*x);
+            false
+        } else {
+            true
+        });
+
+        invulnerables
+    }
+}

primitives/traits/src/lib.rs

@@ -62,3 +62,9 @@ pub trait GetSessionIndex<SessionIndex> {
 pub trait ShouldRotateAllCollators<SessionIndex> {
     fn should_rotate_all_collators(session_index: SessionIndex) -> bool;
 }
+
+/// Helper trait for pallet_collator_assignment to be able to give priority to invulnerables
+pub trait RemoveInvulnerables<AccountId> {
+    /// Remove the first n invulnerables from the list of collators. The order should be respected.
+    fn remove_invulnerables(collators: &mut Vec<AccountId>, num_invulnerables: usize) -> Vec<AccountId>;
+}

runtime/dancebox/src/lib.rs

@@ -29,6 +29,7 @@ use sp_version::NativeVersion;
 
 #[cfg(any(feature = "std", test))]
 pub use sp_runtime::BuildStorage;
+use tp_traits::RemoveInvulnerables;
 
 pub mod migrations;
 pub mod weights;
@@ -650,6 +651,25 @@ impl GetRandomnessForNextBlock<u32> for BabeGetRandomnessForNextBlock {
     }
 }
 
+pub struct RemoveInvulnerablesImpl;
+
+impl RemoveInvulnerables<AccountId> for RemoveInvulnerablesImpl {
+    fn remove_invulnerables(collators: &mut Vec<AccountId>, num_invulnerables: usize) -> Vec<AccountId> {
+        // TODO: check if this works on session changes
+        let all_invulnerables = pallet_invulnerables::Invulnerables::<Runtime>::get();
+        let mut invulnerables = vec![];
+        // TODO: use binary_search when invulnerables are sorted
+        collators.retain(|x| if invulnerables.len() < num_invulnerables && all_invulnerables.contains(x) {
+            invulnerables.push(x.clone());
+            false
+        } else {
+            true
+        });
+
+        invulnerables
+    }
+}
+
 impl pallet_collator_assignment::Config for Runtime {
     type RuntimeEvent = RuntimeEvent;
     type HostConfiguration = Configuration;
@@ -658,6 +678,7 @@ impl pallet_collator_assignment::Config for Runtime {
     type ShouldRotateAllCollators =
         RotateCollatorsEveryNSessions<ConfigurationCollatorRotationSessionPeriod>;
     type GetRandomnessForNextBlock = BabeGetRandomnessForNextBlock;
+    type RemoveInvulnerables = RemoveInvulnerablesImpl;
     type WeightInfo = pallet_collator_assignment::weights::SubstrateWeight<Runtime>;
 }