Expose ClaimQueue via a runtime api and use it in collation-generation (#3580)

The PR adds two things: 1. Runtime API exposing the whole claim queue 2. Consumes the API in `collation-generation` to fetch the next scheduled `ParaEntry` for an occupied core. Related to https://github.com/paritytech/polkadot-sdk/issues/1797
2026-06-19 07:41:02 +00:00 · 2024-03-20 08:55:58 +02:00
parent e659c4b3f7
commit e58e854a32
15 changed files with 532 additions and 52 deletions
@@ -38,21 +38,25 @@ use polkadot_node_primitives::{
 	SubmitCollationParams,
 };
 use polkadot_node_subsystem::{
-	messages::{CollationGenerationMessage, CollatorProtocolMessage},
+	messages::{CollationGenerationMessage, CollatorProtocolMessage, RuntimeApiRequest},
 	overseer, ActiveLeavesUpdate, FromOrchestra, OverseerSignal, RuntimeApiError, SpawnedSubsystem,
 	SubsystemContext, SubsystemError, SubsystemResult,
 };
 use polkadot_node_subsystem_util::{
-	request_async_backing_params, request_availability_cores, request_persisted_validation_data,
-	request_validation_code, request_validation_code_hash, request_validators,
+	has_required_runtime, request_async_backing_params, request_availability_cores,
+	request_claim_queue, request_persisted_validation_data, request_validation_code,
+	request_validation_code_hash, request_validators,
 };
 use polkadot_primitives::{
 	collator_signature_payload, CandidateCommitments, CandidateDescriptor, CandidateReceipt,
-	CollatorPair, CoreState, Hash, Id as ParaId, OccupiedCoreAssumption, PersistedValidationData,
-	ValidationCodeHash,
+	CollatorPair, CoreIndex, CoreState, Hash, Id as ParaId, OccupiedCoreAssumption,
+	PersistedValidationData, ScheduledCore, ValidationCodeHash,
 };
 use sp_core::crypto::Pair;
-use std::sync::Arc;
+use std::{
+	collections::{BTreeMap, VecDeque},
+	sync::Arc,
+};

 mod error;

@@ -223,6 +227,7 @@ async fn handle_new_activations<Context>(
 		let availability_cores = availability_cores??;
 		let n_validators = validators??.len();
 		let async_backing_params = async_backing_params?.ok();
+		let maybe_claim_queue = fetch_claim_queue(ctx.sender(), relay_parent).await?;

 		for (core_idx, core) in availability_cores.into_iter().enumerate() {
 			let _availability_core_timer = metrics.time_new_activations_availability_core();
@@ -239,10 +244,25 @@ async fn handle_new_activations<Context>(
 						// TODO [now]: this assumes that next up == current.
 						// in practice we should only set `OccupiedCoreAssumption::Included`
 						// when the candidate occupying the core is also of the same para.
-						if let Some(scheduled) = occupied_core.next_up_on_available {
-							(scheduled, OccupiedCoreAssumption::Included)
-						} else {
-							continue
+						let res = match maybe_claim_queue {
+							Some(ref claim_queue) => {
+								// read what's in the claim queue for this core
+								fetch_next_scheduled_on_core(
+									claim_queue,
+									CoreIndex(core_idx as u32),
+								)
+							},
+							None => {
+								// Runtime doesn't support claim queue runtime api. Fallback to
+								// `next_up_on_available`
+								occupied_core.next_up_on_available
+							},
+						}
+						.map(|scheduled| (scheduled, OccupiedCoreAssumption::Included));
+
+						match res {
+							Some(res) => res,
+							None => continue,
 						}
 					},
 					_ => {
@@ -600,3 +620,37 @@ 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 })
+}
@@ -25,15 +25,18 @@ use polkadot_node_primitives::{BlockData, Collation, CollationResult, MaybeCompr
 use polkadot_node_subsystem::{
 	errors::RuntimeApiError,
 	messages::{AllMessages, RuntimeApiMessage, RuntimeApiRequest},
+	ActivatedLeaf,
 };
 use polkadot_node_subsystem_test_helpers::{subsystem_test_harness, TestSubsystemContextHandle};
 use polkadot_node_subsystem_util::TimeoutExt;
 use polkadot_primitives::{
-	CollatorPair, HeadData, Id as ParaId, PersistedValidationData, ScheduledCore, ValidationCode,
+	AsyncBackingParams, CollatorPair, HeadData, Id as ParaId, Id, PersistedValidationData,
+	ScheduledCore, ValidationCode,
 };
+use rstest::rstest;
 use sp_keyring::sr25519::Keyring as Sr25519Keyring;
 use std::pin::Pin;
-use test_helpers::{dummy_hash, dummy_head_data, dummy_validator};
+use test_helpers::{dummy_candidate_descriptor, dummy_hash, dummy_head_data, dummy_validator};

 type VirtualOverseer = TestSubsystemContextHandle<CollationGenerationMessage>;

@@ -132,8 +135,10 @@ fn scheduled_core_for<Id: Into<ParaId>>(para_id: Id) -> ScheduledCore {
 	ScheduledCore { para_id: para_id.into(), collator: None }
 }

-#[test]
-fn requests_availability_per_relay_parent() {
+#[rstest]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT - 1)]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT)]
+fn requests_availability_per_relay_parent(#[case] runtime_version: u32) {
 	let activated_hashes: Vec<Hash> =
 		vec![[1; 32].into(), [4; 32].into(), [9; 32].into(), [16; 32].into()];

@@ -159,6 +164,18 @@ fn requests_availability_per_relay_parent() {
 				))) => {
 					tx.send(Err(RuntimeApiError::NotSupported { runtime_api_name: "doesnt_matter" })).unwrap();
 				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::Version(tx),
+				))) => {
+					tx.send(Ok(runtime_version)).unwrap();
+				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::ClaimQueue(tx),
+				))) if runtime_version >= RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT => {
+					tx.send(Ok(BTreeMap::new())).unwrap();
+				},
 				Some(msg) => panic!("didn't expect any other overseer requests given no availability cores; got {:?}", msg),
 			}
 		}
@@ -184,8 +201,10 @@ fn requests_availability_per_relay_parent() {
 	assert_eq!(requested_availability_cores, activated_hashes);
 }

-#[test]
-fn requests_validation_data_for_scheduled_matches() {
+#[rstest]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT - 1)]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT)]
+fn requests_validation_data_for_scheduled_matches(#[case] runtime_version: u32) {
 	let activated_hashes: Vec<Hash> = vec![
 		Hash::repeat_byte(1),
 		Hash::repeat_byte(4),
@@ -242,6 +261,18 @@ fn requests_validation_data_for_scheduled_matches() {
 					}))
 					.unwrap();
 				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::Version(tx),
+				))) => {
+					tx.send(Ok(runtime_version)).unwrap();
+				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::ClaimQueue(tx),
+				))) if runtime_version >= RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT => {
+					tx.send(Ok(BTreeMap::new())).unwrap();
+				},
 				Some(msg) => {
 					panic!("didn't expect any other overseer requests; got {:?}", msg)
 				},
@@ -271,8 +302,10 @@ fn requests_validation_data_for_scheduled_matches() {
 	assert_eq!(requested_validation_data, vec![[4; 32].into()]);
 }

-#[test]
-fn sends_distribute_collation_message() {
+#[rstest]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT - 1)]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT)]
+fn sends_distribute_collation_message(#[case] runtime_version: u32) {
 	let activated_hashes: Vec<Hash> = vec![
 		Hash::repeat_byte(1),
 		Hash::repeat_byte(4),
@@ -339,6 +372,18 @@ fn sends_distribute_collation_message() {
 					}))
 					.unwrap();
 				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::Version(tx),
+				))) => {
+					tx.send(Ok(runtime_version)).unwrap();
+				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::ClaimQueue(tx),
+				))) if runtime_version >= RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT => {
+					tx.send(Ok(BTreeMap::new())).unwrap();
+				},
 				Some(msg @ AllMessages::CollatorProtocol(_)) => {
 					inner_to_collator_protocol.lock().await.push(msg);
 				},
@@ -423,8 +468,10 @@ fn sends_distribute_collation_message() {
 	}
 }

-#[test]
-fn fallback_when_no_validation_code_hash_api() {
+#[rstest]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT - 1)]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT)]
+fn fallback_when_no_validation_code_hash_api(#[case] runtime_version: u32) {
 	// This is a variant of the above test, but with the validation code hash API disabled.

 	let activated_hashes: Vec<Hash> = vec![
@@ -501,9 +548,22 @@ fn fallback_when_no_validation_code_hash_api() {
 					}))
 					.unwrap();
 				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::Version(tx),
+				))) => {
+					tx.send(Ok(runtime_version)).unwrap();
+				},
 				Some(msg @ AllMessages::CollatorProtocol(_)) => {
 					inner_to_collator_protocol.lock().await.push(msg);
 				},
+				Some(AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+					_hash,
+					RuntimeApiRequest::ClaimQueue(tx),
+				))) if runtime_version >= RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT => {
+					let res = BTreeMap::<CoreIndex, VecDeque<ParaId>>::new();
+					tx.send(Ok(res)).unwrap();
+				},
 				Some(msg) => {
 					panic!("didn't expect any other overseer requests; got {:?}", msg)
 				},
@@ -635,3 +695,252 @@ fn submit_collation_leads_to_distribution() {
 		virtual_overseer
 	});
 }
+
+// There is one core in `Occupied` state and async backing is enabled. On new head activation
+// `CollationGeneration` should produce and distribute a new collation.
+#[rstest]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT - 1)]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT)]
+fn distribute_collation_for_occupied_core_with_async_backing_enabled(#[case] runtime_version: u32) {
+	let activated_hash: Hash = [1; 32].into();
+	let para_id = ParaId::from(5);
+
+	// One core, in occupied state. The data in `CoreState` and `ClaimQueue` should match.
+	let cores: Vec<CoreState> = vec![CoreState::Occupied(polkadot_primitives::OccupiedCore {
+		next_up_on_available: Some(ScheduledCore { para_id, collator: None }),
+		occupied_since: 1,
+		time_out_at: 10,
+		next_up_on_time_out: Some(ScheduledCore { para_id, collator: None }),
+		availability: Default::default(), // doesn't matter
+		group_responsible: polkadot_primitives::GroupIndex(0),
+		candidate_hash: Default::default(),
+		candidate_descriptor: dummy_candidate_descriptor(dummy_hash()),
+	})];
+	let claim_queue = BTreeMap::from([(CoreIndex::from(0), VecDeque::from([para_id]))]);
+
+	test_harness(|mut virtual_overseer| async move {
+		helpers::initialize_collator(&mut virtual_overseer, para_id).await;
+		helpers::activate_new_head(&mut virtual_overseer, activated_hash).await;
+		helpers::handle_runtime_calls_on_new_head_activation(
+			&mut virtual_overseer,
+			activated_hash,
+			AsyncBackingParams { max_candidate_depth: 1, allowed_ancestry_len: 1 },
+			cores,
+			runtime_version,
+			claim_queue,
+		)
+		.await;
+		helpers::handle_core_processing_for_a_leaf(
+			&mut virtual_overseer,
+			activated_hash,
+			para_id,
+			// `CoreState` is `Occupied` => `OccupiedCoreAssumption` is `Included`
+			OccupiedCoreAssumption::Included,
+		)
+		.await;
+
+		virtual_overseer
+	});
+}
+
+// There is one core in `Occupied` state and async backing is disabled. On new head activation
+// no new collation should be generated.
+#[rstest]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT - 1)]
+#[case(RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT)]
+fn no_collation_is_distributed_for_occupied_core_with_async_backing_disabled(
+	#[case] runtime_version: u32,
+) {
+	let activated_hash: Hash = [1; 32].into();
+	let para_id = ParaId::from(5);
+
+	// One core, in occupied state. The data in `CoreState` and `ClaimQueue` should match.
+	let cores: Vec<CoreState> = vec![CoreState::Occupied(polkadot_primitives::OccupiedCore {
+		next_up_on_available: Some(ScheduledCore { para_id, collator: None }),
+		occupied_since: 1,
+		time_out_at: 10,
+		next_up_on_time_out: Some(ScheduledCore { para_id, collator: None }),
+		availability: Default::default(), // doesn't matter
+		group_responsible: polkadot_primitives::GroupIndex(0),
+		candidate_hash: Default::default(),
+		candidate_descriptor: dummy_candidate_descriptor(dummy_hash()),
+	})];
+	let claim_queue = BTreeMap::from([(CoreIndex::from(0), VecDeque::from([para_id]))]);
+
+	test_harness(|mut virtual_overseer| async move {
+		helpers::initialize_collator(&mut virtual_overseer, para_id).await;
+		helpers::activate_new_head(&mut virtual_overseer, activated_hash).await;
+		helpers::handle_runtime_calls_on_new_head_activation(
+			&mut virtual_overseer,
+			activated_hash,
+			AsyncBackingParams { max_candidate_depth: 0, allowed_ancestry_len: 0 },
+			cores,
+			runtime_version,
+			claim_queue,
+		)
+		.await;
+
+		virtual_overseer
+	});
+}
+
+mod helpers {
+	use super::*;
+
+	// Sends `Initialize` with a collator config
+	pub async fn initialize_collator(virtual_overseer: &mut VirtualOverseer, para_id: ParaId) {
+		virtual_overseer
+			.send(FromOrchestra::Communication {
+				msg: CollationGenerationMessage::Initialize(test_config(para_id)),
+			})
+			.await;
+	}
+
+	// Sends `ActiveLeaves` for a single leaf with the specified hash. Block number is hardcoded.
+	pub async fn activate_new_head(virtual_overseer: &mut VirtualOverseer, activated_hash: Hash) {
+		virtual_overseer
+			.send(FromOrchestra::Signal(OverseerSignal::ActiveLeaves(ActiveLeavesUpdate {
+				activated: Some(ActivatedLeaf {
+					hash: activated_hash,
+					number: 10,
+					unpin_handle: polkadot_node_subsystem_test_helpers::mock::dummy_unpin_handle(
+						activated_hash,
+					),
+					span: Arc::new(overseer::jaeger::Span::Disabled),
+				}),
+				..Default::default()
+			})))
+			.await;
+	}
+
+	// Handle all runtime calls performed in `handle_new_activations`. Conditionally expects a
+	// `CLAIM_QUEUE_RUNTIME_REQUIREMENT` call if the passed `runtime_version` is greater or equal to
+	// `CLAIM_QUEUE_RUNTIME_REQUIREMENT`
+	pub async fn handle_runtime_calls_on_new_head_activation(
+		virtual_overseer: &mut VirtualOverseer,
+		activated_hash: Hash,
+		async_backing_params: AsyncBackingParams,
+		cores: Vec<CoreState>,
+		runtime_version: u32,
+		claim_queue: BTreeMap<CoreIndex, VecDeque<Id>>,
+	) {
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(hash, RuntimeApiRequest::AvailabilityCores(tx))) => {
+				assert_eq!(hash, activated_hash);
+				let _ = tx.send(Ok(cores));
+			}
+		);
+
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(hash, RuntimeApiRequest::Validators(tx))) => {
+				assert_eq!(hash, activated_hash);
+				let _ = tx.send(Ok(vec![
+					Sr25519Keyring::Alice.public().into(),
+					Sr25519Keyring::Bob.public().into(),
+					Sr25519Keyring::Charlie.public().into(),
+				]));
+			}
+		);
+
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+								hash,
+								RuntimeApiRequest::AsyncBackingParams(
+									tx,
+								),
+							)) => {
+				assert_eq!(hash, activated_hash);
+				let _ = tx.send(Ok(async_backing_params));
+			}
+		);
+
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+								hash,
+								RuntimeApiRequest::Version(tx),
+							)) => {
+				assert_eq!(hash, activated_hash);
+				let _ = tx.send(Ok(runtime_version));
+			}
+		);
+
+		if runtime_version == RuntimeApiRequest::CLAIM_QUEUE_RUNTIME_REQUIREMENT {
+			assert_matches!(
+				overseer_recv(virtual_overseer).await,
+				AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+									hash,
+									RuntimeApiRequest::ClaimQueue(tx),
+								)) => {
+					assert_eq!(hash, activated_hash);
+					let _ = tx.send(Ok(claim_queue));
+				}
+			);
+		}
+	}
+
+	// Handles all runtime requests performed in `handle_new_activations` for the case when a
+	// collation should be prepared for the new leaf
+	pub async fn handle_core_processing_for_a_leaf(
+		virtual_overseer: &mut VirtualOverseer,
+		activated_hash: Hash,
+		para_id: ParaId,
+		expected_occupied_core_assumption: OccupiedCoreAssumption,
+	) {
+		// Some hardcoded data - if needed, extract to parameters
+		let validation_code_hash = ValidationCodeHash::from(Hash::repeat_byte(42));
+		let parent_head = HeadData::from(vec![1, 2, 3]);
+		let pvd = PersistedValidationData {
+			parent_head: parent_head.clone(),
+			relay_parent_number: 10,
+			relay_parent_storage_root: Hash::repeat_byte(1),
+			max_pov_size: 1024,
+		};
+
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(hash, RuntimeApiRequest::PersistedValidationData(id, a, tx))) => {
+				assert_eq!(hash, activated_hash);
+				assert_eq!(id, para_id);
+				assert_eq!(a, expected_occupied_core_assumption);
+
+				let _ = tx.send(Ok(Some(pvd.clone())));
+			}
+		);
+
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::RuntimeApi(RuntimeApiMessage::Request(
+				hash,
+				RuntimeApiRequest::ValidationCodeHash(
+					id,
+					assumption,
+					tx,
+				),
+			)) => {
+				assert_eq!(hash, activated_hash);
+				assert_eq!(id, para_id);
+				assert_eq!(assumption, expected_occupied_core_assumption);
+
+				let _ = tx.send(Ok(Some(validation_code_hash)));
+			}
+		);
+
+		assert_matches!(
+			overseer_recv(virtual_overseer).await,
+			AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation{
+				candidate_receipt,
+				parent_head_data_hash,
+				..
+			}) => {
+				assert_eq!(parent_head_data_hash, parent_head.hash());
+				assert_eq!(candidate_receipt.descriptor().persisted_validation_data_hash, pvd.hash());
+				assert_eq!(candidate_receipt.descriptor().para_head, dummy_head_data().hash());
+				assert_eq!(candidate_receipt.descriptor().validation_code_hash, validation_code_hash);
+			}
+		);
+	}
+}