collator protocol changes for elastic scaling (validator side)

cumulus/client/collator/src/lib.rs

@@ -220,6 +220,7 @@ pub mod relay_chain_driven {
 					this_rx.await.ok().flatten()
 				})
 			})),
+			with_elastic_scaling: false,
 		};
 
 		overseer_handle
@@ -243,8 +244,9 @@ pub async fn initialize_collator_subsystems(
 	key: CollatorPair,
 	para_id: ParaId,
 	reinitialize: bool,
+	with_elastic_scaling: bool,
 ) {
-	let config = CollationGenerationConfig { key, para_id, collator: None };
+	let config = CollationGenerationConfig { key, para_id, collator: None, with_elastic_scaling };
 
 	if reinitialize {
 		overseer_handle

cumulus/client/consensus/aura/src/collators/lookahead.rs

@@ -107,6 +107,9 @@ pub struct Params<BI, CIDP, Client, Backend, RClient, CHP, SO, Proposer, CS> {
 	pub authoring_duration: Duration,
 	/// Whether we should reinitialize the collator config (i.e. we are transitioning to aura).
 	pub reinitialize: bool,
+	/// Whether elastic scaling is enabled for this collation.
+	/// If it is, the collator will send the parent-head data along with the collation.
+	pub with_elastic_scaling: bool,
 }
 
 /// Run async-backing-friendly Aura.
@@ -152,6 +155,7 @@ where
 			params.collator_key,
 			params.para_id,
 			params.reinitialize,
+			params.with_elastic_scaling,
 		)
 		.await;
 

cumulus/polkadot-parachain/src/service.rs

@@ -977,6 +977,7 @@ pub async fn start_rococo_parachain_node(
 				collator_service,
 				authoring_duration: Duration::from_millis(1500),
 				reinitialize: false,
+				with_elastic_scaling: false,
 			};
 
 			let fut = aura::run::<
@@ -1473,6 +1474,7 @@ where
 				collator_service,
 				authoring_duration: Duration::from_millis(1500),
 				reinitialize: false,
+				with_elastic_scaling: false,
 			};
 
 			let fut =
@@ -1768,6 +1770,7 @@ where
 					authoring_duration: Duration::from_millis(1500),
 					reinitialize: true, /* we need to always re-initialize for asset-hub moving
 					                     * to aura */
+					with_elastic_scaling: false,
 				};
 
 				aura::run::<Block, <AuraId as AppCrypto>::Pair, _, _, _, _, _, _, _, _, _>(params)
@@ -1870,6 +1873,7 @@ where
 				collator_service,
 				authoring_duration: Duration::from_millis(1500),
 				reinitialize: false,
+				with_elastic_scaling: false,
 			};
 
 			let fut =
@@ -2180,6 +2184,7 @@ pub async fn start_contracts_rococo_node(
 				// Very limited proposal time.
 				authoring_duration: Duration::from_millis(1500),
 				reinitialize: false,
+				with_elastic_scaling: false,
 			};
 
 			let fut = aura::run::<

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

@@ -350,6 +350,8 @@ async fn handle_new_activations<Context>(
 							},
 						};
 
+					let with_elastic_scaling = task_config.with_elastic_scaling;
+
 					construct_and_distribute_receipt(
 						PreparedCollation {
 							collation,
@@ -358,6 +360,7 @@ async fn handle_new_activations<Context>(
 							validation_data,
 							validation_code_hash,
 							n_validators,
+							with_elastic_scaling,
 						},
 						task_config.key.clone(),
 						&mut task_sender,
@@ -392,6 +395,7 @@ async fn handle_submit_collation<Context>(
 
 	let validators = request_validators(relay_parent, ctx.sender()).await.await??;
 	let n_validators = validators.len();
+	let with_elastic_scaling = config.with_elastic_scaling;
 
 	// We need to swap the parent-head data, but all other fields here will be correct.
 	let mut validation_data = match request_persisted_validation_data(
@@ -424,6 +428,7 @@ async fn handle_submit_collation<Context>(
 		validation_data,
 		validation_code_hash,
 		n_validators,
+		with_elastic_scaling,
 	};
 
 	construct_and_distribute_receipt(
@@ -445,6 +450,7 @@ struct PreparedCollation {
 	validation_data: PersistedValidationData,
 	validation_code_hash: ValidationCodeHash,
 	n_validators: usize,
+	with_elastic_scaling: bool,
 }
 
 /// Takes a prepared collation, along with its context, and produces a candidate receipt
@@ -463,6 +469,7 @@ async fn construct_and_distribute_receipt(
 		validation_data,
 		validation_code_hash,
 		n_validators,
+		with_elastic_scaling,
 	} = collation;
 
 	let persisted_validation_data_hash = validation_data.hash();
@@ -540,23 +547,28 @@ async fn construct_and_distribute_receipt(
 		},
 	};
 
+	let maybe_parent_head_data =
+		if with_elastic_scaling { Some(commitments.head_data.clone()) } else { None };
+
 	gum::debug!(
 		target: LOG_TARGET,
 		candidate_hash = ?ccr.hash(),
 		?pov_hash,
 		?relay_parent,
 		para_id = %para_id,
+		?with_elastic_scaling,
 		"candidate is generated",
 	);
 	metrics.on_collation_generated();
 
 	sender
-		.send_message(CollatorProtocolMessage::DistributeCollation(
-			ccr,
+		.send_message(CollatorProtocolMessage::DistributeCollation {
+			candidate_receipt: ccr,
 			parent_head_data_hash,
 			pov,
+			maybe_parent_head_data,
 			result_sender,
-		))
+		})
 		.await;
 }
 

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

@@ -117,6 +117,7 @@ fn test_config<Id: Into<ParaId>>(para_id: Id) -> CollationGenerationConfig {
 		key: CollatorPair::generate().0,
 		collator: Some(Box::new(|_: Hash, _vd: &PersistedValidationData| TestCollator.boxed())),
 		para_id: para_id.into(),
+		with_elastic_scaling: false,
 	}
 }
 
@@ -125,6 +126,7 @@ fn test_config_no_collator<Id: Into<ParaId>>(para_id: Id) -> CollationGeneration
 		key: CollatorPair::generate().0,
 		collator: None,
 		para_id: para_id.into(),
+		with_elastic_scaling: false,
 	}
 }
 
@@ -390,11 +392,11 @@ fn sends_distribute_collation_message() {
 
 	assert_eq!(to_collator_protocol.len(), 1);
 	match AllMessages::from(to_collator_protocol.pop().unwrap()) {
-		AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation(
-			CandidateReceipt { descriptor, .. },
-			_pov,
-			..,
-		)) => {
+		AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation {
+			candidate_receipt,
+			..
+		}) => {
+			let CandidateReceipt { descriptor, .. } = candidate_receipt;
 			// signature generation is non-deterministic, so we can't just assert that the
 			// expected descriptor is correct. What we can do is validate that the produced
 			// descriptor has a valid signature, then just copy in the generated signature
@@ -529,11 +531,11 @@ fn fallback_when_no_validation_code_hash_api() {
 
 	assert_eq!(to_collator_protocol.len(), 1);
 	match &to_collator_protocol[0] {
-		AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation(
-			CandidateReceipt { descriptor, .. },
-			_pov,
-			..,
-		)) => {
+		AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation {
+			candidate_receipt,
+			..
+		}) => {
+			let CandidateReceipt { descriptor, .. } = candidate_receipt;
 			assert_eq!(expect_validation_code_hash, descriptor.validation_code_hash);
 		},
 		_ => panic!("received wrong message type"),
@@ -619,15 +621,16 @@ fn submit_collation_leads_to_distribution() {
 
 		assert_matches!(
 			overseer_recv(&mut virtual_overseer).await,
-			AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation(
-				ccr,
+			AllMessages::CollatorProtocol(CollatorProtocolMessage::DistributeCollation {
+				candidate_receipt,
 				parent_head_data_hash,
 				..
-			)) => {
+			}) => {
+				let CandidateReceipt { descriptor, .. } = candidate_receipt;
 				assert_eq!(parent_head_data_hash, parent_head.hash());
-				assert_eq!(ccr.descriptor().persisted_validation_data_hash, expected_pvd.hash());
-				assert_eq!(ccr.descriptor().para_head, dummy_head_data().hash());
-				assert_eq!(ccr.descriptor().validation_code_hash, validation_code_hash);
+				assert_eq!(descriptor.persisted_validation_data_hash, expected_pvd.hash());
+				assert_eq!(descriptor.para_head, dummy_head_data().hash());
+				assert_eq!(descriptor.validation_code_hash, validation_code_hash);
 			}
 		);
 

polkadot/node/core/prospective-parachains/src/lib.rs

@@ -771,8 +771,9 @@ fn answer_prospective_validation_data_request(
 		Some(s) => s,
 	};
 
-	let mut head_data =
-		storage.head_data_by_hash(&request.parent_head_data_hash).map(|x| x.clone());
+	let mut head_data = request
+		.maybe_parent_head_data
+		.or_else(|| storage.head_data_by_hash(&request.parent_head_data_hash).map(|x| x.clone()));
 	let mut relay_parent_info = None;
 	let mut max_pov_size = None;
 

polkadot/node/core/prospective-parachains/src/tests.rs

@@ -473,6 +473,7 @@ async fn get_pvd(
 		para_id,
 		candidate_relay_parent,
 		parent_head_data_hash: parent_head_data.hash(),
+		maybe_parent_head_data: None,
 	};
 	let (tx, rx) = oneshot::channel();
 	virtual_overseer

polkadot/node/network/collator-protocol/src/collator_side/collation.rs

@@ -27,7 +27,7 @@ use polkadot_node_network_protocol::{
 	PeerId,
 };
 use polkadot_node_primitives::PoV;
-use polkadot_primitives::{CandidateHash, CandidateReceipt, Hash, Id as ParaId};
+use polkadot_primitives::{CandidateHash, CandidateReceipt, Hash, HeadData, Id as ParaId};
 
 /// The status of a collation as seen from the collator.
 pub enum CollationStatus {
@@ -63,6 +63,8 @@ pub struct Collation {
 	pub parent_head_data_hash: Hash,
 	/// Proof to verify the state transition of the parachain.
 	pub pov: PoV,
+	/// Optional parent head-data needed for elastic scaling.
+	pub maybe_parent_head_data: Option<HeadData>,
 	/// Collation status.
 	pub status: CollationStatus,
 }