Replace NetworkWorker::poll with NetworkWorker::next_action

client/network/src/service.rs

@@ -59,7 +59,7 @@ use sp_utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnbound
 use std::{
 	borrow::Cow,
 	collections::HashSet,
-	fs, io,
+	fs,
 	marker::PhantomData,
 	pin::Pin,
 	str,
@@ -991,10 +991,17 @@ impl Metrics {
 	}
 }
 
-impl<B: BlockT + 'static, H: ExHashT> Future for NetworkWorker<B, H> {
-	type Output = Result<(), io::Error>;
+impl<B: BlockT + 'static, H: ExHashT> NetworkWorker<B, H> {
+	/// Performs one action on the network, then returns.
+	///
+	/// The returned future is designed to be freely cancellable.
+	pub async fn next_action(&mut self) {
+		future::poll_fn(move |cx| Pin::new(&mut *self).poll(cx)).await
+	}
 
-	fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<Self::Output> {
+	/// Implementation of `next_action`. Note that this is not an implementation of `Future` but
+	/// a regular method.
+	fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context) -> Poll<()> {
 		let this = &mut *self;
 
 		// Poll the import queue for actions to perform.
@@ -1019,7 +1026,7 @@ impl<B: BlockT + 'static, H: ExHashT> Future for NetworkWorker<B, H> {
 			// Process the next message coming from the `NetworkService`.
 			let msg = match this.from_worker.poll_next_unpin(cx) {
 				Poll::Ready(Some(msg)) => msg,
-				Poll::Ready(None) => return Poll::Ready(Ok(())),
+				Poll::Ready(None) => return Poll::Ready(()),
 				Poll::Pending => break,
 			};
 

client/network/src/service/tests.rs

@@ -87,7 +87,7 @@ fn build_test_full_node(config: config::NetworkConfiguration)
 		None,
 	));
 
-	let worker = NetworkWorker::new(config::Params {
+	let mut worker = NetworkWorker::new(config::Params {
 		role: config::Role::Full,
 		executor: None,
 		network_config: config,
@@ -109,8 +109,9 @@ fn build_test_full_node(config: config::NetworkConfiguration)
 	let event_stream = service.event_stream("test");
 
 	async_std::task::spawn(async move {
-		futures::pin_mut!(worker);
-		let _ = worker.await;
+		loop {
+			worker.next_action().await;
+		}
 	});
 
 	(service, event_stream)

client/network/test/src/lib.rs

@@ -789,8 +789,12 @@ pub trait TestNetFactory: Sized {
 		self.mut_peers(|peers| {
 			for peer in peers {
 				trace!(target: "sync", "-- Polling {}", peer.id());
-				if let Poll::Ready(res) = Pin::new(&mut peer.network).poll(cx) {
-					res.unwrap();
+				loop {
+					let net_poll_future = peer.network.next_action();
+					futures::pin_mut!(net_poll_future);
+					if let Poll::Pending = net_poll_future.poll(cx) {
+						break;
+					}
 				}
 				trace!(target: "sync", "-- Polling complete {}", peer.id());
 

client/service/src/lib.rs

@@ -445,9 +445,9 @@ fn build_network_future<
 		});
 
 		// Main network polling.
-		if let Poll::Ready(Ok(())) = Pin::new(&mut network).poll(cx).map_err(|err| {
-			warn!(target: "service", "Error in network: {:?}", err);
-		}) {
+		let next_network_action = network.next_action();
+		futures::pin_mut!(next_network_action);
+		if let Poll::Ready(()) = next_network_action.poll(cx) {
 			return Poll::Ready(());
 		}