Import queue API revamp (#2856)

* ImportQueue methods are now mut

* Link methods are now mut

* Remove Arc from BasicSyncQueue

* Fix tests

* Remove BasicSyncQueue

* Change the import queue API

* Add buffered_link

* Remove obsolete tests

* Comments and style improvement pass

* Fix grandpa and comment cleanup

* Update core/consensus/common/src/import_queue.rs

Co-Authored-By: André Silva <andre.beat@gmail.com>

substrate/core/consensus/common/src/block_import.rs

@@ -195,7 +195,7 @@ pub trait JustificationImport<B: BlockT> {
 	type Error: ::std::error::Error + Send + 'static;
 
 	/// Called by the import queue when it is started.
-	fn on_start(&self, _link: &dyn crate::import_queue::Link<B>) { }
+	fn on_start(&self, _link: &mut dyn crate::import_queue::Link<B>) { }
 
 	/// Import a Block justification and finalize the given block.
 	fn import_justification(
@@ -211,7 +211,7 @@ pub trait FinalityProofImport<B: BlockT> {
 	type Error: std::error::Error + Send + 'static;
 
 	/// Called by the import queue when it is started.
-	fn on_start(&self, _link: &dyn crate::import_queue::Link<B>) { }
+	fn on_start(&self, _link: &mut dyn crate::import_queue::Link<B>) { }
 
 	/// Import a Block justification and finalize the given block. Returns finalized block or error.
 	fn import_finality_proof(

substrate/core/consensus/common/src/lib.rs

@@ -26,7 +26,7 @@
 // our error-chain could potentially blow up otherwise
 #![recursion_limit="128"]
 
-#[macro_use] extern crate crossbeam_channel;
+extern crate crossbeam_channel;
 #[macro_use] extern crate log;
 
 use std::sync::Arc;

substrate/core/finality-grandpa/src/import.rs

@@ -76,7 +76,7 @@ impl<B, E, Block: BlockT<Hash=H256>, RA, PRA, SC> JustificationImport<Block>
 {
 	type Error = ConsensusError;
 
-	fn on_start(&self, link: &dyn consensus_common::import_queue::Link<Block>) {
+	fn on_start(&self, link: &mut dyn consensus_common::import_queue::Link<Block>) {
 		let chain_info = self.inner.info().chain;
 
 		// request justifications for all pending changes for which change blocks have already been imported

substrate/core/finality-grandpa/src/light_import.rs

@@ -144,7 +144,7 @@ impl<B, E, Block: BlockT<Hash=H256>, RA> FinalityProofImport<Block>
 {
 	type Error = ConsensusError;
 
-	fn on_start(&self, link: &dyn consensus_common::import_queue::Link<Block>) {
+	fn on_start(&self, link: &mut dyn consensus_common::import_queue::Link<Block>) {
 		let chain_info = self.client.info().chain;
 
 		let data = self.data.read();
@@ -572,7 +572,7 @@ pub mod tests {
 	{
 		type Error = ConsensusError;
 
-		fn on_start(&self, link: &dyn consensus_common::import_queue::Link<Block>) {
+		fn on_start(&self, link: &mut dyn consensus_common::import_queue::Link<Block>) {
 			self.0.on_start(link)
 		}
 

substrate/core/network/src/service.rs

@@ -74,79 +74,6 @@ pub trait TransactionPool<H: ExHashT, B: BlockT>: Send + Sync {
 	fn on_broadcasted(&self, propagations: HashMap<H, Vec<String>>);
 }
 
-/// A link implementation that connects to the network.
-#[derive(Clone)]
-pub struct NetworkLink<B: BlockT, S: NetworkSpecialization<B>> {
-	/// The protocol sender
-	pub(crate) protocol_sender: mpsc::UnboundedSender<ProtocolMsg<B, S>>,
-	/// The network sender
-	pub(crate) network_sender: mpsc::UnboundedSender<NetworkMsg<B>>,
-}
-
-impl<B: BlockT, S: NetworkSpecialization<B>> Link<B> for NetworkLink<B, S> {
-	fn block_imported(&self, hash: &B::Hash, number: NumberFor<B>) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::BlockImportedSync(hash.clone(), number));
-	}
-
-	fn blocks_processed(&self, processed_blocks: Vec<B::Hash>, has_error: bool) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::BlocksProcessed(processed_blocks, has_error));
-	}
-
-	fn justification_imported(&self, who: PeerId, hash: &B::Hash, number: NumberFor<B>, success: bool) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::JustificationImportResult(hash.clone(), number, success));
-		if !success {
-			info!("Invalid justification provided by {} for #{}", who, hash);
-			let _ = self.network_sender.unbounded_send(NetworkMsg::ReportPeer(who.clone(), i32::min_value()));
-			let _ = self.network_sender.unbounded_send(NetworkMsg::DisconnectPeer(who.clone()));
-		}
-	}
-
-	fn clear_justification_requests(&self) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::ClearJustificationRequests);
-	}
-
-	fn request_justification(&self, hash: &B::Hash, number: NumberFor<B>) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::RequestJustification(hash.clone(), number));
-	}
-
-	fn request_finality_proof(&self, hash: &B::Hash, number: NumberFor<B>) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::RequestFinalityProof(
-			hash.clone(),
-			number,
-		));
-	}
-
-	fn finality_proof_imported(
-		&self,
-		who: PeerId,
-		request_block: (B::Hash, NumberFor<B>),
-		finalization_result: Result<(B::Hash, NumberFor<B>), ()>,
-	) {
-		let success = finalization_result.is_ok();
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::FinalityProofImportResult(
-			request_block,
-			finalization_result,
-		));
-		if !success {
-			info!("Invalid finality proof provided by {} for #{}", who, request_block.0);
-			let _ = self.network_sender.unbounded_send(NetworkMsg::ReportPeer(who.clone(), i32::min_value()));
-			let _ = self.network_sender.unbounded_send(NetworkMsg::DisconnectPeer(who.clone()));
-		}
-	}
-
-	fn report_peer(&self, who: PeerId, reputation_change: i32) {
-		let _ = self.network_sender.unbounded_send(NetworkMsg::ReportPeer(who, reputation_change));
-	}
-
-	fn restart(&self) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::RestartSync);
-	}
-
-	fn set_finality_proof_request_builder(&self, request_builder: SharedFinalityProofRequestBuilder<B>) {
-		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::SetFinalityProofRequestBuilder(request_builder));
-	}
-}
-
 /// A cloneable handle for reporting cost/benefits of peers.
 #[derive(Clone)]
 pub struct ReportHandle {
@@ -197,13 +124,6 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>, H: ExHashT> NetworkWorker
 		let (protocol_sender, protocol_rx) = mpsc::unbounded();
 		let status_sinks = Arc::new(Mutex::new(Vec::new()));
 
-		// connect the import-queue to the network service.
-		let link = NetworkLink {
-			protocol_sender: protocol_sender.clone(),
-			network_sender: network_chan.clone(),
-		};
-		params.import_queue.start(Box::new(link))?;
-
 		// Start in off-line mode, since we're not connected to any nodes yet.
 		let is_offline = Arc::new(AtomicBool::new(true));
 		let is_major_syncing = Arc::new(AtomicBool::new(false));
@@ -592,7 +512,7 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>, H: ExHashT> Future for Ne
 	type Error = io::Error;
 
 	fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
-		// Implementation of `protocol::NetworkOut` using the available local variables.
+		// Implementation of `protocol::NetworkOut` trait using the available local variables.
 		struct Context<'a, B: BlockT>(&'a mut Swarm<B>, &'a PeersetHandle);
 		impl<'a, B: BlockT> NetworkOut<B> for Context<'a, B> {
 			fn report_peer(&mut self, who: PeerId, reputation: i32) {
@@ -606,11 +526,74 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>, H: ExHashT> Future for Ne
 			}
 		}
 
+		// Implementation of `import_queue::Link` trait using the available local variables.
+		struct NetworkLink<'a, B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> {
+			protocol: &'a mut Protocol<B, S, H>,
+			context: Context<'a, B>,
+		}
+		impl<'a, B: BlockT, S: NetworkSpecialization<B>, H: ExHashT> Link<B> for NetworkLink<'a, B, S, H> {
+			fn block_imported(&mut self, hash: &B::Hash, number: NumberFor<B>) {
+				self.protocol.block_imported(&hash, number)
+			}
+			fn blocks_processed(&mut self, hashes: Vec<B::Hash>, has_error: bool) {
+				self.protocol.blocks_processed(&mut self.context, hashes, has_error)
+			}
+			fn justification_imported(&mut self, who: PeerId, hash: &B::Hash, number: NumberFor<B>, success: bool) {
+				self.protocol.justification_import_result(hash.clone(), number, success);
+				if !success {
+					info!("Invalid justification provided by {} for #{}", who, hash);
+					self.context.0.user_protocol_mut().disconnect_peer(&who);
+					self.context.1.report_peer(who, i32::min_value());
+				}
+			}
+			fn clear_justification_requests(&mut self) {
+				self.protocol.clear_justification_requests()
+			}
+			fn request_justification(&mut self, hash: &B::Hash, number: NumberFor<B>) {
+				self.protocol.request_justification(&mut self.context, hash, number)
+			}
+			fn request_finality_proof(&mut self, hash: &B::Hash, number: NumberFor<B>) {
+				self.protocol.request_finality_proof(&mut self.context, hash, number)
+			}
+			fn finality_proof_imported(
+				&mut self,
+				who: PeerId,
+				request_block: (B::Hash, NumberFor<B>),
+				finalization_result: Result<(B::Hash, NumberFor<B>), ()>,
+			) {
+				let success = finalization_result.is_ok();
+				self.protocol.finality_proof_import_result(request_block, finalization_result);
+				if !success {
+					info!("Invalid finality proof provided by {} for #{}", who, request_block.0);
+					self.context.0.user_protocol_mut().disconnect_peer(&who);
+					self.context.1.report_peer(who, i32::min_value());
+				}
+			}
+			fn report_peer(&mut self, who: PeerId, reputation_change: i32) {
+				self.context.1.report_peer(who, reputation_change)
+			}
+			fn restart(&mut self) {
+				self.protocol.restart(&mut self.context)
+			}
+			fn set_finality_proof_request_builder(&mut self, builder: SharedFinalityProofRequestBuilder<B>) {
+				self.protocol.set_finality_proof_request_builder(builder)
+			}
+		}
+
 		while let Ok(Async::Ready(_)) = self.status_interval.poll() {
 			let status = self.protocol.status();
 			self.status_sinks.lock().retain(|sink| sink.unbounded_send(status.clone()).is_ok());
 		}
 
+		{
+			let mut network_service = self.network_service.lock();
+			let mut link = NetworkLink {
+				protocol: &mut self.protocol,
+				context: Context(&mut network_service, &self.peerset),
+			};
+			self.import_queue.poll_actions(&mut link);
+		}
+
 		while let Ok(Async::Ready(_)) = self.connected_peers_interval.poll() {
 			let infos = self.protocol.peers_info().map(|(id, info)| {
 				(id.clone(), ConnectedPeer { peer_info: info.clone() })
@@ -618,10 +601,14 @@ impl<B: BlockT + 'static, S: NetworkSpecialization<B>, H: ExHashT> Future for Ne
 			*self.peers.write() = infos;
 		}
 
-		match self.protocol.poll(&mut Context(&mut self.network_service.lock(), &self.peerset), &*self.transaction_pool) {
-			Ok(Async::Ready(v)) => void::unreachable(v),
-			Ok(Async::NotReady) => {}
-			Err(err) => void::unreachable(err),
+		{
+			let mut network_service = self.network_service.lock();
+			let mut ctxt = Context(&mut *network_service, &self.peerset);
+			match self.protocol.poll(&mut ctxt, &*self.transaction_pool) {
+				Ok(Async::Ready(v)) => void::unreachable(v),
+				Ok(Async::NotReady) => {}
+				Err(err) => void::unreachable(err),
+			}
 		}
 
 		// Check for new incoming on-demand requests.

substrate/core/network/src/test/block_import.rs

@@ -77,8 +77,7 @@ fn async_import_queue_drops() {
 	// Perform this test multiple times since it exhibits non-deterministic behavior.
 	for _ in 0..100 {
 		let verifier = Arc::new(PassThroughVerifier(true));
-		let queue = BasicQueue::new(verifier, Arc::new(test_client::new()), None, None, None);
-		queue.start(Box::new(TestLink{})).unwrap();
+		let mut queue = BasicQueue::new(verifier, Arc::new(test_client::new()), None, None, None);
 		drop(queue);
 	}
 }

substrate/core/network/src/test/mod.rs

@@ -41,6 +41,7 @@ use consensus::{Error as ConsensusError, well_known_cache_keys::{self, Id as Cac
 use consensus::{BlockOrigin, ForkChoiceStrategy, ImportBlock, JustificationImport};
 use crate::consensus_gossip::{ConsensusGossip, MessageRecipient as GossipMessageRecipient, TopicNotification};
 use futures::{prelude::*, sync::{mpsc, oneshot}};
+use log::info;
 use crate::message::Message;
 use libp2p::PeerId;
 use parking_lot::{Mutex, RwLock};
@@ -49,7 +50,7 @@ use crate::protocol::{Context, Protocol, ProtocolConfig, ProtocolStatus, CustomM
 use runtime_primitives::generic::{BlockId, OpaqueDigestItemId};
 use runtime_primitives::traits::{Block as BlockT, Header, NumberFor};
 use runtime_primitives::{Justification, ConsensusEngineId};
-use crate::service::{NetworkLink, NetworkMsg, ProtocolMsg, TransactionPool};
+use crate::service::{NetworkMsg, ProtocolMsg, TransactionPool};
 use crate::specialization::NetworkSpecialization;
 use test_client::{self, AccountKeyring};
 
@@ -97,6 +98,79 @@ pub struct NoopLink { }
 
 impl<B: BlockT> Link<B> for NoopLink { }
 
+/// A link implementation that connects to the network.
+#[derive(Clone)]
+pub struct NetworkLink<B: BlockT, S: NetworkSpecialization<B>> {
+	/// The protocol sender
+	pub(crate) protocol_sender: mpsc::UnboundedSender<ProtocolMsg<B, S>>,
+	/// The network sender
+	pub(crate) network_sender: mpsc::UnboundedSender<NetworkMsg<B>>,
+}
+
+impl<B: BlockT, S: NetworkSpecialization<B>> Link<B> for NetworkLink<B, S> {
+	fn block_imported(&mut self, hash: &B::Hash, number: NumberFor<B>) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::BlockImportedSync(hash.clone(), number));
+	}
+
+	fn blocks_processed(&mut self, processed_blocks: Vec<B::Hash>, has_error: bool) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::BlocksProcessed(processed_blocks, has_error));
+	}
+
+	fn justification_imported(&mut self, who: PeerId, hash: &B::Hash, number: NumberFor<B>, success: bool) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::JustificationImportResult(hash.clone(), number, success));
+		if !success {
+			info!("Invalid justification provided by {} for #{}", who, hash);
+			let _ = self.network_sender.unbounded_send(NetworkMsg::ReportPeer(who.clone(), i32::min_value()));
+			let _ = self.network_sender.unbounded_send(NetworkMsg::DisconnectPeer(who.clone()));
+		}
+	}
+
+	fn clear_justification_requests(&mut self) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::ClearJustificationRequests);
+	}
+
+	fn request_justification(&mut self, hash: &B::Hash, number: NumberFor<B>) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::RequestJustification(hash.clone(), number));
+	}
+
+	fn request_finality_proof(&mut self, hash: &B::Hash, number: NumberFor<B>) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::RequestFinalityProof(
+			hash.clone(),
+			number,
+		));
+	}
+
+	fn finality_proof_imported(
+		&mut self,
+		who: PeerId,
+		request_block: (B::Hash, NumberFor<B>),
+		finalization_result: Result<(B::Hash, NumberFor<B>), ()>,
+	) {
+		let success = finalization_result.is_ok();
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::FinalityProofImportResult(
+			request_block,
+			finalization_result,
+		));
+		if !success {
+			info!("Invalid finality proof provided by {} for #{}", who, request_block.0);
+			let _ = self.network_sender.unbounded_send(NetworkMsg::ReportPeer(who.clone(), i32::min_value()));
+			let _ = self.network_sender.unbounded_send(NetworkMsg::DisconnectPeer(who.clone()));
+		}
+	}
+
+	fn report_peer(&mut self, who: PeerId, reputation_change: i32) {
+		let _ = self.network_sender.unbounded_send(NetworkMsg::ReportPeer(who, reputation_change));
+	}
+
+	fn restart(&mut self) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::RestartSync);
+	}
+
+	fn set_finality_proof_request_builder(&mut self, request_builder: SharedFinalityProofRequestBuilder<B>) {
+		let _ = self.protocol_sender.unbounded_send(ProtocolMsg::SetFinalityProofRequestBuilder(request_builder));
+	}
+}
+
 /// The test specialization.
 #[derive(Clone)]
 pub struct DummySpecialization;
@@ -232,24 +306,24 @@ impl<S: NetworkSpecialization<Block>> TestLink<S> {
 }
 
 impl<S: NetworkSpecialization<Block>> Link<Block> for TestLink<S> {
-	fn block_imported(&self, hash: &Hash, number: NumberFor<Block>) {
+	fn block_imported(&mut self, hash: &Hash, number: NumberFor<Block>) {
 		self.link.block_imported(hash, number);
 	}
 
-	fn blocks_processed(&self, processed_blocks: Vec<Hash>, has_error: bool) {
+	fn blocks_processed(&mut self, processed_blocks: Vec<Hash>, has_error: bool) {
 		self.link.blocks_processed(processed_blocks, has_error);
 	}
 
-	fn justification_imported(&self, who: PeerId, hash: &Hash, number:NumberFor<Block>, success: bool) {
+	fn justification_imported(&mut self, who: PeerId, hash: &Hash, number:NumberFor<Block>, success: bool) {
 		self.link.justification_imported(who, hash, number, success);
 	}
 
-	fn request_justification(&self, hash: &Hash, number: NumberFor<Block>) {
+	fn request_justification(&mut self, hash: &Hash, number: NumberFor<Block>) {
 		self.link.request_justification(hash, number);
 	}
 
 	fn finality_proof_imported(
-		&self,
+		&mut self,
 		who: PeerId,
 		request_block: (Hash, NumberFor<Block>),
 		finalization_result: Result<(Hash, NumberFor<Block>), ()>,
@@ -257,19 +331,19 @@ impl<S: NetworkSpecialization<Block>> Link<Block> for TestLink<S> {
 		self.link.finality_proof_imported(who, request_block, finalization_result);
 	}
 
-	fn request_finality_proof(&self, hash: &Hash, number: NumberFor<Block>) {
+	fn request_finality_proof(&mut self, hash: &Hash, number: NumberFor<Block>) {
 		self.link.request_finality_proof(hash, number);
 	}
 
-	fn set_finality_proof_request_builder(&self, request_builder: SharedFinalityProofRequestBuilder<Block>) {
+	fn set_finality_proof_request_builder(&mut self, request_builder: SharedFinalityProofRequestBuilder<Block>) {
 		self.link.set_finality_proof_request_builder(request_builder);
 	}
 
-	fn report_peer(&self, who: PeerId, reputation_change: i32) {
+	fn report_peer(&mut self, who: PeerId, reputation_change: i32) {
 		self.link.report_peer(who, reputation_change);
 	}
 
-	fn restart(&self) {
+	fn restart(&mut self) {
 		self.link.restart();
 	}
 
@@ -278,7 +352,7 @@ impl<S: NetworkSpecialization<Block>> Link<Block> for TestLink<S> {
 	/// The caller should wait for the `Link::synchronized` call to ensure that it has synchronized
 	/// with `ImportQueue`.
 	#[cfg(any(test, feature = "test-helpers"))]
-	fn synchronized(&self) {
+	fn synchronized(&mut self) {
 		drop(self.network_to_protocol_sender.unbounded_send(FromNetworkMsg::Synchronize))
 	}
 }
@@ -292,7 +366,7 @@ pub struct Peer<D, S: NetworkSpecialization<Block>> {
 	/// we allow it to be unused.
 	#[cfg_attr(not(test), allow(unused))]
 	protocol_status: Arc<RwLock<ProtocolStatus<Block>>>,
-	import_queue: Box<BasicQueue<Block>>,
+	import_queue: Arc<Mutex<Box<BasicQueue<Block>>>>,
 	pub data: D,
 	best_hash: Mutex<Option<H256>>,
 	finalized_hash: Mutex<Option<H256>>,
@@ -437,7 +511,7 @@ impl<D, S: NetworkSpecialization<Block>> Peer<D, S> {
 	fn new(
 		protocol_status: Arc<RwLock<ProtocolStatus<Block>>>,
 		client: PeersClient,
-		import_queue: Box<BasicQueue<Block>>,
+		import_queue: Arc<Mutex<Box<BasicQueue<Block>>>>,
 		use_tokio: bool,
 		network_to_protocol_sender: mpsc::UnboundedSender<FromNetworkMsg<Block>>,
 		protocol_sender: mpsc::UnboundedSender<ProtocolMsg<Block, S>>,
@@ -451,12 +525,6 @@ impl<D, S: NetworkSpecialization<Block>> Peer<D, S> {
 			protocol_sender.clone(),
 			network_port,
 		);
-		let network_link = TestLink::new(
-			protocol_sender.clone(),
-			network_to_protocol_sender.clone(),
-			network_sender.clone(),
-		);
-		import_queue.start(Box::new(network_link)).expect("Test ImportQueue always starts");
 		Peer {
 			protocol_status,
 			peer_id: PeerId::random(),
@@ -535,7 +603,7 @@ impl<D, S: NetworkSpecialization<Block>> Peer<D, S> {
 	/// Synchronize with import queue.
 	#[cfg(any(test, feature = "test-helpers"))]
 	pub fn import_queue_sync(&self) {
-		self.import_queue.synchronize();
+		self.import_queue.lock().synchronize();
 		let _ = self.net_proto_channel.wait_sync();
 	}
 
@@ -663,7 +731,7 @@ impl<D, S: NetworkSpecialization<Block>> Peer<D, S> {
 			);
 			let header = block.header.clone();
 			at = hash;
-			self.import_queue.import_blocks(
+			self.import_queue.lock().import_blocks(
 				origin,
 				vec![IncomingBlock {
 					origin: None,
@@ -806,10 +874,12 @@ pub trait TestNetFactory: Sized {
 	fn add_peer(
 		&mut self,
 		protocol_status: Arc<RwLock<ProtocolStatus<Block>>>,
-		import_queue: Box<BasicQueue<Block>>,
+		import_queue: Arc<Mutex<Box<BasicQueue<Block>>>>,
 		tx_pool: EmptyTransactionPool,
 		finality_proof_provider: Option<Arc<dyn FinalityProofProvider<Block>>>,
 		mut protocol: Protocol<Block, Self::Specialization, Hash>,
+		protocol_sender: mpsc::UnboundedSender<ProtocolMsg<Block, Self::Specialization>>,
+		network_to_protocol_sender: mpsc::UnboundedSender<FromNetworkMsg<Block>>,
 		network_sender: mpsc::UnboundedSender<NetworkMsg<Block>>,
 		mut network_to_protocol_rx: mpsc::UnboundedReceiver<FromNetworkMsg<Block>>,
 		mut protocol_rx: mpsc::UnboundedReceiver<ProtocolMsg<Block, Self::Specialization>>,
@@ -831,6 +901,12 @@ pub trait TestNetFactory: Sized {
 			}
 
 			tokio::runtime::current_thread::run(futures::future::poll_fn(move || {
+				import_queue.lock().poll_actions(&mut TestLink::new(
+					protocol_sender.clone(),
+					network_to_protocol_sender.clone(),
+					network_sender.clone(),
+				));
+
 				while let Async::Ready(msg) = network_to_protocol_rx.poll().unwrap() {
 					let outcome = match msg {
 						Some(FromNetworkMsg::PeerConnected(peer_id)) => {
@@ -858,11 +934,11 @@ pub trait TestNetFactory: Sized {
 
 					match outcome {
 						CustomMessageOutcome::BlockImport(origin, blocks) =>
-							import_queue.import_blocks(origin, blocks),
+							import_queue.lock().import_blocks(origin, blocks),
 						CustomMessageOutcome::JustificationImport(origin, hash, nb, justification) =>
-							import_queue.import_justification(origin, hash, nb, justification),
+							import_queue.lock().import_justification(origin, hash, nb, justification),
 						CustomMessageOutcome::FinalityProofImport(origin, hash, nb, proof) =>
-							import_queue.import_finality_proof(origin, hash, nb, proof),
+							import_queue.lock().import_finality_proof(origin, hash, nb, proof),
 						CustomMessageOutcome::None => {}
 					}
 				}
@@ -959,13 +1035,13 @@ pub trait TestNetFactory: Sized {
 			= self.make_block_import(PeersClient::Full(client.clone()));
 		let (network_sender, network_port) = mpsc::unbounded();
 
-		let import_queue = Box::new(BasicQueue::new(
+		let import_queue = Arc::new(Mutex::new(Box::new(BasicQueue::new(
 			verifier,
 			block_import,
 			justification_import,
 			finality_proof_import,
 			finality_proof_request_builder,
-		));
+		))));
 		let specialization = self::SpecializationFactory::create();
 
 		let (network_to_protocol_sender, network_to_protocol_rx) = mpsc::unbounded();
@@ -985,6 +1061,8 @@ pub trait TestNetFactory: Sized {
 			EmptyTransactionPool,
 			self.make_finality_proof_provider(PeersClient::Full(client.clone())),
 			protocol,
+			protocol_sender.clone(),
+			network_to_protocol_sender.clone(),
 			network_sender.clone(),
 			network_to_protocol_rx,
 			protocol_rx,
@@ -1013,13 +1091,13 @@ pub trait TestNetFactory: Sized {
 			= self.make_block_import(PeersClient::Light(client.clone()));
 		let (network_sender, network_port) = mpsc::unbounded();
 
-		let import_queue = Box::new(BasicQueue::new(
+		let import_queue = Arc::new(Mutex::new(Box::new(BasicQueue::new(
 			verifier,
 			block_import,
 			justification_import,
 			finality_proof_import,
 			finality_proof_request_builder,
-		));
+		))));
 		let specialization = self::SpecializationFactory::create();
 
 		let (network_to_protocol_sender, network_to_protocol_rx) = mpsc::unbounded();
@@ -1039,6 +1117,8 @@ pub trait TestNetFactory: Sized {
 			EmptyTransactionPool,
 			self.make_finality_proof_provider(PeersClient::Light(client.clone())),
 			protocol,
+			protocol_sender.clone(),
+			network_to_protocol_sender.clone(),
 			network_sender.clone(),
 			network_to_protocol_rx,
 			protocol_rx,

substrate/core/service/src/chain_ops.rs

@@ -17,7 +17,7 @@
 //! Chain utilities.
 
 use std::{self, io::{Read, Write}};
-use futures::Future;
+use futures::prelude::*;
 use log::{info, warn};
 
 use runtime_primitives::generic::{SignedBlock, BlockId};
@@ -99,21 +99,20 @@ pub fn export_blocks<F, E, W>(
 }
 
 struct WaitLink {
-	wait_send: std::sync::mpsc::Sender<()>,
+	imported_blocks: u64,
 }
 
 impl WaitLink {
-	fn new(wait_send: std::sync::mpsc::Sender<()>) -> WaitLink {
+	fn new() -> WaitLink {
 		WaitLink {
-			wait_send,
+			imported_blocks: 0,
 		}
 	}
 }
 
 impl<B: Block> Link<B> for WaitLink {
-	fn block_imported(&self, _hash: &B::Hash, _number: NumberFor<B>) {
-		self.wait_send.send(())
-			.expect("Unable to notify main process; if the main process panicked then this thread would already be dead as well. qed.");
+	fn block_imported(&mut self, _hash: &B::Hash, _number: NumberFor<B>) {
+		self.imported_blocks += 1;
 	}
 }
 
@@ -128,11 +127,7 @@ pub fn import_blocks<F, E, R>(
 	let client = new_client::<F>(&config)?;
 	// FIXME #1134 this shouldn't need a mutable config.
 	let select_chain = components::FullComponents::<F>::build_select_chain(&mut config, client.clone())?;
-	let queue = components::FullComponents::<F>::build_import_queue(&mut config, client.clone(), select_chain)?;
-
-	let (wait_send, wait_recv) = std::sync::mpsc::channel();
-	let wait_link = WaitLink::new(wait_send);
-	queue.start(Box::new(wait_link))?;
+	let mut queue = components::FullComponents::<F>::build_import_queue(&mut config, client.clone(), select_chain)?;
 
 	let (exit_send, exit_recv) = std::sync::mpsc::channel();
 	::std::thread::spawn(move || {
@@ -179,19 +174,23 @@ pub fn import_blocks<F, E, R>(
 		}
 	}
 
-	let mut blocks_imported = 0;
-	while blocks_imported < count {
-		wait_recv.recv()
-			.expect("Importing thread has panicked. Then the main process will die before this can be reached. qed.");
-		blocks_imported += 1;
-		if blocks_imported % 1000 == 0 {
+	let mut link = WaitLink::new();
+	tokio::run(futures::future::poll_fn(move || {
+		let blocks_before = link.imported_blocks;
+		queue.poll_actions(&mut link);
+		if link.imported_blocks / 1000 != blocks_before / 1000 {
 			info!(
 				"#{} blocks were imported (#{} left)",
-				blocks_imported,
-				count - blocks_imported
+				link.imported_blocks,
+				count - link.imported_blocks
 			);
 		}
-	}
+		if link.imported_blocks >= count {
+			Ok(Async::Ready(()))
+		} else {
+			Ok(Async::NotReady)
+		}
+	}));
 
 	info!("Imported {} blocks. Best: #{}", block_count, client.info().chain.best_number);