Support parathread registration and collator assignment (#393)

pallets/collator-assignment/src/lib.rs

@@ -166,14 +166,23 @@ pub mod pallet {
             let session_delay = T::SessionIndex::one();
             let target_session_index = current_session_index.saturating_add(session_delay);
             // We get the containerChains that we will have at the target session
-            let mut container_chain_ids =
+            let container_chains =
                 T::ContainerChains::session_container_chains(target_session_index);
-            let mut parathreads = T::ContainerChains::session_parathreads(target_session_index);
-            let num_total_registered_paras = container_chain_ids.len() as u32;
+            let num_total_registered_paras =
+                (container_chains.parachains.len() + container_chains.parathreads.len()) as u32;
+            let mut container_chain_ids = container_chains.parachains;
+            let mut parathreads: Vec<_> = container_chains
+                .parathreads
+                .into_iter()
+                .map(|(para_id, _)| para_id)
+                .collect();
             // Remove the containerChains that do not have enough credits for block production
             T::RemoveParaIdsWithNoCredits::remove_para_ids_with_no_credits(
                 &mut container_chain_ids,
             );
+            // TODO: parathreads should be treated a bit differently, they don't need to have the same amount of credits
+            // as paratherads because they will not be producing blocks on every slot.
+            T::RemoveParaIdsWithNoCredits::remove_para_ids_with_no_credits(&mut parathreads);
 
             // If the random_seed is all zeros, we don't shuffle the list of collators nor the list
             // of container chains.

pallets/collator-assignment/src/mock.rs

@@ -31,7 +31,10 @@ use {
         BuildStorage,
     },
     sp_std::collections::btree_map::BTreeMap,
-    tp_traits::{ParaId, RemoveInvulnerables, RemoveParaIdsWithNoCredits},
+    tp_traits::{
+        ParaId, ParathreadParams, RemoveInvulnerables, RemoveParaIdsWithNoCredits,
+        SessionContainerChains,
+    },
     tracing_subscriber::{layer::SubscriberExt, FmtSubscriber},
 };
 
@@ -161,22 +164,32 @@ impl GetCollators<u64, u32> for CollatorsGetter {
 pub struct ContainerChainsGetter;
 
 impl tp_traits::GetSessionContainerChains<u32> for ContainerChainsGetter {
-    fn session_container_chains(_session_index: u32) -> Vec<ParaId> {
-        MockData::mock()
+    fn session_container_chains(_session_index: u32) -> SessionContainerChains {
+        let parachains = MockData::mock()
             .container_chains
             .iter()
             .cloned()
-            .map(ParaId::from)
-            .collect()
-    }
+            .map(|para_id| ParaId::from(para_id))
+            .collect();
 
-    fn session_parathreads(_session_index: u32) -> Vec<ParaId> {
-        MockData::mock()
+        let parathreads = MockData::mock()
             .parathreads
             .iter()
             .cloned()
-            .map(ParaId::from)
-            .collect()
+            .map(|para_id| {
+                (
+                    ParaId::from(para_id),
+                    ParathreadParams {
+                        slot_frequency: Default::default(),
+                    },
+                )
+            })
+            .collect();
+
+        SessionContainerChains {
+            parachains,
+            parathreads,
+        }
     }
 
     #[cfg(feature = "runtime-benchmarks")]

pallets/collator-assignment/src/tests.rs

@@ -821,18 +821,18 @@ fn assign_collators_rotation_parathreads_are_shuffled() {
 
             // 4 collators so we can only assign to one parathread
             m.collators = vec![1, 2, 3, 4];
-            m.parathreads = vec![5001, 5002];
+            m.parathreads = vec![3001, 3002];
         });
         assert_eq!(assigned_collators(), initial_collators(),);
         run_to_block(11);
 
         let initial_assignment =
-            BTreeMap::from_iter(vec![(1, 1000), (2, 1000), (3, 5001), (4, 5001)]);
+            BTreeMap::from_iter(vec![(1, 1000), (2, 1000), (3, 3001), (4, 3001)]);
 
         assert_eq!(assigned_collators(), initial_assignment,);
 
         MockData::mutate(|m| {
-            // Seed chosen manually to see the case where parathread 5002 is given priority
+            // Seed chosen manually to see the case where parathread 3002 is given priority
             m.random_seed = [2; 32];
         });
 
@@ -841,7 +841,7 @@ fn assign_collators_rotation_parathreads_are_shuffled() {
         // Random assignment depends on the seed, shouldn't change unless the algorithm changes
         // Test that container chains are shuffled because 1001 does not have priority
         let shuffled_assignment =
-            BTreeMap::from_iter(vec![(1, 1000), (2, 5002), (3, 1000), (4, 5002)]);
+            BTreeMap::from_iter(vec![(1, 1000), (2, 3002), (3, 1000), (4, 3002)]);
 
         assert_eq!(assigned_collators(), shuffled_assignment,);
     });
@@ -1231,3 +1231,57 @@ fn collator_assignment_includes_empty_chains() {
         assert_eq!(assigned_collators, expected);
     });
 }
+
+#[test]
+fn collator_assignment_remove_parachains_without_credits() {
+    new_test_ext().execute_with(|| {
+        run_to_block(1);
+
+        MockData::mutate(|m| {
+            m.collators_per_container = 2;
+            m.collators_per_parathread = 2;
+            m.min_orchestrator_chain_collators = 2;
+            m.max_orchestrator_chain_collators = 5;
+
+            m.collators = vec![1, 2, 3, 4, 5, 6, 7];
+            m.container_chains = vec![2000, 5001, 5002];
+            m.parathreads = vec![]
+        });
+        assert_eq!(assigned_collators(), initial_collators(),);
+        run_to_block(11);
+
+        let assigned_collators = CollatorContainerChain::<Test>::get();
+        let expected = AssignedCollators {
+            orchestrator_chain: vec![1, 2, 3, 4, 5],
+            container_chains: BTreeMap::from_iter(vec![(2000.into(), vec![6, 7])]),
+        };
+        assert_eq!(assigned_collators, expected);
+    });
+}
+
+#[test]
+fn collator_assignment_remove_parathreads_without_credits() {
+    new_test_ext().execute_with(|| {
+        run_to_block(1);
+
+        MockData::mutate(|m| {
+            m.collators_per_container = 2;
+            m.collators_per_parathread = 2;
+            m.min_orchestrator_chain_collators = 2;
+            m.max_orchestrator_chain_collators = 5;
+
+            m.collators = vec![1, 2, 3, 4, 5, 6, 7];
+            m.container_chains = vec![];
+            m.parathreads = vec![3000, 5001, 5002]
+        });
+        assert_eq!(assigned_collators(), initial_collators(),);
+        run_to_block(11);
+
+        let assigned_collators = CollatorContainerChain::<Test>::get();
+        let expected = AssignedCollators {
+            orchestrator_chain: vec![1, 2, 3, 4, 5],
+            container_chains: BTreeMap::from_iter(vec![(3000.into(), vec![6, 7])]),
+        };
+        assert_eq!(assigned_collators, expected);
+    });
+}

pallets/configuration/src/lib.rs

@@ -131,7 +131,7 @@ impl HostConfiguration {
         if self.max_orchestrator_collators < self.min_orchestrator_collators {
             return Err(InconsistentError::MaxCollatorsLowerThanMinCollators);
         }
-        if self.collators_per_parathread != 1 || self.parathreads_per_collator != 1 {
+        if self.parathreads_per_collator != 1 {
             return Err(InconsistentError::UnimplementedParameter);
         }
         Ok(())

pallets/registrar/src/benchmarks.rs

@@ -25,7 +25,7 @@ use {
     sp_core::Get,
     sp_std::{vec, vec::Vec},
     tp_container_chain_genesis_data::{ContainerChainGenesisData, ContainerChainGenesisDataItem},
-    tp_traits::ParaId,
+    tp_traits::{ParaId, SlotFrequency},
 };
 
 /// Create a funded user.
@@ -355,5 +355,93 @@ mod benchmarks {
         assert!(Pallet::<T>::registered_para_ids().contains(&ParaId::from(1000)));
     }
 
+    #[benchmark]
+    fn register_parathread(x: Linear<5, 3_000_000>, y: Linear<1, 50>, z: Linear<1, 10>) {
+        let mut data = vec![];
+        // Number of keys
+        for _i in 1..z {
+            data.push((b"code".to_vec(), vec![1; (x / z) as usize]).into())
+        }
+
+        let slot_frequency = SlotFrequency::default();
+        let storage = new_genesis_data(data);
+
+        for i in 1..y {
+            // Twice the deposit just in case
+            let (caller, _deposit_amount) =
+                create_funded_user::<T>("caller", i, T::DepositAmount::get());
+            Pallet::<T>::register_parathread(
+                RawOrigin::Signed(caller.clone()).into(),
+                i.into(),
+                slot_frequency.clone(),
+                storage.clone(),
+            )
+            .unwrap();
+        }
+
+        // We should have registered y-1
+        assert_eq!(Pallet::<T>::pending_verification().len(), (y - 1) as usize);
+
+        let (caller, _deposit_amount) =
+            create_funded_user::<T>("caller", 0, T::DepositAmount::get());
+
+        #[extrinsic_call]
+        Pallet::<T>::register_parathread(
+            RawOrigin::Signed(caller),
+            Default::default(),
+            slot_frequency,
+            storage,
+        );
+
+        // verification code
+        assert_eq!(Pallet::<T>::pending_verification().len(), y as usize);
+        assert!(Pallet::<T>::registrar_deposit(ParaId::default()).is_some());
+    }
+
+    #[benchmark]
+    fn set_parathread_params(y: Linear<1, 50>) {
+        let storage = vec![(vec![1; 4], vec![1; 3_000_000usize]).into()];
+        let storage = new_genesis_data(storage);
+        let slot_frequency = SlotFrequency::default();
+
+        // Deregister all the existing chains to avoid conflicts with the new ones
+        for para_id in Pallet::<T>::registered_para_ids() {
+            Pallet::<T>::deregister(RawOrigin::Root.into(), para_id).unwrap();
+        }
+
+        for i in 0..y {
+            // Twice the deposit just in case
+            let (caller, _deposit_amount) =
+                create_funded_user::<T>("caller", i, T::DepositAmount::get());
+            Pallet::<T>::register_parathread(
+                RawOrigin::Signed(caller.clone()).into(),
+                i.into(),
+                slot_frequency.clone(),
+                storage.clone(),
+            )
+            .unwrap();
+            T::RegistrarHooks::benchmarks_ensure_valid_for_collating(i.into());
+            Pallet::<T>::mark_valid_for_collating(RawOrigin::Root.into(), i.into()).unwrap();
+        }
+
+        let new_slot_frequency = SlotFrequency { min: 2, max: 2 };
+
+        #[extrinsic_call]
+        Pallet::<T>::set_parathread_params(
+            RawOrigin::Root,
+            (y - 1).into(),
+            new_slot_frequency.clone(),
+        );
+
+        // Start a new session
+        Pallet::<T>::initializer_on_new_session(&T::SessionDelay::get());
+
+        // Check y-1 has new slot frequency
+        assert_eq!(
+            Pallet::<T>::parathread_params(&ParaId::from(y - 1)).map(|x| x.slot_frequency),
+            Some(new_slot_frequency)
+        );
+    }
+
     impl_benchmark_test_suite!(Pallet, crate::mock::new_test_ext(), crate::mock::Test);
 }