Cancel previous task when new one is created

collator/src/lib.rs

@@ -16,7 +16,7 @@
 
 //! Cumulus Collator implementation for Substrate.
 
-use cumulus_network::wait_to_announce;
+use cumulus_network::WaitToAnnounce;
 use cumulus_runtime::ParachainBlockData;
 
 use sc_client::Client;
@@ -63,6 +63,7 @@ pub struct Collator<Block: BlockT, PF, BI> {
 	block_import: Arc<Mutex<BI>>,
 	spawner: Arc<dyn Spawn + Send + Sync>,
 	network: Arc<NetworkService<Block, <Block as BlockT>::Hash>>,
+	wait_to_announce: Arc<Mutex<WaitToAnnounce<Block>>>,
 }
 
 impl<Block: BlockT, PF, BI> Collator<Block, PF, BI> {
@@ -75,14 +76,22 @@ impl<Block: BlockT, PF, BI> Collator<Block, PF, BI> {
 		spawner: Arc<dyn Spawn + Send + Sync>,
 		network: Arc<NetworkService<Block, <Block as BlockT>::Hash>>,
 	) -> Self {
+		let collator_network = Arc::new(collator_network);
+		let wait_to_announce = Arc::new(Mutex::new(WaitToAnnounce::new(
+			spawner.clone(),
+			network.clone(),
+			collator_network.clone(),
+		)));
+
 		Self {
 			proposer_factory: Arc::new(Mutex::new(proposer_factory)),
 			inherent_data_providers,
 			_phantom: PhantomData,
-			collator_network: Arc::new(collator_network),
+			collator_network,
 			block_import: Arc::new(Mutex::new(block_import)),
 			spawner,
 			network,
+			wait_to_announce,
 		}
 	}
 }
@@ -97,6 +106,7 @@ impl<Block: BlockT, PF, BI> Clone for Collator<Block, PF, BI> {
 			block_import: self.block_import.clone(),
 			spawner: self.spawner.clone(),
 			network: self.network.clone(),
+			wait_to_announce: self.wait_to_announce.clone(),
 		}
 	}
 }
@@ -143,9 +153,7 @@ where
 			.lock()
 			.init(&last_head.header);
 
-		let spawner = self.spawner.clone();
-		let network = self.network.clone();
-		let collator_network = self.collator_network.clone();
+		let wait_to_announce = self.wait_to_announce.clone();
 
 		Box::pin(async move {
 			let parent_state_root = *last_head.header.state_root();
@@ -244,20 +252,11 @@ where
 				parachain::HeadData(head_data.encode()),
 			);
 
-			if let Err(err) = spawner.spawn_obj(Box::pin(async move {
-				wait_to_announce(
-					hash,
-					relay_chain_parent,
-					network, collator_network,
-					&encoded_header,
-				).await
-			}).into()) {
-				error!(
-					target: "cumulus-collator",
-					"Could not spawn a new task to wait for the announce block: {:?}",
-					err
-				);
-			}
+			wait_to_announce.lock().wait_to_announce(
+				hash,
+				relay_chain_parent,
+				encoded_header,
+			);
 
 			trace!(target: "cumulus-collator", "Produced candidate: {:?}", candidate);
 

network/Cargo.toml

@@ -24,3 +24,5 @@ polkadot-network = { git = "https://github.com/paritytech/polkadot", branch = "c
 # other deps
 codec = { package = "parity-scale-codec", version = "1.3.0", features = [ "derive" ] }
 futures = { version = "0.3.1", features = ["compat"] }
+stream-cancel = "0.5"
+log = "0.4.8"

network/src/lib.rs

@@ -34,7 +34,11 @@ use polkadot_statement_table::{SignedStatement, Statement};
 use polkadot_validation::check_statement;
 
 use codec::{Decode, Encode};
-use futures::StreamExt;
+use futures::{pin_mut, select, StreamExt};
+use futures::future::FutureExt;
+use futures::task::Spawn;
+use stream_cancel::{Trigger, Tripwire};
+use log::{error, trace};
 
 use std::{marker::PhantomData, sync::Arc};
 
@@ -145,7 +149,73 @@ where
 	}
 }
 
-pub async fn wait_to_announce<Block: BlockT>(
+pub struct WaitToAnnounce<Block: BlockT> {
+	spawner: Arc<dyn Spawn + Send + Sync>,
+	network: Arc<NetworkService<Block, <Block as BlockT>::Hash>>,
+	collator_network: Arc<dyn CollatorNetwork>,
+	current_trigger: Trigger,
+}
+
+impl<Block: BlockT> WaitToAnnounce<Block> {
+	pub fn new(
+		spawner: Arc<dyn Spawn + Send + Sync>,
+		network: Arc<NetworkService<Block, <Block as BlockT>::Hash>>,
+		collator_network: Arc<dyn CollatorNetwork>,
+	) -> WaitToAnnounce<Block> {
+		let (current_trigger, _tripwire) = Tripwire::new();
+
+		WaitToAnnounce {
+			spawner,
+			network,
+			collator_network,
+			current_trigger,
+		}
+	}
+
+	pub fn wait_to_announce(
+		&mut self,
+		hash: <Block as BlockT>::Hash,
+		relay_chain_leaf: PHash,
+		head_data: Vec<u8>,
+	) {
+		let (trigger, tripwire) = Tripwire::new();
+		let network = self.network.clone();
+		let collator_network = self.collator_network.clone();
+
+		std::mem::replace(&mut self.current_trigger, trigger).cancel();
+
+		if let Err(err) = self.spawner.spawn_obj(Box::pin(async move {
+			let t1 = wait_to_announce(
+				hash,
+				relay_chain_leaf,
+				network,
+				collator_network,
+				&head_data,
+			).fuse();
+			let t2 = tripwire.fuse();
+
+			pin_mut!(t1, t2);
+
+			select! {
+				_ = t1 => {},
+				_ = t2 => {
+					trace!(
+						target: "cumulus-network",
+						"previous task that waits for announce block has been canceled",
+					);
+				}
+			}
+		}).into()) {
+			error!(
+				target: "cumulus-network",
+				"Could not spawn a new task to wait for the announce block: {:?}",
+				err
+			);
+		}
+	}
+}
+
+async fn wait_to_announce<Block: BlockT>(
 	hash: <Block as BlockT>::Hash,
 	relay_chain_leaf: PHash,
 	network: Arc<NetworkService<Block, <Block as BlockT>::Hash>>,