From d8dc0ea53614e2a0ed881dbfa854c125def20cc9 Mon Sep 17 00:00:00 2001 From: Tomasz Polaczyk Date: Tue, 22 Jul 2025 15:53:11 +0200 Subject: [PATCH 1/7] Add helper methods to AssignedCollators to avoid is_empty checks --- .../authorities-noting/src/lib.rs | 3 +- .../authorities-noting/src/tests.rs | 24 ++-- .../authorities-noting-inherent/src/mock.rs | 11 +- primitives/collator-assignment/Cargo.toml | 2 +- primitives/collator-assignment/src/lib.rs | 106 +++++++++++++++++- 5 files changed, 119 insertions(+), 27 deletions(-) diff --git a/container-chain-pallets/authorities-noting/src/lib.rs b/container-chain-pallets/authorities-noting/src/lib.rs index 401a2b8..f5133b4 100644 --- a/container-chain-pallets/authorities-noting/src/lib.rs +++ b/container-chain-pallets/authorities-noting/src/lib.rs @@ -371,8 +371,7 @@ impl Pallet { // Read those authorities assigned to this chain let authorities = assignment - .container_chains - .get(¶_id) + .get_container_chain(¶_id) .ok_or(Error::::NoAuthoritiesFound)?; Ok(authorities.clone()) } diff --git a/container-chain-pallets/authorities-noting/src/tests.rs b/container-chain-pallets/authorities-noting/src/tests.rs index 39f6947..cca7725 100644 --- a/container-chain-pallets/authorities-noting/src/tests.rs +++ b/container-chain-pallets/authorities-noting/src/tests.rs @@ -51,8 +51,7 @@ fn test_authorities_insertion_right_para_id() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get(), vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get(), vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof(); @@ -85,8 +84,7 @@ fn test_authorities_insertion_wrong_para_id() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get() + 1, vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get() + 1, vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof(); @@ -119,8 +117,7 @@ fn test_authorities_insertion_right_para_id_solochain() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get(), vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get(), vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof_solochain(); @@ -141,8 +138,7 @@ fn test_not_inserting_inherent() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get(), vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get(), vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof(); @@ -184,8 +180,7 @@ fn encode_proof_for_benchmarks() { assignment .authority_assignment - .container_chains - .insert(container_chain_para_id, vec![10u64, 11u64]); + .insert_container_chain(container_chain_para_id, vec![10u64, 11u64]); assignment.session_index = 0; // TODO let (root_b, proof_b) = assignment.into_state_root_and_proof(); @@ -230,8 +225,7 @@ fn test_set_authorities() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get(), vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get(), vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof(); @@ -270,8 +264,7 @@ fn test_set_orchestrator_para_id() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get(), vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get(), vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof(); @@ -334,8 +327,7 @@ fn weights_assigned_to_extrinsics_are_correct() { let mut assignment = AuthorityAssignmentSproofBuilder::::default(); assignment .authority_assignment - .container_chains - .insert(ParachainId::get(), vec![10u64, 11u64]); + .insert_container_chain(ParachainId::get(), vec![10u64, 11u64]); let (orchestrator_chain_root, orchestrator_chain_state) = assignment.into_state_root_and_proof(); diff --git a/container-chain-primitives/authorities-noting-inherent/src/mock.rs b/container-chain-primitives/authorities-noting-inherent/src/mock.rs index 55e1fe8..4581245 100644 --- a/container-chain-primitives/authorities-noting-inherent/src/mock.rs +++ b/container-chain-primitives/authorities-noting-inherent/src/mock.rs @@ -36,7 +36,6 @@ use { dp_collator_assignment::AssignedCollators, nimbus_primitives::NimbusId, sp_inherents::{InherentData, InherentDataProvider}, - sp_std::collections::btree_map::BTreeMap, test_relay_sproof_builder::{ AuthorityAssignmentSproofBuilder, HeaderAs, ParaHeaderSproofBuilder, ParaHeaderSproofBuilderItem, @@ -134,13 +133,11 @@ impl MockAuthoritiesNotingInherentDataProvider { pub fn build_sproof_builder(&self) -> (ParaHeaderSproofBuilder, sp_trie::StorageProof) { let mut sproof_builder = ParaHeaderSproofBuilder::default(); - let container_chains = - BTreeMap::from_iter([(self.container_para_id, self.authorities.clone())]); + let mut authority_assignment = AssignedCollators::default(); + authority_assignment + .insert_container_chain(self.container_para_id, self.authorities.clone()); let assignment = AuthorityAssignmentSproofBuilder:: { - authority_assignment: AssignedCollators { - orchestrator_chain: vec![], - container_chains, - }, + authority_assignment, session_index: 0, }; diff --git a/primitives/collator-assignment/Cargo.toml b/primitives/collator-assignment/Cargo.toml index c84aeb7..7414f4c 100644 --- a/primitives/collator-assignment/Cargo.toml +++ b/primitives/collator-assignment/Cargo.toml @@ -23,7 +23,7 @@ frame-support = { workspace = true, optional = true } parity-scale-codec = { workspace = true } scale-info = { workspace = true } sp-core = { workspace = true } -sp-runtime = { workspace = true, optional = true } +sp-runtime = { workspace = true } sp-state-machine = { workspace = true, optional = true } sp-std = { workspace = true } sp-trie = { workspace = true, optional = true } diff --git a/primitives/collator-assignment/src/lib.rs b/primitives/collator-assignment/src/lib.rs index be279fa..acf3acc 100644 --- a/primitives/collator-assignment/src/lib.rs +++ b/primitives/collator-assignment/src/lib.rs @@ -20,14 +20,18 @@ use { cumulus_primitives_core::ParaId, parity_scale_codec::{Decode, Encode}, scale_info::prelude::collections::BTreeMap, + scale_info::prelude::collections::BTreeSet, + sp_runtime::Saturating, sp_std::vec::Vec, }; #[derive(Clone, Encode, Decode, PartialEq, sp_core::RuntimeDebug, scale_info::TypeInfo)] #[cfg_attr(feature = "std", derive(serde::Serialize, serde::Deserialize))] pub struct AssignedCollators { + // This must be a Vec and not a BTreeSet because the order is important pub orchestrator_chain: Vec, - pub container_chains: BTreeMap>, + // This is private to try to handle the edge case of empty vec here instead of in caller code + container_chains: BTreeMap>, } // Manual default impl that does not require AccountId: Default @@ -44,6 +48,11 @@ impl AssignedCollators where AccountId: PartialEq, { + /// Find the `ParaId` where collator `x` is assigned to. Returns `None` if not assigned to any. + /// + /// `orchestrator_chain_para_id` is used to simplify the return value: returning `Some` means + /// the collator is assigned somewhere, but it could be a container chain or the orchestrator + /// chain. pub fn para_id_of(&self, x: &AccountId, orchestrator_chain_para_id: ParaId) -> Option { for (id, cs) in self.container_chains.iter() { if cs.contains(x) { @@ -58,6 +67,9 @@ where None } + /// Map the collator type. Returns all collators in the same order as the input. + // TODO: if we didn't need to support this method we could change all the `Vec` into + // `BTreeSet`. pub fn map(&self, mut f: F) -> AssignedCollators where F: FnMut(&AccountId) -> T, @@ -74,4 +86,96 @@ where a } + + /// Get collators assigned to container chain `para_id`. Handles the edge case of an empty list. + /// If this returns Some, the Vec can be assumed to not be empty. + pub fn get_container_chain(&self, para_id: &ParaId) -> Option<&Vec> { + let x = self.container_chains.get(para_id); + + // Filter out empty assignment, return None instead + match x { + Some(x) if x.is_empty() => None, + x => x, + } + } + + /// Return container_chains map with all the chains that have at least 1 assigned collator. + /// Ignores orchestrator chain. + pub fn into_container_chains_with_collators(mut self) -> BTreeMap> { + self.cleanup_empty(); + + self.container_chains + } + + /// Remove all the container chains whose list of assigned collators is empty. That is logically + /// equivalent to that para id not being in the map. + pub fn cleanup_empty(&mut self) { + self.container_chains.retain(|_, v| { + // Keep the entries whose value is not an empty list + !v.is_empty() + }); + } + + /// Merge `orchestrator_chain` into `container_chains` map as `orchestrator_para_id`, and return + /// the resulting map. + pub fn into_single_map( + mut self, + orchestrator_para_id: ParaId, + ) -> BTreeMap> { + self.container_chains + .insert(orchestrator_para_id, self.orchestrator_chain); + + self.container_chains + } + + /// Create `Self` from a single map, removing `orchestrator_para_id`. + /// This calls `Self::cleanup_empty` internally, so the resulting assignment will only include + /// chains with collators. + pub fn from_single_map( + mut container_chains: BTreeMap>, + orchestrator_para_id: &ParaId, + ) -> Self { + let orchestrator_chain = container_chains + .remove(orchestrator_para_id) + .unwrap_or_default(); + + let mut x = Self { + orchestrator_chain, + container_chains, + }; + + x.cleanup_empty(); + + x + } + + /// Return the total number of collators assigned to all chains, orchestrator + containers + pub fn count_collators(&self) -> usize { + let mut num_collators: usize = 0; + num_collators.saturating_accrue(self.orchestrator_chain.len()); + for collators in self.container_chains.values() { + num_collators.saturating_accrue(collators.len()); + } + + num_collators + } + + /// Return all the container chain para ids with at least 1 collator assigned + pub fn container_para_ids(&self) -> BTreeSet { + self.container_chains + .iter() + .filter_map(|(k, v)| if v.is_empty() { None } else { Some(k.clone()) }) + .collect() + } + + /// If `v` is not empty, insert into `self.container_chains` + pub fn insert_container_chain(&mut self, k: ParaId, v: Vec) { + if !v.is_empty() { + self.container_chains.insert(k, v); + } + } + /// Remove container chain + pub fn remove_container_chain(&mut self, k: &ParaId) -> Option> { + self.container_chains.remove(k) + } } From c4f78b32724c7df95e4d8f25f6caafe0d4401a0e Mon Sep 17 00:00:00 2001 From: Tomasz Polaczyk Date: Tue, 22 Jul 2025 17:38:42 +0200 Subject: [PATCH 2/7] Add into_collators --- primitives/collator-assignment/src/lib.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/primitives/collator-assignment/src/lib.rs b/primitives/collator-assignment/src/lib.rs index acf3acc..4ca8c61 100644 --- a/primitives/collator-assignment/src/lib.rs +++ b/primitives/collator-assignment/src/lib.rs @@ -179,3 +179,21 @@ where self.container_chains.remove(k) } } + +impl AssignedCollators +where + AccountId: core::cmp::Ord, +{ + /// Return all collators assigned to some chain. Includes orchestartor chain. + pub fn into_collators(mut self) -> BTreeSet { + let mut collators = BTreeSet::new(); + collators.extend(core::mem::take(&mut self.orchestrator_chain)); + collators.extend( + self.into_container_chains_with_collators() + .into_values() + .flatten(), + ); + + collators + } +} From cf98374eaf1481dc4513f25f7386acd58049ec92 Mon Sep 17 00:00:00 2001 From: Tomasz Polaczyk Date: Tue, 22 Jul 2025 17:03:25 +0200 Subject: [PATCH 3/7] Update setup-protoc@v3 --- .github/workflows/build.yml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 193fdc8..8406555 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -137,7 +137,7 @@ jobs: run: ./.github/scripts/check_toml_lints.sh - name: Install Protoc - uses: arduino/setup-protoc@v1 + uses: arduino/setup-protoc@v3 - name: Clippy run: SKIP_WASM_BUILD=1 cargo clippy --all-targets --locked --workspace --features try-runtime,runtime-benchmarks @@ -191,7 +191,7 @@ jobs: with: ref: ${{ needs.set-tags.outputs.git_ref }} - name: Install Protoc - uses: arduino/setup-protoc@v1 + uses: arduino/setup-protoc@v3 - name: Cargo build uses: ./.github/workflow-templates/cargo-build @@ -209,7 +209,7 @@ jobs: with: ref: ${{ needs.set-tags.outputs.git_ref }} - name: Install Protoc - uses: arduino/setup-protoc@v1 + uses: arduino/setup-protoc@v3 - name: Run sccache-cache uses: mozilla-actions/sccache-action@v0.0.3 - name: Setup Variables From 51e2a48e3b284b6f1a13a6ecdd940104f6cedbda Mon Sep 17 00:00:00 2001 From: Tomasz Polaczyk Date: Tue, 22 Jul 2025 18:26:10 +0200 Subject: [PATCH 4/7] invert_map --- primitives/collator-assignment/src/lib.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/primitives/collator-assignment/src/lib.rs b/primitives/collator-assignment/src/lib.rs index 4ca8c61..ea96ee0 100644 --- a/primitives/collator-assignment/src/lib.rs +++ b/primitives/collator-assignment/src/lib.rs @@ -196,4 +196,22 @@ where collators } + + /// Invert the map relation and return a map of collator to para id. + /// Useful for testing and for checking the assignment of multiple collators at once. + pub fn invert_map(mut self, orchestrator_para_id: ParaId) -> BTreeMap { + let mut x = BTreeMap::new(); + + for collator in core::mem::take(&mut self.orchestrator_chain) { + x.insert(collator, orchestrator_para_id); + } + + for (para_id, collators) in self.container_chains { + for collator in collators { + x.insert(collator, para_id); + } + } + + x + } } From ec554bdc15a3916c588de0f93ad2e8eb3d7c01ae Mon Sep 17 00:00:00 2001 From: Tomasz Polaczyk Date: Wed, 23 Jul 2025 10:06:34 +0200 Subject: [PATCH 5/7] container_para_id_of --- primitives/collator-assignment/src/lib.rs | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/primitives/collator-assignment/src/lib.rs b/primitives/collator-assignment/src/lib.rs index ea96ee0..c3aadfc 100644 --- a/primitives/collator-assignment/src/lib.rs +++ b/primitives/collator-assignment/src/lib.rs @@ -54,10 +54,8 @@ where /// the collator is assigned somewhere, but it could be a container chain or the orchestrator /// chain. pub fn para_id_of(&self, x: &AccountId, orchestrator_chain_para_id: ParaId) -> Option { - for (id, cs) in self.container_chains.iter() { - if cs.contains(x) { - return Some(*id); - } + if let Some(id) = self.container_para_id_of(x) { + return Some(id); } if self.orchestrator_chain.contains(x) { @@ -67,6 +65,18 @@ where None } + /// Find the container `ParaId` where collator `x` is assigned to. Returns `None` if + /// not assigned to any. The collator could still be assigned to the orchestrator chain. + pub fn container_para_id_of(&self, x: &AccountId) -> Option { + for (id, cs) in self.container_chains.iter() { + if cs.contains(x) { + return Some(*id); + } + } + + None + } + /// Map the collator type. Returns all collators in the same order as the input. // TODO: if we didn't need to support this method we could change all the `Vec` into // `BTreeSet`. From ac41de30d93fc563f36f96d65d498b81a82e74ba Mon Sep 17 00:00:00 2001 From: Tomasz Polaczyk Date: Tue, 5 Aug 2025 12:27:07 +0200 Subject: [PATCH 6/7] Add iterator and some cleanup --- primitives/collator-assignment/src/lib.rs | 61 ++++++++++++++++------- 1 file changed, 43 insertions(+), 18 deletions(-) diff --git a/primitives/collator-assignment/src/lib.rs b/primitives/collator-assignment/src/lib.rs index c3aadfc..ac90a35 100644 --- a/primitives/collator-assignment/src/lib.rs +++ b/primitives/collator-assignment/src/lib.rs @@ -66,7 +66,7 @@ where } /// Find the container `ParaId` where collator `x` is assigned to. Returns `None` if - /// not assigned to any. The collator could still be assigned to the orchestrator chain. + /// not assigned to any. If this returns `None`, the collator could still be assigned to the orchestrator chain. pub fn container_para_id_of(&self, x: &AccountId) -> Option { for (id, cs) in self.container_chains.iter() { if cs.contains(x) { @@ -109,16 +109,8 @@ where } } - /// Return container_chains map with all the chains that have at least 1 assigned collator. - /// Ignores orchestrator chain. - pub fn into_container_chains_with_collators(mut self) -> BTreeMap> { - self.cleanup_empty(); - - self.container_chains - } - /// Remove all the container chains whose list of assigned collators is empty. That is logically - /// equivalent to that para id not being in the map. + /// equivalent to that para id not being in the map. Call this before serializing this type. pub fn cleanup_empty(&mut self) { self.container_chains.retain(|_, v| { // Keep the entries whose value is not an empty list @@ -126,14 +118,26 @@ where }); } + /// Return container_chains map with all the chains that have at least 1 assigned collator. + /// Ignores orchestrator chain. + pub fn into_container_chains_with_collators(mut self) -> BTreeMap> { + self.cleanup_empty(); + + self.container_chains + } + /// Merge `orchestrator_chain` into `container_chains` map as `orchestrator_para_id`, and return - /// the resulting map. + /// the resulting map. Empty chains will be removed, including orchestrator chain if its empty. pub fn into_single_map( mut self, orchestrator_para_id: ParaId, ) -> BTreeMap> { - self.container_chains - .insert(orchestrator_para_id, self.orchestrator_chain); + self.container_chains.insert( + orchestrator_para_id, + core::mem::take(&mut self.orchestrator_chain), + ); + + self.cleanup_empty(); self.container_chains } @@ -170,12 +174,32 @@ where num_collators } + /// Iterate over all the non-empty container chains. + pub fn container_chains_iter(&self) -> impl Iterator)> { + self.container_chains.iter().filter_map( + |(k, v)| { + if v.is_empty() { + None + } else { + Some((k, v)) + } + }, + ) + } + + /// Convenience method to get a new `BTreeMap` of all the non-empty container chains. + /// Prefer using some other method to avoid creating this temporary map: + /// * `container_chains_iter` if you just need to iterate + /// * `get_container_chain` to query 1 chain + /// * `insert_container_chain` / `remove_container_chain` to add new chains + /// * `into_single_map` / `from_single_map` if you prefer to work on a raw `BTreeMap` + pub fn container_chains(&self) -> BTreeMap<&ParaId, &Vec> { + self.container_chains_iter().collect() + } + /// Return all the container chain para ids with at least 1 collator assigned pub fn container_para_ids(&self) -> BTreeSet { - self.container_chains - .iter() - .filter_map(|(k, v)| if v.is_empty() { None } else { Some(k.clone()) }) - .collect() + self.container_chains_iter().map(|(k, _v)| *k).collect() } /// If `v` is not empty, insert into `self.container_chains` @@ -184,6 +208,7 @@ where self.container_chains.insert(k, v); } } + /// Remove container chain pub fn remove_container_chain(&mut self, k: &ParaId) -> Option> { self.container_chains.remove(k) @@ -192,7 +217,7 @@ where impl AssignedCollators where - AccountId: core::cmp::Ord, + AccountId: Ord, { /// Return all collators assigned to some chain. Includes orchestartor chain. pub fn into_collators(mut self) -> BTreeSet { From 9d44f06231e46fd601b8d30b9aba3d0c1d4ecf51 Mon Sep 17 00:00:00 2001 From: tmpolaczyk <44604217+tmpolaczyk@users.noreply.github.com> Date: Wed, 6 Aug 2025 12:49:41 +0200 Subject: [PATCH 7/7] Update mock.rs --- .../authorities-noting-inherent/src/mock.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/container-chain-primitives/authorities-noting-inherent/src/mock.rs b/container-chain-primitives/authorities-noting-inherent/src/mock.rs index 5fca88a..4581245 100644 --- a/container-chain-primitives/authorities-noting-inherent/src/mock.rs +++ b/container-chain-primitives/authorities-noting-inherent/src/mock.rs @@ -29,7 +29,6 @@ use { crate::ContainerChainAuthoritiesInherentData, - alloc::collections::btree_map::BTreeMap, cumulus_primitives_core::ParaId, cumulus_primitives_parachain_inherent::{ ParachainInherentData, INHERENT_IDENTIFIER as PARACHAIN_SYSTEM_INHERENT_IDENTIFIER,