statement-distribution: fix filtering of statements for elastic parachains (#3879)

fixes https://github.com/paritytech/polkadot-sdk/issues/3775

Additionally moves the claim queue fetch utilities into
`subsystem-util`.

TODO:
- [x] fix tests
- [x] add elastic scaling tests

---------

Signed-off-by: Andrei Sandu <andrei-mihail@parity.io>

polkadot/node/collation-generation/src/error.rs

@@ -27,6 +27,8 @@ pub enum Error {
 	#[error(transparent)]
 	Util(#[from] polkadot_node_subsystem_util::Error),
 	#[error(transparent)]
+	UtilRuntime(#[from] polkadot_node_subsystem_util::runtime::Error),
+	#[error(transparent)]
 	Erasure(#[from] polkadot_erasure_coding::Error),
 	#[error("Parachain backing state not available in runtime.")]
 	MissingParaBackingState,

polkadot/node/collation-generation/src/lib.rs

@@ -38,25 +38,23 @@ use polkadot_node_primitives::{
 	SubmitCollationParams,
 };
 use polkadot_node_subsystem::{
-	messages::{CollationGenerationMessage, CollatorProtocolMessage, RuntimeApiRequest},
+	messages::{CollationGenerationMessage, CollatorProtocolMessage},
 	overseer, ActiveLeavesUpdate, FromOrchestra, OverseerSignal, RuntimeApiError, SpawnedSubsystem,
 	SubsystemContext, SubsystemError, SubsystemResult,
 };
 use polkadot_node_subsystem_util::{
-	has_required_runtime, request_async_backing_params, request_availability_cores,
-	request_claim_queue, request_para_backing_state, request_persisted_validation_data,
-	request_validation_code, request_validation_code_hash, request_validators,
+	request_async_backing_params, request_availability_cores, request_para_backing_state,
+	request_persisted_validation_data, request_validation_code, request_validation_code_hash,
+	request_validators,
+	vstaging::{fetch_claim_queue, fetch_next_scheduled_on_core},
 };
 use polkadot_primitives::{
 	collator_signature_payload, CandidateCommitments, CandidateDescriptor, CandidateReceipt,
 	CollatorPair, CoreIndex, CoreState, Hash, Id as ParaId, OccupiedCoreAssumption,
-	PersistedValidationData, ScheduledCore, ValidationCodeHash,
+	PersistedValidationData, ValidationCodeHash,
 };
 use sp_core::crypto::Pair;
-use std::{
-	collections::{BTreeMap, VecDeque},
-	sync::Arc,
-};
+use std::sync::Arc;
 
 mod error;
 
@@ -228,7 +226,9 @@ async fn handle_new_activations<Context>(
 		let availability_cores = availability_cores??;
 		let async_backing_params = async_backing_params?.ok();
 		let n_validators = validators??.len();
-		let maybe_claim_queue = fetch_claim_queue(ctx.sender(), relay_parent).await?;
+		let maybe_claim_queue = fetch_claim_queue(ctx.sender(), relay_parent)
+			.await
+			.map_err(crate::error::Error::UtilRuntime)?;
 
 		// The loop bellow will fill in cores that the para is allowed to build on.
 		let mut cores_to_build_on = Vec::new();
@@ -655,37 +655,3 @@ fn erasure_root(
 	let chunks = polkadot_erasure_coding::obtain_chunks_v1(n_validators, &available_data)?;
 	Ok(polkadot_erasure_coding::branches(&chunks).root())
 }
-
-// Checks if the runtime supports `request_claim_queue` and executes it. Returns `Ok(None)`
-// otherwise. Any [`RuntimeApiError`]s are bubbled up to the caller.
-async fn fetch_claim_queue(
-	sender: &mut impl overseer::CollationGenerationSenderTrait,
-	relay_parent: Hash,
-) -> crate::error::Result<Option<BTreeMap<CoreIndex, VecDeque<ParaId>>>> {
-	if has_required_runtime(
-		sender,
-		relay_parent,
-		RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT,
-	)
-	.await
-	{
-		let res = request_claim_queue(relay_parent, sender).await.await??;
-		Ok(Some(res))
-	} else {
-		gum::trace!(target: LOG_TARGET, "Runtime doesn't support `request_claim_queue`");
-		Ok(None)
-	}
-}
-
-// Returns the next scheduled `ParaId` for a core in the claim queue, wrapped in `ScheduledCore`.
-// This function is supposed to be used in `handle_new_activations` hence the return type.
-fn fetch_next_scheduled_on_core(
-	claim_queue: &BTreeMap<CoreIndex, VecDeque<ParaId>>,
-	core_idx: CoreIndex,
-) -> Option<ScheduledCore> {
-	claim_queue
-		.get(&core_idx)?
-		.front()
-		.cloned()
-		.map(|para_id| ScheduledCore { para_id, collator: None })
-}

polkadot/node/collation-generation/src/tests.rs

@@ -28,7 +28,7 @@ use polkadot_node_subsystem::{
 	ActivatedLeaf,
 };
 use polkadot_node_subsystem_test_helpers::{subsystem_test_harness, TestSubsystemContextHandle};
-use polkadot_node_subsystem_util::TimeoutExt;
+use polkadot_node_subsystem_util::{vstaging::ClaimQueueSnapshot, TimeoutExt};
 use polkadot_primitives::{
 	async_backing::{BackingState, CandidatePendingAvailability},
 	AsyncBackingParams, BlockNumber, CollatorPair, HeadData, PersistedValidationData,
@@ -36,7 +36,10 @@ use polkadot_primitives::{
 };
 use rstest::rstest;
 use sp_keyring::sr25519::Keyring as Sr25519Keyring;
-use std::pin::Pin;
+use std::{
+	collections::{BTreeMap, VecDeque},
+	pin::Pin,
+};
 use test_helpers::{
 	dummy_candidate_descriptor, dummy_hash, dummy_head_data, dummy_validator, make_candidate,
 };
@@ -617,7 +620,7 @@ fn fallback_when_no_validation_code_hash_api(#[case] runtime_version: u32) {
 					_hash,
 					RuntimeApiRequest::ClaimQueue(tx),
 				))) if runtime_version >= RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT => {
-					let res = BTreeMap::<CoreIndex, VecDeque<ParaId>>::new();
+					let res = ClaimQueueSnapshot::new();
 					tx.send(Ok(res)).unwrap();
 				},
 				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
@@ -780,7 +783,7 @@ fn distribute_collation_for_occupied_core_with_async_backing_enabled(#[case] run
 		candidate_hash: Default::default(),
 		candidate_descriptor: dummy_candidate_descriptor(dummy_hash()),
 	})];
-	let claim_queue = BTreeMap::from([(CoreIndex::from(0), VecDeque::from([para_id]))]);
+	let claim_queue = ClaimQueueSnapshot::from([(CoreIndex::from(0), VecDeque::from([para_id]))]);
 
 	test_harness(|mut virtual_overseer| async move {
 		helpers::initialize_collator(&mut virtual_overseer, para_id).await;
@@ -962,7 +965,7 @@ fn no_collation_is_distributed_for_occupied_core_with_async_backing_disabled(
 		candidate_hash: Default::default(),
 		candidate_descriptor: dummy_candidate_descriptor(dummy_hash()),
 	})];
-	let claim_queue = BTreeMap::from([(CoreIndex::from(0), VecDeque::from([para_id]))]);
+	let claim_queue = ClaimQueueSnapshot::from([(CoreIndex::from(0), VecDeque::from([para_id]))]);
 
 	test_harness(|mut virtual_overseer| async move {
 		helpers::initialize_collator(&mut virtual_overseer, para_id).await;
@@ -1050,7 +1053,7 @@ mod helpers {
 		async_backing_params: AsyncBackingParams,
 		cores: Vec<CoreState>,
 		runtime_version: u32,
-		claim_queue: BTreeMap<CoreIndex, VecDeque<ParaId>>,
+		claim_queue: ClaimQueueSnapshot,
 	) {
 		assert_matches!(
 			overseer_recv(virtual_overseer).await,