From 21b32849db4ac5f284d7ca2f93927b0d35bb783a Mon Sep 17 00:00:00 2001
From: Sebastian Kunert <skunert49@gmail.com>
Date: Thu, 19 Oct 2023 10:07:29 +0200
Subject: [PATCH] Cumulus: Allow aura to use initialized collation request
 receiver (#1911)

When launching our [small
network](https://github.com/paritytech/polkadot-sdk/blob/master/cumulus/zombienet/examples/small_network.toml)
for testing the node was crashing here shortly after launch:

https://github.com/paritytech/polkadot-sdk/blob/5cdd819ed295645958afd9d937d989978fd0c84e/polkadot/node/collation-generation/src/lib.rs#L140

After changes in #1788 for the asset hub collator we are waiting for
blocks of the shell runtime to pass before we initialize aura. However,
this means that we attempted to initialize the collation related relay
chain subsystems twice, leading to the error.

I modified Aura to let it optionally take an already initialized stream
of collation requests.
---
 .../consensus/aura/src/collators/basic.rs     | 26 +++++++++++++------
 .../parachain-template/node/src/service.rs    |  1 +
 cumulus/polkadot-parachain/src/service.rs     |  4 +++
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/cumulus/client/consensus/aura/src/collators/basic.rs b/cumulus/client/consensus/aura/src/collators/basic.rs
index 3c904915cea..dc0078b0d6a 100644
--- a/cumulus/client/consensus/aura/src/collators/basic.rs
+++ b/cumulus/client/consensus/aura/src/collators/basic.rs
@@ -23,7 +23,9 @@
 //! For more information about AuRa, the Substrate crate should be checked.
 
 use codec::{Codec, Decode};
-use cumulus_client_collator::service::ServiceInterface as CollatorServiceInterface;
+use cumulus_client_collator::{
+	relay_chain_driven::CollationRequest, service::ServiceInterface as CollatorServiceInterface,
+};
 use cumulus_client_consensus_common::ParachainBlockImportMarker;
 use cumulus_client_consensus_proposer::ProposerInterface;
 use cumulus_primitives_core::{relay_chain::BlockId as RBlockId, CollectCollationInfo};
@@ -33,7 +35,7 @@ use polkadot_node_primitives::CollationResult;
 use polkadot_overseer::Handle as OverseerHandle;
 use polkadot_primitives::{CollatorPair, Id as ParaId};
 
-use futures::prelude::*;
+use futures::{channel::mpsc::Receiver, prelude::*};
 use sc_client_api::{backend::AuxStore, BlockBackend, BlockOf};
 use sc_consensus::BlockImport;
 use sp_api::ProvideRuntimeApi;
@@ -81,6 +83,10 @@ pub struct Params<BI, CIDP, Client, RClient, SO, Proposer, CS> {
 	pub collator_service: CS,
 	/// The amount of time to spend authoring each block.
 	pub authoring_duration: Duration,
+	/// Receiver for collation requests. If `None`, Aura consensus will establish a new receiver.
+	/// Should be used when a chain migrates from a different consensus algorithm and was already
+	/// processing collation requests before initializing Aura.
+	pub collation_request_receiver: Option<Receiver<CollationRequest>>,
 }
 
 /// Run bare Aura consensus as a relay-chain-driven collator.
@@ -110,12 +116,16 @@ where
 	P::Signature: TryFrom<Vec<u8>> + Member + Codec,
 {
 	async move {
-		let mut collation_requests = cumulus_client_collator::relay_chain_driven::init(
-			params.collator_key,
-			params.para_id,
-			params.overseer_handle,
-		)
-		.await;
+		let mut collation_requests = match params.collation_request_receiver {
+			Some(receiver) => receiver,
+			None =>
+				cumulus_client_collator::relay_chain_driven::init(
+					params.collator_key,
+					params.para_id,
+					params.overseer_handle,
+				)
+				.await,
+		};
 
 		let mut collator = {
 			let params = collator_util::Params {
diff --git a/cumulus/parachain-template/node/src/service.rs b/cumulus/parachain-template/node/src/service.rs
index 9fa6d60c2e7..84dcd6dd1b3 100644
--- a/cumulus/parachain-template/node/src/service.rs
+++ b/cumulus/parachain-template/node/src/service.rs
@@ -411,6 +411,7 @@ fn start_consensus(
 		collator_service,
 		// Very limited proposal time.
 		authoring_duration: Duration::from_millis(500),
+		collation_request_receiver: None,
 	};
 
 	let fut =
diff --git a/cumulus/polkadot-parachain/src/service.rs b/cumulus/polkadot-parachain/src/service.rs
index fa61f534784..438d09a4c77 100644
--- a/cumulus/polkadot-parachain/src/service.rs
+++ b/cumulus/polkadot-parachain/src/service.rs
@@ -988,6 +988,7 @@ pub async fn start_rococo_parachain_node(
 				collator_service,
 				// Very limited proposal time.
 				authoring_duration: Duration::from_millis(500),
+				collation_request_receiver: None,
 			};
 
 			let fut = basic_aura::run::<
@@ -1380,6 +1381,7 @@ where
 				collator_service,
 				// Very limited proposal time.
 				authoring_duration: Duration::from_millis(500),
+				collation_request_receiver: None,
 			};
 
 			let fut =
@@ -1520,6 +1522,7 @@ where
 					collator_service,
 					// Very limited proposal time.
 					authoring_duration: Duration::from_millis(500),
+					collation_request_receiver: Some(request_stream),
 				};
 
 				basic_aura::run::<Block, <AuraId as AppCrypto>::Pair, _, _, _, _, _, _, _>(params)
@@ -1925,6 +1928,7 @@ pub async fn start_contracts_rococo_node(
 				collator_service,
 				// Very limited proposal time.
 				authoring_duration: Duration::from_millis(500),
+				collation_request_receiver: None,
 			};
 
 			let fut = basic_aura::run::<
-- 
GitLab