From 6e78f347fb9c53de83878c6d0ef4716a9bebdbe0 Mon Sep 17 00:00:00 2001
From: Sebastian Kunert <skunert49@gmail.com>
Date: Tue, 17 Oct 2023 11:45:33 +0200
Subject: [PATCH] Allow aura to use initialized collation request receiver

---
 .../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 3c904915ceafb..dc0078b0d6a98 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 9fa6d60c2e74a..84dcd6dd1b39d 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 fa61f534784e0..438d09a4c7771 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::<