extract utility function into util crate

node/network/collator-protocol/Cargo.toml

@@ -8,7 +8,6 @@ edition = "2018"
 futures = "0.3.5"
 log = "0.4.11"
 derive_more = "0.99.9"
-pin-project = "0.4.22"
 
 codec = { package="parity-scale-codec", version = "1.3.4", features = ["std"]  }
 

node/network/collator-protocol/src/collator_side.rs

@@ -15,19 +15,17 @@
 // along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
 
 use std::collections::HashMap;
-use std::pin::Pin;
 
 use super::{TARGET,  Result};
 
-use futures::channel::{mpsc, oneshot};
-use futures::stream::{self, StreamExt as _};
-use futures::task::{Poll, self};
+use futures::channel::oneshot;
+use futures::stream::StreamExt as _;
+use futures::task::Poll;
 use log::warn;
-use pin_project::pin_project;
 
 use polkadot_primitives::v1::{
 	CollatorId, CoreIndex, CoreState, Hash, Id as ParaId, CandidateReceipt,
-	PoV, ValidatorId, AuthorityDiscoveryId,
+	PoV, ValidatorId,
 };
 use polkadot_subsystem::{
 	FromOverseer, OverseerSignal, SubsystemContext,
@@ -40,6 +38,7 @@ use polkadot_node_network_protocol::{
 	v1 as protocol_v1, View, PeerId, NetworkBridgeEvent, RequestId,
 };
 use polkadot_node_subsystem_util::{
+	validator_discovery,
 	request_validators_ctx,
 	request_validator_groups_ctx,
 };
@@ -77,7 +76,7 @@ struct State {
 	known_validators: HashMap<PeerId, ValidatorId>,
 
 	/// Use to await for the next validator connection and revoke the request.
-	last_connection_request: Option<ConnectionRequest>,
+	last_connection_request: Option<validator_discovery::ConnectionRequest>,
 }
 
 /// Distribute a collation.
@@ -261,116 +260,17 @@ where
 		request.revoke();
 	}
 
-	let request = connect_to_validators_impl(ctx, relay_parent, validators).await?;
+	let request = validator_discovery::connect_to_validators(
+		ctx,
+		relay_parent,
+		validators,
+	).await?;
 
 	state.last_connection_request = Some(request);
 
 	Ok(())
 }
 
-// TODO (ordian): make this code reusable by other subsystems
-// put in subsystem-util?
-async fn connect_to_validators_impl<Context: SubsystemContext>(
-	ctx: &mut Context,
-	relay_parent: Hash,
-	validators: Vec<ValidatorId>,
-) -> Result<ConnectionRequest> {
-	// ValidatorId -> AuthorityDiscoveryId
-	let (tx, rx) = oneshot::channel();
-
-	ctx.send_message(AllMessages::RuntimeApi(
-		RuntimeApiMessage::Request(
-			relay_parent,
-			RuntimeApiRequest::ValidatorDiscovery(validators.clone(), tx),
-		)
-	)).await?;
-
-	let maybe_authorities = rx.await??;
-	let authorities: Vec<_> = maybe_authorities.iter()
-		.cloned()
-		.filter_map(|id| id)
-		.collect();
-
-	let validator_map = validators.into_iter()
-		.zip(maybe_authorities.into_iter())
-		.filter_map(|(k, v)| v.map(|v| (v, k)))
-		.collect::<HashMap<AuthorityDiscoveryId, ValidatorId>>();
-
-	let (connections, revoke) = connect_to_authorities(ctx, authorities).await?;
-
-	Ok(ConnectionRequest {
-		validator_map,
-		connections,
-		revoke,
-	})
-}
-
-async fn connect_to_authorities<Context: SubsystemContext>(
-	ctx: &mut Context,
-	validator_ids: Vec<AuthorityDiscoveryId>,
-) -> Result<(mpsc::Receiver<(AuthorityDiscoveryId, PeerId)>, oneshot::Sender<()>)> {
-	// do we need a bounded channel at all?
-	const PEERS_CAPACITY: usize = 8;
-
-	let (revoke_tx, revoke) = oneshot::channel();
-	let (connected, connected_rx) = mpsc::channel(PEERS_CAPACITY);
-
-	ctx.send_message(AllMessages::NetworkBridge(
-		NetworkBridgeMessage::ConnectToValidators {
-			validator_ids,
-			connected,
-			revoke,
-		}
-	)).await?;
-
-	Ok((connected_rx, revoke_tx))
-}
-
-#[pin_project]
-struct ConnectionRequest {
-	#[pin]
-	validator_map: HashMap<AuthorityDiscoveryId, ValidatorId>,
-	#[pin]
-	#[must_use = "streams do nothing unless polled"]
-	connections: mpsc::Receiver<(AuthorityDiscoveryId, PeerId)>,
-	#[must_use = "a request should be revoked at some point"]
-	revoke: oneshot::Sender<()>,
-}
-
-impl stream::Stream for ConnectionRequest {
-	type Item = (ValidatorId, PeerId);
-
-	fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
-		let mut this = self.project();
-		if this.validator_map.is_empty() {
-			return Poll::Ready(None);
-		}
-		match this.connections.poll_next(cx) {
-			Poll::Ready(Some((id, peer_id))) => {
-				if let Some(validator_id) = this.validator_map.remove(&id) {
-					return Poll::Ready(Some((validator_id, peer_id)));
-				} else {
-					// unknown authority_id
-					// should be unreachable
-				}
-			}
-			_ => {},
-		}
-		Poll::Pending
-	}
-}
-
-impl ConnectionRequest {
-	pub fn revoke(self) {
-		if let Err(_) = self.revoke.send(()) {
-			warn!(
-				target: TARGET,
-				"Failed to revoke a validator connection request",
-			);
-		}
-	}
-}
-
 /// Advertise collation to a set of relay chain validators.
 async fn advertise_collation<Context>(
 	ctx: &mut Context,
@@ -690,7 +590,7 @@ where
 				Communication { msg } => process_msg(&mut ctx, &mut state, msg).await?,
 				Signal(ActiveLeaves(_update)) => {}
 				Signal(BlockFinalized(_)) => {}
-				Signal(Conclude) => break,	
+				Signal(Conclude) => break,
 			}
 		}
 

node/network/collator-protocol/src/lib.rs

@@ -55,6 +55,16 @@ enum Error {
 	UtilError(util::Error),
 }
 
+impl From<util::validator_discovery::Error> for Error {
+	fn from(me: util::validator_discovery::Error) -> Self {
+		match me {
+			util::validator_discovery::Error::Subsystem(s) => Error::Subsystem(s),
+			util::validator_discovery::Error::RuntimeApi(ra) => Error::RuntimeApi(ra),
+			util::validator_discovery::Error::Oneshot(c) => Error::Oneshot(c),
+		}
+	}
+}
+
 type Result<T> = std::result::Result<T, Error>;
 
 enum ProtocolSide {

node/subsystem-util/src/lib.rs

@@ -55,6 +55,8 @@ use std::{
 };
 use streamunordered::{StreamUnordered, StreamYield};
 
+pub mod validator_discovery;
+
 /// These reexports are required so that external crates can use the `delegated_subsystem` macro properly.
 pub mod reexports {
 	pub use sp_core::traits::SpawnNamed;

node/subsystem-util/src/validator_discovery.rs

@@ -0,0 +1,155 @@
+// Copyright 2020 Parity Technologies (UK) Ltd.
+// This file is part of Polkadot.
+
+// Polkadot is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// Polkadot is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with Polkadot.  If not, see <http://www.gnu.org/licenses/>.
+
+//! Utility function to make it easier to connect to validators.
+
+use std::collections::HashMap;
+use std::pin::Pin;
+
+use futures::{
+	channel::{mpsc, oneshot},
+	task::{Poll, self},
+	stream,
+};
+use pin_project::pin_project;
+
+use polkadot_node_subsystem::{
+	errors::RuntimeApiError, SubsystemError,
+	messages::{AllMessages, RuntimeApiMessage, RuntimeApiRequest, NetworkBridgeMessage},
+	SubsystemContext,
+};
+use polkadot_primitives::v1::{Hash, ValidatorId, AuthorityDiscoveryId};
+use sc_network::PeerId;
+
+/// Error when making a request to connect to validators. 
+#[derive(Debug, derive_more::From)]
+pub enum Error {
+	/// Attempted to send or receive on a oneshot channel which had been canceled
+	#[from]
+	Oneshot(oneshot::Canceled),
+	/// A subsystem error.
+	#[from]
+	Subsystem(SubsystemError),
+	/// An error in the Runtime API.
+	#[from]
+	RuntimeApi(RuntimeApiError),
+}
+
+/// Utility function to make it easier to connect to validators.
+pub async fn connect_to_validators<Context: SubsystemContext>(
+	ctx: &mut Context,
+	relay_parent: Hash,
+	validators: Vec<ValidatorId>,
+) -> Result<ConnectionRequest, Error> {
+	// ValidatorId -> AuthorityDiscoveryId
+	let (tx, rx) = oneshot::channel();
+
+	ctx.send_message(AllMessages::RuntimeApi(
+		RuntimeApiMessage::Request(
+			relay_parent,
+			RuntimeApiRequest::ValidatorDiscovery(validators.clone(), tx),
+		)
+	)).await?;
+
+	let maybe_authorities = rx.await??;
+	let authorities: Vec<_> = maybe_authorities.iter()
+		.cloned()
+		.filter_map(|id| id)
+		.collect();
+
+	let validator_map = validators.into_iter()
+		.zip(maybe_authorities.into_iter())
+		.filter_map(|(k, v)| v.map(|v| (v, k)))
+		.collect::<HashMap<AuthorityDiscoveryId, ValidatorId>>();
+
+	let (connections, revoke) = connect_to_authorities(ctx, authorities).await?;
+
+	Ok(ConnectionRequest {
+		validator_map,
+		connections,
+		revoke,
+	})
+}
+
+async fn connect_to_authorities<Context: SubsystemContext>(
+	ctx: &mut Context,
+	validator_ids: Vec<AuthorityDiscoveryId>,
+) -> Result<(mpsc::Receiver<(AuthorityDiscoveryId, PeerId)>, oneshot::Sender<()>), Error> {
+	const PEERS_CAPACITY: usize = 8;
+
+	let (revoke_tx, revoke) = oneshot::channel();
+	let (connected, connected_rx) = mpsc::channel(PEERS_CAPACITY);
+
+	ctx.send_message(AllMessages::NetworkBridge(
+		NetworkBridgeMessage::ConnectToValidators {
+			validator_ids,
+			connected,
+			revoke,
+		}
+	)).await?;
+
+	Ok((connected_rx, revoke_tx))
+}
+
+/// A pending connection request to validators.
+/// This struct implements `Stream` to allow for asynchronous
+/// discovery of validator addresses.
+///
+/// NOTE: you should call `revoke` on this struct
+/// when you're no longer interested in the requested validators.
+#[pin_project]
+pub struct ConnectionRequest {
+	#[pin]
+	validator_map: HashMap<AuthorityDiscoveryId, ValidatorId>,
+	#[pin]
+	#[must_use = "streams do nothing unless polled"]
+	connections: mpsc::Receiver<(AuthorityDiscoveryId, PeerId)>,
+	#[must_use = "a request should be revoked at some point"]
+	revoke: oneshot::Sender<()>,
+}
+
+impl stream::Stream for ConnectionRequest {
+	type Item = (ValidatorId, PeerId);
+
+	fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
+		let mut this = self.project();
+		if this.validator_map.is_empty() {
+			return Poll::Ready(None);
+		}
+		match this.connections.poll_next(cx) {
+			Poll::Ready(Some((id, peer_id))) => {
+				if let Some(validator_id) = this.validator_map.remove(&id) {
+					return Poll::Ready(Some((validator_id, peer_id)));
+				} else {
+					// unknown authority_id
+					// should be unreachable
+				}
+			}
+			_ => {},
+		}
+		Poll::Pending
+	}
+}
+
+impl ConnectionRequest {
+	pub fn revoke(self) {
+		if let Err(_) = self.revoke.send(()) {
+			log::warn!(
+				"Failed to revoke a validator connection request",
+			);
+		}
+	}
+}