A0-1761: Enhance gossip network interface

finality-aleph/src/network/gossip/mod.rs

@@ -1,5 +1,6 @@
 //! A P2P-based gossip network, for now only for sending broadcasts.
 use std::{
+    collections::HashSet,
     fmt::{Debug, Display},
     hash::Hash,
 };
@@ -15,18 +16,34 @@ mod service;
 pub use service::Service;
 
 #[async_trait::async_trait]
-/// Interface for the gossip network, currently only supports broadcasting and receiving data.
+/// Interface for the gossip network. This represents a P2P network and a lot of the properties of
+/// this interface result from that. In particular we might know the ID of a given peer, but not be
+/// connected to them directly.
 pub trait Network<D: Data>: Send + 'static {
     type Error: Display + Send;
+    type PeerId: Clone + Debug + Eq + Hash + Send + 'static;
+
+    /// Attempt to send data to a peer. Might silently fail if we are not connected to them.
+    fn send(&mut self, data: D, peer_id: Self::PeerId) -> Result<(), Self::Error>;
+
+    /// Send data to a random peer, preferably from a list. It should send the data to a randomly
+    /// chosen peer from the provided list, but if it cannot (e.g. because it's not connected) it
+    /// will send to a random available peer. No guarantees any peer gets it even if no errors are
+    /// returned, retry appropriately.
+    fn send_to_random(
+        &mut self,
+        data: D,
+        peer_ids: HashSet<Self::PeerId>,
+    ) -> Result<(), Self::Error>;
 
     /// Broadcast data to all directly connected peers. Network-wide broadcasts have to be
     /// implemented on top of this abstraction. Note that there might be no currently connected
     /// peers, so there are no guarantees any single call sends anything even if no errors are
     /// returned, retry appropriately.
     fn broadcast(&mut self, data: D) -> Result<(), Self::Error>;
 
-    /// Receive some data from the network.
-    async fn next(&mut self) -> Result<D, Self::Error>;
+    /// Receive some data from the network, including information about who sent it.
+    async fn next(&mut self) -> Result<(D, Self::PeerId), Self::Error>;
 }
 
 /// The Authentication protocol is used for validator discovery.
@@ -51,7 +68,7 @@ pub trait NetworkSender: Send + Sync + 'static {
 pub enum Event<P> {
     StreamOpened(P, Protocol),
     StreamClosed(P, Protocol),
-    Messages(Vec<(Protocol, Bytes)>),
+    Messages(P, Vec<(Protocol, Bytes)>),
 }
 
 #[async_trait::async_trait]
@@ -63,7 +80,7 @@ pub trait EventStream<P> {
 pub trait RawNetwork: Clone + Send + Sync + 'static {
     type SenderError: std::error::Error;
     type NetworkSender: NetworkSender;
-    type PeerId: Clone + Debug + Eq + Hash + Send;
+    type PeerId: Clone + Debug + Eq + Hash + Send + 'static;
     type EventStream: EventStream<Self::PeerId>;
 
     /// Returns a stream of events representing what happens on the network.

finality-aleph/src/network/gossip/service.rs

@@ -1,11 +1,13 @@
 use std::{
     collections::{HashMap, HashSet},
-    fmt::{Display, Error as FmtError, Formatter},
+    fmt::{Debug, Display, Error as FmtError, Formatter},
     future::Future,
+    hash::Hash,
 };
 
 use futures::{channel::mpsc, StreamExt};
 use log::{debug, error, info, trace, warn};
+use rand::{seq::IteratorRandom, thread_rng};
 use sc_service::SpawnTaskHandle;
 use sc_utils::mpsc::{tracing_unbounded, TracingUnboundedReceiver, TracingUnboundedSender};
 use tokio::time;
@@ -18,6 +20,12 @@ use crate::{
     STATUS_REPORT_INTERVAL,
 };
 
+enum Command<D: Data, P: Clone + Debug + Eq + Hash + Send + 'static> {
+    Send(D, P),
+    SendToRandom(D, HashSet<P>),
+    Broadcast(D),
+}
+
 /// A service managing all the direct interaction with the underlying network implementation. It
 /// handles:
 /// 1. Incoming network events
@@ -26,16 +34,16 @@ use crate::{
 /// 3. Outgoing messages, sending them out, using 1.2. to broadcast.
 pub struct Service<N: RawNetwork, D: Data> {
     network: N,
-    messages_from_user: mpsc::UnboundedReceiver<D>,
-    messages_for_user: mpsc::UnboundedSender<D>,
+    messages_from_user: mpsc::UnboundedReceiver<Command<D, N::PeerId>>,
+    messages_for_user: mpsc::UnboundedSender<(D, N::PeerId)>,
     authentication_connected_peers: HashSet<N::PeerId>,
     authentication_peer_senders: HashMap<N::PeerId, TracingUnboundedSender<D>>,
     spawn_handle: SpawnTaskHandle,
 }
 
-struct ServiceInterface<D: Data> {
-    messages_from_service: mpsc::UnboundedReceiver<D>,
-    messages_for_service: mpsc::UnboundedSender<D>,
+struct ServiceInterface<D: Data, P: Clone + Debug + Eq + Hash + Send + 'static> {
+    messages_from_service: mpsc::UnboundedReceiver<(D, P)>,
+    messages_for_service: mpsc::UnboundedSender<Command<D, P>>,
 }
 
 /// What can go wrong when receiving or sending data.
@@ -56,16 +64,33 @@ impl Display for Error {
 }
 
 #[async_trait::async_trait]
-impl<D: Data> Network<D> for ServiceInterface<D> {
+impl<D: Data, P: Clone + Debug + Eq + Hash + Send + 'static> Network<D> for ServiceInterface<D, P> {
     type Error = Error;
+    type PeerId = P;
+
+    fn send(&mut self, data: D, peer_id: Self::PeerId) -> Result<(), Self::Error> {
+        self.messages_for_service
+            .unbounded_send(Command::Send(data, peer_id))
+            .map_err(|_| Error::ServiceStopped)
+    }
+
+    fn send_to_random(
+        &mut self,
+        data: D,
+        peer_ids: HashSet<Self::PeerId>,
+    ) -> Result<(), Self::Error> {
+        self.messages_for_service
+            .unbounded_send(Command::SendToRandom(data, peer_ids))
+            .map_err(|_| Error::ServiceStopped)
+    }
 
     fn broadcast(&mut self, data: D) -> Result<(), Self::Error> {
         self.messages_for_service
-            .unbounded_send(data)
+            .unbounded_send(Command::Broadcast(data))
             .map_err(|_| Error::ServiceStopped)
     }
 
-    async fn next(&mut self) -> Result<D, Self::Error> {
+    async fn next(&mut self) -> Result<(D, Self::PeerId), Self::Error> {
         self.messages_from_service
             .next()
             .await
@@ -83,7 +108,10 @@ impl<N: RawNetwork, D: Data> Service<N, D> {
     pub fn new(
         network: N,
         spawn_handle: SpawnTaskHandle,
-    ) -> (Service<N, D>, impl Network<D, Error = Error>) {
+    ) -> (
+        Service<N, D>,
+        impl Network<D, Error = Error, PeerId = N::PeerId>,
+    ) {
         let (messages_for_user, messages_from_service) = mpsc::unbounded();
         let (messages_for_service, messages_from_user) = mpsc::unbounded();
         (
@@ -170,21 +198,55 @@ impl<N: RawNetwork, D: Data> Service<N, D> {
         }
     }
 
-    fn broadcast(&mut self, data: D, protocol: Protocol) {
-        let peers = match protocol {
-            Protocol::Authentication => self.authentication_connected_peers.clone(),
+    fn send(&mut self, data: D, peer_id: N::PeerId, protocol: Protocol) {
+        if let Err(e) = self.send_to_peer(data, peer_id.clone(), protocol) {
+            trace!(target: "aleph-network", "Failed to send to peer{:?}, {:?}", peer_id, e);
+        }
+    }
+
+    fn protocol_peers(&self, protocol: Protocol) -> &HashSet<N::PeerId> {
+        match protocol {
+            Protocol::Authentication => &self.authentication_connected_peers,
+        }
+    }
+
+    fn random_peer<'a>(
+        &'a self,
+        peer_ids: &'a HashSet<N::PeerId>,
+        protocol: Protocol,
+    ) -> Option<&'a N::PeerId> {
+        peer_ids
+            .intersection(self.protocol_peers(protocol))
+            .into_iter()
+            .choose(&mut thread_rng())
+            .or(self
+                .protocol_peers(protocol)
+                .iter()
+                .choose(&mut thread_rng()))
+    }
+
+    fn send_to_random(&mut self, data: D, peer_ids: HashSet<N::PeerId>, protocol: Protocol) {
+        let peer_id = match self.random_peer(&peer_ids, protocol) {
+            Some(peer_id) => peer_id.clone(),
+            None => {
+                trace!(target: "aleph-network", "Failed to send to random peer, no peers are available.");
+                return;
+            }
         };
+        self.send(data, peer_id, protocol);
+    }
+
+    fn broadcast(&mut self, data: D, protocol: Protocol) {
+        let peers = self.protocol_peers(protocol).clone();
         for peer in peers {
-            if let Err(e) = self.send_to_peer(data.clone(), peer.clone(), protocol) {
-                trace!(target: "aleph-network", "Failed to send broadcast to peer{:?}, {:?}", peer, e);
-            }
+            self.send(data.clone(), peer, protocol);
         }
     }
 
     fn handle_network_event(
         &mut self,
         event: Event<N::PeerId>,
-    ) -> Result<(), mpsc::TrySendError<D>> {
+    ) -> Result<(), mpsc::TrySendError<(D, N::PeerId)>> {
         use Event::*;
         match event {
             StreamOpened(peer, protocol) => {
@@ -212,11 +274,13 @@ impl<N: RawNetwork, D: Data> Service<N, D> {
                     }
                 }
             }
-            Messages(messages) => {
+            Messages(peer_id, messages) => {
                 for (protocol, data) in messages.into_iter() {
                     match protocol {
                         Protocol::Authentication => match D::decode(&mut &data[..]) {
-                            Ok(data) => self.messages_for_user.unbounded_send(data)?,
+                            Ok(data) => self
+                                .messages_for_user
+                                .unbounded_send((data, peer_id.clone()))?,
                             Err(e) => {
                                 warn!(target: "aleph-network", "Error decoding authentication protocol message: {}", e)
                             }
@@ -256,7 +320,9 @@ impl<N: RawNetwork, D: Data> Service<N, D> {
                     }
                 },
                 maybe_message = self.messages_from_user.next() => match maybe_message {
-                    Some(message) => self.broadcast(message, Protocol::Authentication),
+                    Some(Command::Broadcast(message)) => self.broadcast(message, Protocol::Authentication),
+                    Some(Command::SendToRandom(message, peer_ids)) => self.send_to_random(message, peer_ids, Protocol::Authentication),
+                    Some(Command::Send(message, peer_id)) => self.send(message, peer_id, Protocol::Authentication),
                     None => {
                         error!(target: "aleph-network", "User message stream ended.");
                         return;
@@ -281,7 +347,7 @@ mod tests {
 
     use super::{Error, Service};
     use crate::network::{
-        clique::mock::random_peer_id,
+        clique::mock::{random_peer_id, MockPublicKey},
         gossip::{
             mock::{MockEvent, MockRawNetwork, MockSenderError},
             Network,
@@ -294,7 +360,7 @@ mod tests {
 
     pub struct TestData {
         pub network: MockRawNetwork,
-        gossip_network: Box<dyn Network<MockData, Error = Error>>,
+        gossip_network: Box<dyn Network<MockData, Error = Error, PeerId = MockPublicKey>>,
         pub service: Service<MockRawNetwork, MockData>,
         // `TaskManager` can't be dropped for `SpawnTaskHandle` to work
         _task_manager: TaskManager,
@@ -330,12 +396,25 @@ mod tests {
     #[async_trait::async_trait]
     impl Network<MockData> for TestData {
         type Error = Error;
+        type PeerId = MockPublicKey;
+
+        fn send(&mut self, data: MockData, peer_id: Self::PeerId) -> Result<(), Self::Error> {
+            self.gossip_network.send(data, peer_id)
+        }
+
+        fn send_to_random(
+            &mut self,
+            data: MockData,
+            peer_ids: HashSet<Self::PeerId>,
+        ) -> Result<(), Self::Error> {
+            self.gossip_network.send_to_random(data, peer_ids)
+        }
 
         fn broadcast(&mut self, data: MockData) -> Result<(), Self::Error> {
             self.gossip_network.broadcast(data)
         }
 
-        async fn next(&mut self) -> Result<MockData, Self::Error> {
+        async fn next(&mut self) -> Result<(MockData, Self::PeerId), Self::Error> {
             self.gossip_network.next().await
         }
     }
@@ -514,16 +593,14 @@ mod tests {
 
         test_data
             .service
-            .handle_network_event(MockEvent::Messages(vec![(
-                PROTOCOL,
-                message.clone().encode().into(),
-            )]))
+            .handle_network_event(MockEvent::Messages(
+                random_peer_id(),
+                vec![(PROTOCOL, message.clone().encode().into())],
+            ))
             .expect("Should handle");
 
-        assert_eq!(
-            test_data.next().await.expect("Should receive message"),
-            message,
-        );
+        let (received_message, _) = test_data.next().await.expect("Should receive message");
+        assert_eq!(received_message, message);
 
         test_data.cleanup().await
     }

finality-aleph/src/network/session/service.rs

@@ -385,7 +385,7 @@ where
                     }
                 },
                 maybe_authentication = self.gossip_network.next() => {
-                    let authentication = maybe_authentication.map_err(Error::GossipNetwork)?;
+                    let (authentication, _) = maybe_authentication.map_err(Error::GossipNetwork)?;
                     trace!(target: "aleph-network", "Manager received an authentication from network");
                     match authentication.try_into() {
                         Ok(message) => {

finality-aleph/src/network/substrate.rs

@@ -171,8 +171,9 @@ impl<B: Block, H: ExHashT> EventStream<PeerId> for NetworkEventStream<B, H> {
                             Err(_) => continue,
                         }
                     }
-                    NotificationsReceived { messages, .. } => {
+                    NotificationsReceived { messages, remote } => {
                         return Some(Messages(
+                            remote,
                             messages
                                 .into_iter()
                                 .filter_map(|(protocol, data)| {

finality-aleph/src/testing/network.rs

@@ -227,10 +227,13 @@ impl TestData {
             for versioned_authentication in
                 Vec::<VersionedAuthentication<_, _>>::from(handler.authentication().unwrap())
             {
-                self.network.emit_event(MockEvent::Messages(vec![(
-                    Protocol::Authentication,
-                    versioned_authentication.encode().into(),
-                )]));
+                self.network.emit_event(MockEvent::Messages(
+                    authority.auth_peer_id(),
+                    vec![(
+                        Protocol::Authentication,
+                        versioned_authentication.encode().into(),
+                    )],
+                ));
             }
         }
     }
@@ -334,10 +337,13 @@ async fn test_forwards_authentication_broadcast() {
     for versioned_authentication in
         Vec::<VersionedAuthentication<_, _>>::from(sending_peer_handler.authentication().unwrap())
     {
-        test_data.network.emit_event(MockEvent::Messages(vec![(
-            Protocol::Authentication,
-            versioned_authentication.encode().into(),
-        )]));
+        test_data.network.emit_event(MockEvent::Messages(
+            sending_peer.auth_peer_id(),
+            vec![(
+                Protocol::Authentication,
+                versioned_authentication.encode().into(),
+            )],
+        ));
     }
 
     for _ in 0..2 {