From 233b347a5849a2ece9e54c4df97cb64ebd366365 Mon Sep 17 00:00:00 2001
From: Sergei Shulepov <sergei@parity.io>
Date: Thu, 3 Dec 2020 23:02:19 +0100
Subject: [PATCH] Plumb polkadot client into the collator struct (#255)

* plumb polkadot_client into Collator

* plumb para_id into Collator

* promote retrieve_dmq_contents to a method

* remove the retrieve_dmq_contents closure
---
 cumulus/collator/src/lib.rs | 77 +++++++++++++++++++++----------------
 1 file changed, 44 insertions(+), 33 deletions(-)

diff --git a/cumulus/collator/src/lib.rs b/cumulus/collator/src/lib.rs
index 729a5ea294..4d5ddd1f4d 100644
--- a/cumulus/collator/src/lib.rs
+++ b/cumulus/collator/src/lib.rs
@@ -60,20 +60,22 @@ type TransactionFor<E, Block> =
 	<<E as Environment<Block>>::Proposer as Proposer<Block>>::Transaction;
 
 /// The implementation of the Cumulus `Collator`.
-pub struct Collator<Block: BlockT, PF, BI, BS, Backend> {
+pub struct Collator<Block: BlockT, PF, BI, BS, Backend, PBackend, PClient> {
+	para_id: ParaId,
 	proposer_factory: Arc<Mutex<PF>>,
-	_phantom: PhantomData<Block>,
+	_phantom: PhantomData<(Block, PBackend)>,
 	inherent_data_providers: InherentDataProviders,
 	block_import: Arc<Mutex<BI>>,
 	block_status: Arc<BS>,
 	wait_to_announce: Arc<Mutex<WaitToAnnounce<Block>>>,
 	backend: Arc<Backend>,
-	retrieve_dmq_contents: Arc<dyn Fn(PHash) -> Option<DownwardMessagesType> + Send + Sync>,
+	polkadot_client: Arc<PClient>,
 }
 
-impl<Block: BlockT, PF, BI, BS, Backend> Clone for Collator<Block, PF, BI, BS, Backend> {
+impl<Block: BlockT, PF, BI, BS, Backend, PBackend, PClient> Clone for Collator<Block, PF, BI, BS, Backend, PBackend, PClient> {
 	fn clone(&self) -> Self {
 		Self {
+			para_id: self.para_id.clone(),
 			proposer_factory: self.proposer_factory.clone(),
 			inherent_data_providers: self.inherent_data_providers.clone(),
 			_phantom: PhantomData,
@@ -81,12 +83,12 @@ impl<Block: BlockT, PF, BI, BS, Backend> Clone for Collator<Block, PF, BI, BS, B
 			block_status: self.block_status.clone(),
 			wait_to_announce: self.wait_to_announce.clone(),
 			backend: self.backend.clone(),
-			retrieve_dmq_contents: self.retrieve_dmq_contents.clone(),
+			polkadot_client: self.polkadot_client.clone(),
 		}
 	}
 }
 
-impl<Block, PF, BI, BS, Backend> Collator<Block, PF, BI, BS, Backend>
+impl<Block, PF, BI, BS, Backend, PBackend, PApi, PClient> Collator<Block, PF, BI, BS, Backend, PBackend, PClient>
 where
 	Block: BlockT,
 	PF: Environment<Block> + 'static + Send,
@@ -100,9 +102,14 @@ where
 		+ 'static,
 	BS: BlockBackend<Block>,
 	Backend: sc_client_api::Backend<Block> + 'static,
+	PBackend: sc_client_api::Backend<PBlock> + 'static,
+	PBackend::State: StateBackend<BlakeTwo256>,
+	PApi: RuntimeApiCollection<StateBackend = PBackend::State>,
+	PClient: polkadot_service::AbstractClient<PBlock, PBackend, Api = PApi> + 'static,
 {
 	/// Create a new instance.
 	fn new(
+		para_id: ParaId,
 		proposer_factory: PF,
 		inherent_data_providers: InherentDataProviders,
 		overseer_handler: OverseerHandler,
@@ -111,7 +118,7 @@ where
 		spawner: Arc<dyn SpawnNamed + Send + Sync>,
 		announce_block: Arc<dyn Fn(Block::Hash, Vec<u8>) + Send + Sync>,
 		backend: Arc<Backend>,
-		retrieve_dmq_contents: Arc<dyn Fn(PHash) -> Option<DownwardMessagesType> + Send + Sync>,
+		polkadot_client: Arc<PClient>,
 	) -> Self {
 		let wait_to_announce = Arc::new(Mutex::new(WaitToAnnounce::new(
 			spawner,
@@ -120,6 +127,7 @@ where
 		)));
 
 		Self {
+			para_id,
 			proposer_factory: Arc::new(Mutex::new(proposer_factory)),
 			inherent_data_providers,
 			_phantom: PhantomData,
@@ -127,10 +135,33 @@ where
 			block_status,
 			wait_to_announce,
 			backend,
-			retrieve_dmq_contents,
+			polkadot_client,
 		}
 	}
 
+	/// Returns the whole contents of the downward message queue for the parachain we are collating
+	/// for.
+	///
+	/// Returns `None` in case of an error.
+	fn retrieve_dmq_contents(&self, relay_parent: PHash) -> Option<DownwardMessagesType> {
+		self
+			.polkadot_client
+			.runtime_api()
+			.dmq_contents_with_context(
+				&BlockId::hash(relay_parent),
+				sp_core::ExecutionContext::Importing,
+				self.para_id,
+			)
+			.map_err(|e| {
+				error!(
+					target: "cumulus-collator",
+					"An error occured during requesting the downward messages for {}: {:?}",
+					relay_parent, e,
+				);
+			})
+			.ok()
+	}
+
 	/// Get the inherent data with validation function parameters injected
 	fn inherent_data(
 		&mut self,
@@ -160,7 +191,7 @@ where
 			})
 			.ok()?;
 
-		let downward_messages = (self.retrieve_dmq_contents)(relay_parent)?;
+		let downward_messages = self.retrieve_dmq_contents(relay_parent)?;
 		inherent_data
 			.put_data(DOWNWARD_MESSAGES_IDENTIFIER, &downward_messages)
 			.map_err(|e| {
@@ -459,36 +490,15 @@ where
 	for<'a> &'a Client: BlockImport<Block>,
 	BS: BlockBackend<Block> + Send + Sync + 'static,
 	Spawner: SpawnNamed + Clone + Send + Sync + 'static,
-	PBackend: sc_client_api::Backend<PBlock>,
+	PBackend: sc_client_api::Backend<PBlock> + 'static,
 	PBackend::State: StateBackend<BlakeTwo256>,
 	PApi: RuntimeApiCollection<StateBackend = PBackend::State>,
 	PClient: polkadot_service::AbstractClient<PBlock, PBackend, Api = PApi> + 'static,
 {
-	let retrieve_dmq_contents = {
-		let polkadot_client = polkadot_client.clone();
-		move |relay_parent: PHash| {
-			polkadot_client
-				.runtime_api()
-				.dmq_contents_with_context(
-					&BlockId::hash(relay_parent),
-					sp_core::ExecutionContext::Importing,
-					para_id,
-				)
-				.map_err(|e| {
-					error!(
-						target: "cumulus-collator",
-						"An error occured during requesting the downward messages for {}: {:?}",
-						relay_parent, e,
-					);
-				})
-				.ok()
-		}
-	};
-
 	let follow = match cumulus_consensus::follow_polkadot(
 		para_id,
 		client,
-		polkadot_client,
+		polkadot_client.clone(),
 		announce_block.clone(),
 	) {
 		Ok(follow) => follow,
@@ -498,6 +508,7 @@ where
 	spawner.spawn("cumulus-follow-polkadot", follow.map(|_| ()).boxed());
 
 	let collator = Collator::new(
+		para_id,
 		proposer_factory,
 		inherent_data_providers,
 		overseer_handler.clone(),
@@ -506,7 +517,7 @@ where
 		Arc::new(spawner),
 		announce_block,
 		backend,
-		Arc::new(retrieve_dmq_contents),
+		polkadot_client,
 	);
 
 	let config = CollationGenerationConfig {