Fix Asset Hub collator crashing when starting from genesis

cumulus/client/collator/src/lib.rs

@@ -69,14 +69,14 @@ where
 	RA::Api: CollectCollationInfo<Block>,
 {
 	/// Create a new instance.
-	fn new(
+	pub fn new(
 		collator_service: CollatorService<Block, BS, RA>,
 		parachain_consensus: Box<dyn ParachainConsensus<Block>>,
 	) -> Self {
 		Self { service: collator_service, parachain_consensus }
 	}
 
-	async fn produce_candidate(
+	pub async fn produce_candidate(
 		mut self,
 		relay_parent: PHash,
 		validation_data: PersistedValidationData,
@@ -142,6 +142,27 @@ where
 
 		Some(CollationResult { collation, result_sender: Some(result_sender) })
 	}
+
+	pub fn header_hash(&self, pvd: PersistedValidationData) -> Option<Block::Hash> {
+		let header = match Block::Header::decode(&mut &pvd.parent_head.0[..]) {
+			Ok(x) => x,
+			Err(e) => {
+				tracing::error!(
+					target: LOG_TARGET,
+					error = ?e,
+					"Could not decode the head data."
+				);
+				return None
+			},
+		};
+
+		let header_hash = header.hash();
+		if !self.service.check_block_status(header_hash, &header) {
+			return None
+		}
+
+		Some(header_hash)
+	}
 }
 
 /// Relay-chain-driven collators are those whose block production is driven purely

cumulus/polkadot-parachain/src/command.rs

@@ -836,21 +836,21 @@ pub fn run() -> Result<()> {
 				info!("Is collating: {}", if config.role.is_authority() { "yes" } else { "no" });
 
 				match config.chain_spec.runtime() {
-					Runtime::AssetHubPolkadot => crate::service::start_generic_aura_node::<
+					Runtime::AssetHubPolkadot => crate::service::start_asset_hub_node::<
 						asset_hub_polkadot_runtime::RuntimeApi,
 						AssetHubPolkadotAuraId,
 					>(config, polkadot_config, collator_options, id, hwbench)
 					.await
 					.map(|r| r.0)
 					.map_err(Into::into),
-					Runtime::AssetHubKusama => crate::service::start_generic_aura_node::<
+					Runtime::AssetHubKusama => crate::service::start_asset_hub_node::<
 						asset_hub_kusama_runtime::RuntimeApi,
 						AuraId,
 					>(config, polkadot_config, collator_options, id, hwbench)
 					.await
 					.map(|r| r.0)
 					.map_err(Into::into),
-					Runtime::AssetHubWestend => crate::service::start_generic_aura_node::<
+					Runtime::AssetHubWestend => crate::service::start_asset_hub_node::<
 						asset_hub_westend_runtime::RuntimeApi,
 						AuraId,
 					>(config, polkadot_config, collator_options, id, hwbench)

cumulus/polkadot-parachain/src/service.rs

@@ -43,8 +43,9 @@ use jsonrpsee::RpcModule;
 use crate::rpc;
 pub use parachains_common::{AccountId, Balance, Block, BlockNumber, Hash, Header, Nonce};
 
+use cumulus_client_collator::Collator;
 use cumulus_client_consensus_relay_chain::Verifier as RelayChainVerifier;
-use futures::lock::Mutex;
+use futures::{lock::Mutex, prelude::*};
 use sc_consensus::{
 	import_queue::{BasicQueue, Verifier as VerifierT},
 	BlockImportParams, ImportQueue,
@@ -54,7 +55,7 @@ use sc_network::{config::FullNetworkConfiguration, NetworkBlock};
 use sc_network_sync::SyncingService;
 use sc_service::{Configuration, PartialComponents, TFullBackend, TFullClient, TaskManager};
 use sc_telemetry::{Telemetry, TelemetryHandle, TelemetryWorker, TelemetryWorkerHandle};
-use sp_api::{ApiExt, ConstructRuntimeApi};
+use sp_api::{ApiExt, ConstructRuntimeApi, ProvideRuntimeApi};
 use sp_consensus_aura::AuraApi;
 use sp_keystore::KeystorePtr;
 use sp_runtime::{app_crypto::AppCrypto, traits::Header as HeaderT};
@@ -1389,6 +1390,188 @@ where
 	.await
 }
 
+/// Start a shell node which should later transition into an Aura powered parachain node. Asset Hub
+/// uses this because at genesis, Asset Hub was on the `shell` runtime which didn't have Aura and
+/// needs to sync and upgrade before it can run `AuraApi` functions.
+pub async fn start_asset_hub_node<RuntimeApi, AuraId: AppCrypto + Send + Codec + Sync>(
+	parachain_config: Configuration,
+	polkadot_config: Configuration,
+	collator_options: CollatorOptions,
+	para_id: ParaId,
+	hwbench: Option<sc_sysinfo::HwBench>,
+) -> sc_service::error::Result<(TaskManager, Arc<ParachainClient<RuntimeApi>>)>
+where
+	RuntimeApi: ConstructRuntimeApi<Block, ParachainClient<RuntimeApi>> + Send + Sync + 'static,
+	RuntimeApi::RuntimeApi: sp_transaction_pool::runtime_api::TaggedTransactionQueue<Block>
+		+ sp_api::Metadata<Block>
+		+ sp_session::SessionKeys<Block>
+		+ sp_api::ApiExt<Block>
+		+ sp_offchain::OffchainWorkerApi<Block>
+		+ sp_block_builder::BlockBuilder<Block>
+		+ cumulus_primitives_core::CollectCollationInfo<Block>
+		+ sp_consensus_aura::AuraApi<Block, <<AuraId as AppCrypto>::Pair as Pair>::Public>
+		+ pallet_transaction_payment_rpc::TransactionPaymentRuntimeApi<Block, Balance>
+		+ frame_rpc_system::AccountNonceApi<Block, AccountId, Nonce>,
+	<<AuraId as AppCrypto>::Pair as Pair>::Signature:
+		TryFrom<Vec<u8>> + std::hash::Hash + sp_runtime::traits::Member + Codec,
+{
+	start_node_impl::<RuntimeApi, _, _, _>(
+		parachain_config,
+		polkadot_config,
+		collator_options,
+		CollatorSybilResistance::Resistant, // Aura
+		para_id,
+		|_| Ok(RpcModule::new(())),
+		aura_build_import_queue::<_, AuraId>,
+		|client,
+		 block_import,
+		 prometheus_registry,
+		 telemetry,
+		 task_manager,
+		 relay_chain_interface,
+		 transaction_pool,
+		 sync_oracle,
+		 keystore,
+		 relay_chain_slot_duration,
+		 para_id,
+		 collator_key,
+		 overseer_handle,
+		 announce_block| {
+			let proposer_factory = sc_basic_authorship::ProposerFactory::with_proof_recording(
+				task_manager.spawn_handle(),
+				client.clone(),
+				transaction_pool.clone(),
+				prometheus_registry,
+				telemetry.clone(),
+			);
+			let relay_chain_interface2 = relay_chain_interface.clone();
+
+			// Build the free-for-all consensus necessary for the initial `shell` runtime.
+			let free_for_all = cumulus_client_consensus_relay_chain::build_relay_chain_consensus(
+				cumulus_client_consensus_relay_chain::BuildRelayChainConsensusParams {
+					para_id,
+					proposer_factory,
+					block_import: block_import.clone(),
+					relay_chain_interface: relay_chain_interface.clone(),
+					create_inherent_data_providers: move |_, (relay_parent, validation_data)| {
+						let relay_chain_interface = relay_chain_interface.clone();
+						async move {
+							let parachain_inherent =
+							cumulus_primitives_parachain_inherent::ParachainInherentData::create_at(
+								relay_parent,
+								&relay_chain_interface,
+								&validation_data,
+								para_id,
+							).await;
+							let parachain_inherent = parachain_inherent.ok_or_else(|| {
+								Box::<dyn std::error::Error + Send + Sync>::from(
+									"Failed to create parachain inherent",
+								)
+							})?;
+							Ok(parachain_inherent)
+						}
+					},
+				},
+			);
+
+			let proposer_factory = sc_basic_authorship::ProposerFactory::with_proof_recording(
+				task_manager.spawn_handle(),
+				client.clone(),
+				transaction_pool,
+				prometheus_registry,
+				telemetry.clone(),
+			);
+			let proposer = Proposer::new(proposer_factory);
+
+			let spawner = task_manager.spawn_handle();
+			let collator_service = CollatorService::new(
+				client.clone(),
+				Arc::new(spawner),
+				announce_block.clone(),
+				client.clone(),
+			);
+
+			let collator = Collator::new(collator_service.clone(), free_for_all);
+
+			let collation_future = Box::pin(async move {
+				// Start collating with the `shell` runtime while waiting for an upgrade to an Aura
+				// compatible runtime.
+				let mut request_stream = cumulus_client_collator::relay_chain_driven::init(
+					collator_key.clone(),
+					para_id,
+					overseer_handle.clone(),
+				)
+				.await;
+				while let Some(request) = request_stream.next().await {
+					let collation = {
+						let last_head_hash = match collator
+							.header_hash(request.persisted_validation_data().clone())
+						{
+							Some(header_hash) => header_hash,
+							None => {
+								request.complete(None);
+								continue
+							},
+						};
+						// Check if we have upgraded to an Aura compatible runtime and transition if
+						// necessary.
+						if client
+							.runtime_api()
+							.has_api::<dyn AuraApi<Block, AuraId>>(last_head_hash)
+							.unwrap_or(false)
+						{
+							// Respond to this request before transitioning to Aura.
+							request.complete(None);
+							break
+						}
+
+						collator
+							.clone()
+							.produce_candidate(
+								*request.relay_parent(),
+								request.persisted_validation_data().clone(),
+							)
+							.await
+					};
+
+					request.complete(collation);
+				}
+
+				// Move to Aura consensus.
+				let slot_duration = cumulus_client_consensus_aura::slot_duration(&*client).unwrap();
+
+				let params = BasicAuraParams {
+					create_inherent_data_providers: move |_, ()| async move { Ok(()) },
+					block_import,
+					para_client: client,
+					relay_client: relay_chain_interface2,
+					sync_oracle,
+					keystore,
+					collator_key,
+					para_id,
+					overseer_handle,
+					slot_duration,
+					relay_chain_slot_duration,
+					proposer,
+					collator_service,
+					// Very limited proposal time.
+					authoring_duration: Duration::from_millis(500),
+				};
+
+				basic_aura::run::<Block, <AuraId as AppCrypto>::Pair, _, _, _, _, _, _, _>(params)
+					.await
+			});
+
+			let spawner = task_manager.spawn_handle();
+			spawner.spawn("cumulus-asset-hub-collator", None, collation_future);
+
+			Ok(())
+		},
+		hwbench,
+	)
+	.await
+}
+
 /// Start an aura powered parachain node which uses the lookahead collator to support async backing.
 /// This node is basic in the sense that its runtime api doesn't include common contents such as
 /// transaction payment. Used for aura glutton.