Repository: mintlayer/mintlayer-core
Pull request #110 (2022-02-25 06:31:06)

net: Split NetworkService into two separate services
Split NetworkService into ConnectivityService and FloodsubService
and instead of returning the NetworkService object, return two handles
which can be used to exchange connectivity and floodsub-related events
with the network service provider.

The ultimate goal is separate P2P object into Sync and Swarm-like
objects that handle syncing and connectivity-related events,
respectively. This requires that multiple objects can communicate with
the network service provider and as RefCell/Mutex/RwLock are not optimal
solutions, the cleanest way to implement this is to take advantage of
the frontend/backend split and acquire two handles to the network service
provider and have each handle provide only some of the functionality of
selected network service.

Depending on how the subsystem communication is introduced to P2P,
it may be worthwhile to implement the FloodsubHandle using
tokio::broadcast and provide the handle for both Sync and P2P objects.


Patch:
diff --git a/p2p/src/lib.rs b/p2p/src/lib.rs
index a686bdc16..4ae57588f 100644
--- a/p2p/src/lib.rs
+++ b/p2p/src/lib.rs
@@ -19,7 +19,7 @@
 use crate::{
     error::P2pError,
     event::{Event, PeerEvent, PeerEventType},
-    net::NetworkService,
+    net::{ConnectivityService, FloodsubService, NetworkService},
     peer::{Peer, PeerRole},
 };
 use common::chain::ChainConfig;
@@ -52,12 +52,12 @@ enum PeerState {
 
 #[allow(unused)]
 #[derive(Debug)]
-struct PeerContext<NetworkingBackend>
+struct PeerContext<T>
 where
-    NetworkingBackend: NetworkService,
+    T: NetworkService,
 {
     /// Unique peer ID
-    id: NetworkingBackend::PeerId,
+    id: T::PeerId,
 
     /// Peer state
     state: PeerState,
@@ -90,36 +90,35 @@ where
 }
 
 #[allow(unused)]
-pub struct P2P<NetworkingBackend>
+pub struct P2P<T>
 where
-    NetworkingBackend: NetworkService,
+    T: NetworkService,
 {
     /// Network backend (libp2p, mock)
-    network: NetworkingBackend,
+    network: (T::ConnectivityHandle, T::FloodsubHandle),
 
     /// Chain config
     config: Arc<ChainConfig>,
 
     /// Hashmap for peer information
-    peers: HashMap<NetworkingBackend::PeerId, PeerContext<NetworkingBackend>>,
+    peers: HashMap<T::PeerId, PeerContext<T>>,
 
     /// Hashmap of discovered peers we don't have an active connection with
-    discovered: HashMap<NetworkingBackend::PeerId, PeerAddrInfo<NetworkingBackend>>,
+    discovered: HashMap<T::PeerId, PeerAddrInfo<T>>,
 
     /// Peer backlog maximum size
     peer_backlock: usize,
 
     /// Channel for p2p<->peers communication
-    mgr_chan: (
-        Sender<PeerEvent<NetworkingBackend>>,
-        Receiver<PeerEvent<NetworkingBackend>>,
-    ),
+    mgr_chan: (Sender<PeerEvent<T>>, Receiver<PeerEvent<T>>),
 }
 
 #[allow(unused)]
-impl<NetworkingBackend> P2P<NetworkingBackend>
+impl<T> P2P<T>
 where
-    NetworkingBackend: 'static + NetworkService,
+    T: 'static + NetworkService,
+    T::ConnectivityHandle: ConnectivityService<T>,
+    T::FloodsubHandle: FloodsubService<T>,
 {
     /// Create new P2P
     ///
@@ -128,11 +127,11 @@ where
     pub async fn new(
         mgr_backlog: usize,
         peer_backlock: usize,
-        addr: NetworkingBackend::Address,
+        addr: T::Address,
         config: Arc<ChainConfig>,
     ) -> error::Result<Self> {
         Ok(Self {
-            network: NetworkingBackend::new(addr, &[], &[]).await?,
+            network: T::start(addr, &[], &[]).await?,
             config,
             peer_backlock,
             peers: HashMap::with_capacity(MAX_ACTIVE_CONNECTIONS),
@@ -150,10 +149,7 @@ where
     /// The event is wrapped in an `Option` because the peer might have ungracefully
     /// failed and reading from the closed channel might gives a `None` value, indicating
     /// a protocol error which should be handled accordingly.
-    async fn on_peer_event(
-        &mut self,
-        event: Option<PeerEvent<NetworkingBackend>>,
-    ) -> error::Result<()> {
+    async fn on_peer_event(&mut self, event: Option<PeerEvent<T>>) -> error::Result<()> {
         let event = event.ok_or(P2pError::ChannelClosed)?;
         match event.event {
             PeerEventType::HandshakeFailed => {
@@ -166,7 +162,7 @@ where
             PeerEventType::HandshakeSucceeded => match self.peers.get_mut(&event.peer_id) {
                 Some(peer) => {
                     log::info!("new peer joined, peer id {:?}", event.peer_id);
-                    self.network.register_peer(event.peer_id).await?;
+                    self.network.0.register_peer(event.peer_id).await?;
                     (*peer).state = PeerState::Active;
                     Ok(())
                 }
@@ -183,10 +179,7 @@ where
     /// This may be a socket event (new peer, `accept()` failed) or it may be
     /// a connection request from some other part of the system indicating that
     /// P2P should try to establish a connection with a specific remote peer.
-    async fn on_connectivity_event(
-        &mut self,
-        event: ConnectivityEvent<NetworkingBackend>,
-    ) -> error::Result<()> {
+    async fn on_connectivity_event(&mut self, event: ConnectivityEvent<T>) -> error::Result<()> {
         match event {
             ConnectivityEvent::Accept(peer_id, socket) => {
                 log::debug!("accept incoming connection, peer id {:?}", peer_id);
@@ -198,7 +191,7 @@ where
                     address
                 );
 
-                self.network.connect(address).await.map(|(peer_id, socket)| {
+                self.network.0.connect(address).await.map(|(peer_id, socket)| {
                     self.create_peer(peer_id, socket, PeerRole::Outbound)
                 })?
             }
@@ -237,7 +230,7 @@ where
         let mut iter = self.discovered.iter();
 
         #[allow(clippy::needless_collect)]
-        let peers: Vec<(NetworkingBackend::PeerId, Arc<NetworkingBackend::Address>)> = (0..npeers)
+        let peers: Vec<(T::PeerId, Arc<T::Address>)> = (0..npeers)
             .map(|i| {
                 let peer_info = iter.nth(i).expect("Peer to exist");
 
@@ -268,7 +261,7 @@ where
     }
 
     /// Update the list of peers we know about or update a known peers list of addresses
-    fn peer_discovered(&mut self, peers: &[net::AddrInfo<NetworkingBackend>]) -> error::Result<()> {
+    fn peer_discovered(&mut self, peers: &[net::AddrInfo<T>]) -> error::Result<()> {
         log::info!("discovered {} new peers", peers.len());
 
         for info in peers.iter() {
@@ -292,16 +285,18 @@ where
         Ok(())
     }
 
-    fn peer_expired(&mut self, peers: &[net::AddrInfo<NetworkingBackend>]) -> error::Result<()> {
+    fn peer_expired(&mut self, peers: &[net::AddrInfo<T>]) -> error::Result<()> {
         Ok(())
     }
 
     /// Handle floodsub event
-    fn on_floodsub_event(
-        &mut self,
-        topic: net::FloodsubTopic,
-        message: message::Message,
-    ) -> error::Result<()> {
+    fn on_floodsub_event(&mut self, event: net::FloodsubEvent<T>) -> error::Result<()> {
+        let net::FloodsubEvent::MessageReceived {
+            peer_id: _,
+            topic,
+            message,
+        } = event;
+
         match topic {
             net::FloodsubTopic::Transactions => {
                 log::debug!("received new transaction: {:#?}", message);
@@ -315,22 +310,13 @@ where
     }
 
     /// Handle network event received from the network service provider
-    async fn on_network_event(
-        &mut self,
-        event: net::Event<NetworkingBackend>,
-    ) -> error::Result<()> {
+    async fn on_network_event(&mut self, event: net::ConnectivityEvent<T>) -> error::Result<()> {
         match event {
-            net::Event::Connectivity(event) => match event {
-                net::ConnectivityEvent::IncomingConnection(peer_id, socket) => {
-                    self.on_connectivity_event(ConnectivityEvent::Accept(peer_id, socket)).await
-                }
-                net::ConnectivityEvent::PeerDiscovered(peers) => self.peer_discovered(&peers),
-                net::ConnectivityEvent::PeerExpired(peers) => self.peer_expired(&peers),
-            },
-            net::Event::Floodsub(event) => {
-                let net::FloodsubEvent::MessageReceived(topic, message) = event;
-                self.on_floodsub_event(topic, message)
+            net::ConnectivityEvent::IncomingConnection { peer_id, socket } => {
+                self.on_connectivity_event(ConnectivityEvent::Accept(peer_id, socket)).await
             }
+            net::ConnectivityEvent::PeerDiscovered { peers } => self.peer_discovered(&peers),
+            net::ConnectivityEvent::PeerExpired { peers } => self.peer_expired(&peers),
         }
     }
 
@@ -340,11 +326,16 @@ where
 
         loop {
             tokio::select! {
-                res = self.network.poll_next() => {
+                res = self.network.0.poll_next() => {
                     res.map(|event| async {
                         self.on_network_event(event).await
                     })?;
                 }
+                res = self.network.1.poll_next() => {
+                    res.map(|event| {
+                        self.on_floodsub_event(event)
+                    })?;
+                }
                 event = self.mgr_chan.1.recv().fuse() => {
                     self.on_peer_event(event).await?;
                 }
@@ -353,12 +344,7 @@ where
     }
 
     /// Create `Peer` object from a socket object and spawn task for it
-    fn create_peer(
-        &mut self,
-        id: NetworkingBackend::PeerId,
-        socket: NetworkingBackend::Socket,
-        role: PeerRole,
-    ) {
+    fn create_peer(&mut self, id: T::PeerId, socket: T::Socket, role: PeerRole) {
         let config = self.config.clone();
         let mgr_tx = self.mgr_chan.0.clone();
         let (tx, rx) = tokio::sync::mpsc::channel(self.peer_backlock);
@@ -375,7 +361,7 @@ where
         log::debug!("spawning a task for peer {:?}, role {:?}", id, role);
 
         tokio::spawn(async move {
-            Peer::<NetworkingBackend>::new(id, role, config, socket, mgr_tx, rx).run().await;
+            Peer::<T>::new(id, role, config, socket, mgr_tx, rx).run().await;
         });
     }
 }
diff --git a/p2p/src/net/libp2p/mod.rs b/p2p/src/net/libp2p/mod.rs
index adc56b049..e224a6706 100644
--- a/p2p/src/net/libp2p/mod.rs
+++ b/p2p/src/net/libp2p/mod.rs
@@ -19,7 +19,7 @@
 #![allow(unused)]
 use crate::{
     error::{self, Libp2pError, P2pError},
-    net::{self, Event, FloodsubTopic, NetworkService, SocketService},
+    net::{self, FloodsubTopic, NetworkService, SocketService},
 };
 use async_trait::async_trait;
 use futures::prelude::*;
diff --git a/p2p/src/net/mock/backend.rs b/p2p/src/net/mock/backend.rs
index 5075eb5f8..1d7203a5e 100644
--- a/p2p/src/net/mock/backend.rs
+++ b/p2p/src/net/mock/backend.rs
@@ -18,7 +18,7 @@
 use crate::{
     error::{self, P2pError},
     net::mock::types,
-    net::{Event, FloodsubTopic, NetworkService, SocketService},
+    net::{FloodsubTopic, NetworkService, SocketService},
     peer::Peer,
 };
 use async_trait::async_trait;
@@ -37,6 +37,9 @@ use tokio::{
 };
 
 pub struct Backend {
+    /// Socket address of the backend
+    addr: SocketAddr,
+
     /// Socket for listening to incoming connections
     socket: TcpListener,
 
@@ -52,12 +55,14 @@ pub struct Backend {
 
 impl Backend {
     pub fn new(
+        addr: SocketAddr,
         socket: TcpListener,
         cmd_rx: mpsc::Receiver<types::Command>,
         conn_tx: mpsc::Sender<types::ConnectivityEvent>,
         _flood_tx: mpsc::Sender<types::FloodsubEvent>,
     ) -> Self {
         Self {
+            addr,
             socket,
             cmd_rx,
             conn_tx,
@@ -80,6 +85,11 @@ impl Backend {
                 },
                 event = self.cmd_rx.recv().fuse() => match event.ok_or(P2pError::ChannelClosed)? {
                     types::Command::Connect { addr, response } => {
+                        if self.addr == addr {
+                            let _ = response.send(Err(P2pError::SocketError(ErrorKind::AddrNotAvailable)));
+                            continue;
+                        }
+
                         let _ = match TcpStream::connect(addr).await {
                             Ok(socket) => response.send(Ok(socket)),
                             Err(e) => response.send(Err(e.into())),
diff --git a/p2p/src/net/mock/mod.rs b/p2p/src/net/mock/mod.rs
index a2bbb08c0..a3390738e 100644
--- a/p2p/src/net/mock/mod.rs
+++ b/p2p/src/net/mock/mod.rs
@@ -17,7 +17,10 @@
 #![allow(dead_code, unused_variables, unused_imports)]
 use crate::{
     error::{self, P2pError},
-    net::{ConnectivityEvent, Event, FloodsubEvent, FloodsubTopic, NetworkService, SocketService},
+    net::{
+        ConnectivityEvent, ConnectivityService, FloodsubEvent, FloodsubService, FloodsubTopic,
+        NetworkService, SocketService,
+    },
     peer::Peer,
 };
 use async_trait::async_trait;
@@ -42,19 +45,7 @@ pub mod types;
 pub enum MockStrategy {}
 
 #[derive(Debug)]
-pub struct MockService {
-    /// Address the mock service has been bind to
-    addr: SocketAddr,
-
-    /// TX channel for sending commands to mock backend
-    cmd_tx: mpsc::Sender<types::Command>,
-
-    /// RX channel for receiving connectivity events from mock backend
-    conn_rx: mpsc::Receiver<types::ConnectivityEvent>,
-
-    /// RX channel for receiving floodsub events from mock backend
-    flood_rx: mpsc::Receiver<types::FloodsubEvent>,
-}
+pub struct MockService;
 
 #[derive(Debug)]
 pub struct MockSocket {
@@ -67,46 +58,77 @@ impl MockSocket {
     }
 }
 
+pub struct MockConnectivityHandle<T>
+where
+    T: NetworkService,
+{
+    /// TX channel for sending commands to mock backend
+    cmd_tx: mpsc::Sender<types::Command>,
+
+    /// RX channel for receiving connectivity events from mock backend
+    conn_rx: mpsc::Receiver<types::ConnectivityEvent>,
+    _marker: std::marker::PhantomData<T>,
+}
+
+pub struct MockFloodsubHandle<T>
+where
+    T: NetworkService,
+{
+    /// TX channel for sending commands to mock backend
+    cmd_tx: mpsc::Sender<types::Command>,
+
+    /// RX channel for receiving floodsub events from mock backend
+    _flood_rx: mpsc::Receiver<types::FloodsubEvent>,
+    _marker: std::marker::PhantomData<T>,
+}
+
 #[async_trait]
 impl NetworkService for MockService {
     type Address = SocketAddr;
     type Socket = MockSocket;
     type Strategy = MockStrategy;
     type PeerId = SocketAddr;
+    type ConnectivityHandle = MockConnectivityHandle<Self>;
+    type FloodsubHandle = MockFloodsubHandle<Self>;
 
-    async fn new(
+    async fn start(
         addr: Self::Address,
         _strategies: &[Self::Strategy],
         _topics: &[FloodsubTopic],
-    ) -> error::Result<Self> {
+    ) -> error::Result<(Self::ConnectivityHandle, Self::FloodsubHandle)> {
         let (cmd_tx, cmd_rx) = mpsc::channel(16);
         let (conn_tx, conn_rx) = mpsc::channel(16);
-        let (flood_tx, flood_rx) = mpsc::channel(16);
+        let (flood_tx, _flood_rx) = mpsc::channel(16);
         let socket = TcpListener::bind(addr).await?;
 
         tokio::spawn(async move {
-            let mut mock = backend::Backend::new(socket, cmd_rx, conn_tx, flood_tx);
+            let mut mock = backend::Backend::new(addr, socket, cmd_rx, conn_tx, flood_tx);
             let _ = mock.run().await;
         });
 
-        Ok(Self {
-            addr,
-            cmd_tx,
-            conn_rx,
-            flood_rx,
-        })
+        Ok((
+            Self::ConnectivityHandle {
+                cmd_tx: cmd_tx.clone(),
+                conn_rx,
+                _marker: Default::default(),
+            },
+            Self::FloodsubHandle {
+                cmd_tx,
+                _flood_rx,
+                _marker: Default::default(),
+            },
+        ))
     }
+}
 
-    async fn connect(
-        &mut self,
-        addr: Self::Address,
-    ) -> error::Result<(Self::PeerId, Self::Socket)> {
+#[async_trait]
+impl<T> ConnectivityService<T> for MockConnectivityHandle<T>
+where
+    T: NetworkService<Address = SocketAddr, PeerId = SocketAddr, Socket = MockSocket> + Send,
+{
+    async fn connect(&mut self, addr: T::Address) -> error::Result<(T::PeerId, T::Socket)> {
         log::debug!("try to establish outbound connection, address {:?}", addr);
 
-        if self.addr == addr {
-            return Err(P2pError::SocketError(ErrorKind::AddrNotAvailable));
-        }
-
         let (tx, rx) = oneshot::channel();
         self.cmd_tx.send(types::Command::Connect { addr, response: tx }).await?;
 
@@ -118,37 +140,40 @@ impl NetworkService for MockService {
         Ok((addr, MockSocket::new(socket)))
     }
 
-    async fn poll_next<T>(&mut self) -> error::Result<Event<T>>
-    where
-        T: NetworkService<Socket = MockSocket, PeerId = SocketAddr>,
-    {
-        tokio::select! {
-            event = self.conn_rx.recv().fuse() => match event.ok_or(P2pError::ChannelClosed)? {
-                types::ConnectivityEvent::IncomingConnection { peer_id, socket } => Ok(Event::Connectivity(
-                    ConnectivityEvent::IncomingConnection(peer_id, MockSocket { socket }),
-                )),
-            },
-            event = self.flood_rx.recv().fuse() => match event.ok_or(P2pError::ChannelClosed)? {
-                types::FloodsubEvent::MessageReceived { peer_id: _, topic, message } => {
-                    Ok(Event::Floodsub(FloodsubEvent::MessageReceived(topic, message)))
-                }
+    async fn poll_next(&mut self) -> error::Result<ConnectivityEvent<T>> {
+        match self.conn_rx.recv().await.ok_or(P2pError::ChannelClosed)? {
+            types::ConnectivityEvent::IncomingConnection { peer_id, socket } => {
+                Ok(ConnectivityEvent::IncomingConnection {
+                    peer_id,
+                    socket: MockSocket { socket },
+                })
             }
         }
     }
 
-    async fn publish<T>(&mut self, topic: FloodsubTopic, data: &T) -> error::Result<()>
-    where
-        T: Sync + Send + Encode,
-    {
-        todo!();
+    async fn register_peer(&mut self, peer: T::PeerId) -> error::Result<()> {
+        Ok(())
     }
 
-    async fn register_peer(&mut self, peer: Self::PeerId) -> error::Result<()> {
+    async fn unregister_peer(&mut self, peer: T::PeerId) -> error::Result<()> {
         Ok(())
     }
+}
 
-    async fn unregister_peer(&mut self, peer: Self::PeerId) -> error::Result<()> {
-        Ok(())
+#[async_trait]
+impl<T> FloodsubService<T> for MockFloodsubHandle<T>
+where
+    T: NetworkService<PeerId = SocketAddr> + Send,
+{
+    async fn publish<U>(&mut self, topic: FloodsubTopic, data: &U) -> error::Result<()>
+    where
+        U: Sync + Send + Encode,
+    {
+        todo!();
+    }
+
+    async fn poll_next(&mut self) -> error::Result<FloodsubEvent<T>> {
+        todo!();
     }
 }
 
@@ -185,7 +210,7 @@ impl SocketService for MockSocket {
 mod tests {
     use super::*;
     use crate::net::mock::MockService;
-    use crate::net::{ConnectivityEvent, Event};
+    use crate::net::ConnectivityEvent;
     use crate::peer::{Peer, PeerRole};
     use common::chain::config;
     use parity_scale_codec::{Decode, Encode};
@@ -202,19 +227,19 @@ mod tests {
 
     #[tokio::test]
     async fn test_new() {
-        let srv_ipv4 = MockService::new("127.0.0.1:5555".parse().unwrap(), &[], &[]).await;
+        let srv_ipv4 = MockService::start("127.0.0.1:5555".parse().unwrap(), &[], &[]).await;
         assert!(srv_ipv4.is_ok());
 
         // address already in use
-        let err = MockService::new("127.0.0.1:5555".parse().unwrap(), &[], &[]).await;
+        let err = MockService::start("127.0.0.1:5555".parse().unwrap(), &[], &[]).await;
         assert!(err.is_err());
 
         // bind to IPv6 localhost
-        let srv_ipv6 = MockService::new("[::1]:5555".parse().unwrap(), &[], &[]).await;
+        let srv_ipv6 = MockService::start("[::1]:5555".parse().unwrap(), &[], &[]).await;
         assert!(srv_ipv6.is_ok());
 
         // address already in use
-        let s_ipv6 = MockService::new("[::1]:5555".parse().unwrap(), &[], &[]).await;
+        let s_ipv6 = MockService::start("[::1]:5555".parse().unwrap(), &[], &[]).await;
         assert!(s_ipv6.is_err());
     }
 
@@ -233,12 +258,13 @@ mod tests {
         });
 
         // create service that is used for testing `connect()`
-        let srv = MockService::new("127.0.0.1:7777".parse().unwrap(), &[], &[]).await;
+        let srv = MockService::start("127.0.0.1:7777".parse().unwrap(), &[], &[]).await;
         assert!(srv.is_ok());
-        let mut srv = srv.unwrap();
+        let (mut srv, _) = srv.unwrap();
 
         // try to connect to self, should fail
         let res = srv.connect("127.0.0.1:7777".parse().unwrap()).await;
+        println!("{:?}", res);
         assert!(res.is_err());
 
         // try to connect to an address that (hopefully)
@@ -255,12 +281,13 @@ mod tests {
     async fn test_accept() {
         // create service that is used for testing `accept()`
         let addr: SocketAddr = "[::1]:9999".parse().unwrap();
-        let mut srv = MockService::new("[::1]:9999".parse().unwrap(), &[], &[]).await.unwrap();
+        let (mut srv, _) =
+            MockService::start("[::1]:9999".parse().unwrap(), &[], &[]).await.unwrap();
 
         let (acc, con) = tokio::join!(srv.poll_next(), TcpStream::connect(addr));
         assert!(acc.is_ok());
         assert!(con.is_ok());
-        let acc: Event<MockService> = acc.unwrap();
+        let acc: ConnectivityEvent<MockService> = acc.unwrap();
 
         // TODO: is there any sensible way to make `accept()` fail?
     }
@@ -268,16 +295,16 @@ mod tests {
     #[tokio::test]
     async fn test_peer_send() {
         let addr: SocketAddr = "[::1]:11112".parse().unwrap();
-        let mut server = MockService::new(addr, &[], &[]).await.unwrap();
+        let (mut server, _) = MockService::start(addr, &[], &[]).await.unwrap();
         let remote_fut = TcpStream::connect(addr);
 
         let (server_res, remote_res) = tokio::join!(server.poll_next(), remote_fut);
         assert!(server_res.is_ok());
         assert!(remote_res.is_ok());
 
-        let server_res: Event<MockService> = server_res.unwrap();
+        let server_res: ConnectivityEvent<MockService> = server_res.unwrap();
         let server_res = match server_res {
-            Event::Connectivity(ConnectivityEvent::IncomingConnection(_, socket)) => socket,
+            ConnectivityEvent::IncomingConnection { peer_id: _, socket } => socket,
             _ => panic!("invalid event received, expected incoming connection"),
         };
 
@@ -313,16 +340,16 @@ mod tests {
     async fn test_peer_recv() {
         // create a `MockService`, connect to it with a `TcpStream` and exchange data
         let addr: SocketAddr = "[::1]:11113".parse().unwrap();
-        let mut server = MockService::new(addr, &[], &[]).await.unwrap();
+        let (mut server, _) = MockService::start(addr, &[], &[]).await.unwrap();
         let remote_fut = TcpStream::connect(addr);
 
         let (server_res, remote_res) = tokio::join!(server.poll_next(), remote_fut);
         assert!(server_res.is_ok());
         assert!(remote_res.is_ok());
 
-        let server_res: Event<MockService> = server_res.unwrap();
+        let server_res: ConnectivityEvent<MockService> = server_res.unwrap();
         let server_res = match server_res {
-            Event::Connectivity(ConnectivityEvent::IncomingConnection(_, socket)) => socket,
+            ConnectivityEvent::IncomingConnection { peer_id: _, socket } => socket,
             _ => panic!("invalid event received, expected incoming connection"),
         };
 
diff --git a/p2p/src/net/mod.rs b/p2p/src/net/mod.rs
index 44882f33c..eddd23a35 100644
--- a/p2p/src/net/mod.rs
+++ b/p2p/src/net/mod.rs
@@ -42,27 +42,29 @@ where
     T: NetworkService,
 {
     /// Incoming connection from remote peer
-    IncomingConnection(T::PeerId, T::Socket),
+    IncomingConnection {
+        peer_id: T::PeerId,
+        socket: T::Socket,
+    },
 
     /// One or more peers discovered
-    PeerDiscovered(Vec<AddrInfo<T>>),
+    PeerDiscovered { peers: Vec<AddrInfo<T>> },
 
     /// One one more peers have expired
-    PeerExpired(Vec<AddrInfo<T>>),
+    PeerExpired { peers: Vec<AddrInfo<T>> },
 }
 
 // TODO: separate events for blocks and transactions?
-pub enum FloodsubEvent {
-    /// Message received from a Floodsub topic
-    MessageReceived(FloodsubTopic, message::Message),
-}
-
-pub enum Event<T>
+pub enum FloodsubEvent<T>
 where
     T: NetworkService,
 {
-    Connectivity(ConnectivityEvent<T>),
-    Floodsub(FloodsubEvent),
+    /// Message received from a Floodsub topic
+    MessageReceived {
+        peer_id: T::PeerId,
+        topic: FloodsubTopic,
+        message: message::Message,
+    },
 }
 
 #[derive(Debug, PartialEq, Eq, Hash)]
@@ -94,20 +96,32 @@ pub trait NetworkService {
     // Enum of different peer discovery strategies that the implementation provides
     type Strategy;
 
+    /// Handle for sending/receiving connecitivity-related events
+    type ConnectivityHandle: Send;
+
+    /// Handle for sending/receiving floodsub-related events
+    type FloodsubHandle: Send;
+
     /// Initialize the network service provider
     ///
     /// # Arguments
     /// `addr` - socket address for incoming P2P traffic
     /// `strategies` - list of strategies that are used for peer discovery
     /// `topics` - list of floodsub topics that the implementation should subscribe to
-    async fn new(
+    async fn start(
         addr: Self::Address,
         strategies: &[Self::Strategy],
         topics: &[FloodsubTopic],
-    ) -> error::Result<Self>
-    where
-        Self: Sized;
+    ) -> error::Result<(Self::ConnectivityHandle, Self::FloodsubHandle)>;
+}
 
+/// ConnectivityService provides an interface through which objects can send
+/// and receive connectivity-related events to/from the network service provider
+#[async_trait]
+pub trait ConnectivityService<T>
+where
+    T: NetworkService,
+{
     /// Connect to a remote node
     ///
     /// If the connection succeeds, the socket object is returned
@@ -115,33 +129,41 @@ pub trait NetworkService {
     ///
     /// # Arguments
     /// `addr` - socket address of the peer
-    async fn connect(&mut self, addr: Self::Address)
-        -> error::Result<(Self::PeerId, Self::Socket)>;
+    async fn connect(&mut self, address: T::Address) -> error::Result<(T::PeerId, T::Socket)>;
 
     /// Poll events from the network service provider
     ///
     /// There are three types of events that can be received:
     /// - incoming peer connections
-    /// - incoming messages from floodsub topics
     /// - new discovered peers
-    async fn poll_next<T>(&mut self) -> error::Result<Event<T>>
-    where
-        T: NetworkService<Socket = Self::Socket, Address = Self::Address, PeerId = Self::PeerId>;
+    /// - peer expiration events
+    async fn poll_next(&mut self) -> error::Result<ConnectivityEvent<T>>;
 
+    /// Register peer to the network service provider
+    async fn register_peer(&mut self, peer: T::PeerId) -> error::Result<()>;
+
+    /// Unregister peer from the network service provider
+    async fn unregister_peer(&mut self, peer: T::PeerId) -> error::Result<()>;
+}
+
+/// FloodsubService provides an interface through which objects can send
+/// and receive floodsub-related events to/from the network service provider
+#[async_trait]
+pub trait FloodsubService<T>
+where
+    T: NetworkService,
+{
     /// Publish data in a given floodsub topic
     ///
     /// # Arguments
     /// `topic` - identifier for the topic
     /// `data` - generic data to send
-    async fn publish<T>(&mut self, topic: FloodsubTopic, data: &T) -> error::Result<()>
+    async fn publish<U>(&mut self, topic: FloodsubTopic, data: &U) -> error::Result<()>
     where
-        T: Sync + Send + Encode;
+        U: Sync + Send + Encode;
 
-    /// Register peer to the network service provider
-    async fn register_peer(&mut self, peer: Self::PeerId) -> error::Result<()>;
-
-    /// Unregister peer from the network service provider
-    async fn unregister_peer(&mut self, peer: Self::PeerId) -> error::Result<()>;
+    /// Poll floodsub-related events from the network service provider
+    async fn poll_next(&mut self) -> error::Result<FloodsubEvent<T>>;
 }
 
 /// `SocketService` provides the low-level socket interface that


Test patch:
diff --git a/p2p/test-utils/src/lib.rs b/p2p/test-utils/src/lib.rs
index 394064dee..264925a6b 100644
--- a/p2p/test-utils/src/lib.rs
+++ b/p2p/test-utils/src/lib.rs
@@ -23,7 +23,7 @@ use p2p::{
         // libp2p::Libp2pService,
         mock::{MockService, MockSocket},
         ConnectivityEvent,
-        Event,
+        ConnectivityService,
         NetworkService,
     },
     peer::*,
@@ -67,13 +67,13 @@ pub async fn create_two_mock_peers(
     config: Arc<ChainConfig>,
 ) -> (Peer<MockService>, Peer<MockService>) {
     let addr: SocketAddr = make_address("[::1]:");
-    let mut server = MockService::new(addr, &[], &[]).await.unwrap();
+    let (mut server, _) = MockService::start(addr, &[], &[]).await.unwrap();
     let peer_fut = TcpStream::connect(addr);
 
     let (remote_res, local_res) = tokio::join!(server.poll_next(), peer_fut);
-    let remote_res: Event<MockService> = remote_res.unwrap();
+    let remote_res: ConnectivityEvent<MockService> = remote_res.unwrap();
     let remote_res = match remote_res {
-        Event::Connectivity(ConnectivityEvent::IncomingConnection(_, socket)) => socket,
+        ConnectivityEvent::IncomingConnection { peer_id: _, socket } => socket,
         _ => panic!("invalid event received, expected incoming connection"),
     };
     let local_res = local_res.unwrap();
diff --git a/p2p/tests/peer.rs b/p2p/tests/peer.rs
index 8185feb1a..3045849c6 100644
--- a/p2p/tests/peer.rs
+++ b/p2p/tests/peer.rs
@@ -19,7 +19,9 @@ extern crate test_utils;
 
 use common::chain::config;
 // use libp2p::Multiaddr;
-use p2p::net::{self, /*libp2p::Libp2pService, */ mock::MockService, NetworkService};
+use p2p::net::{
+    self, /*libp2p::Libp2pService, */ mock::MockService, ConnectivityService, NetworkService,
+};
 use p2p::peer::{Peer, PeerRole};
 use std::net::SocketAddr;
 use std::sync::Arc;
@@ -30,16 +32,16 @@ use tokio::net::TcpStream;
 async fn test_peer_new_mock() {
     let config = Arc::new(config::create_mainnet());
     let addr: SocketAddr = test_utils::make_address("[::1]:");
-    let mut server = MockService::new(addr, &[], &[]).await.unwrap();
+    let (mut server, _) = MockService::start(addr, &[], &[]).await.unwrap();
     let peer_fut = TcpStream::connect(addr);
 
     let (server_res, peer_res) = tokio::join!(server.poll_next(), peer_fut);
     assert!(server_res.is_ok());
     assert!(peer_res.is_ok());
 
-    let server_res: net::Event<MockService> = server_res.unwrap();
+    let server_res: net::ConnectivityEvent<MockService> = server_res.unwrap();
     let server_res = match server_res {
-        net::Event::Connectivity(net::ConnectivityEvent::IncomingConnection(_, socket)) => socket,
+        net::ConnectivityEvent::IncomingConnection { peer_id: _, socket } => socket,
         _ => panic!("invalid event received, expected incoming connection"),
     };