Merge pull request #104 from zeromq/alex/monitor-sockets

Initial version of socket monitor impl. Related to #103
zeromq · Dec 10, 2020 · f5f53a6 · f5f53a6
2 parents 8e75a13 + ae103f3
commit f5f53a6
Show file tree

Hide file tree

Showing 12 changed files with 178 additions and 32 deletions.
diff --git a/src/backend.rs b/src/backend.rs
@@ -1,9 +1,10 @@
 use crate::codec::{FramedIo, Message, ZmqFramedRead, ZmqFramedWrite};
 use crate::fair_queue::QueueInner;
 use crate::util::PeerIdentity;
-use crate::{MultiPeerBackend, SocketBackend, SocketType, ZmqError, ZmqResult};
+use crate::{MultiPeerBackend, SocketBackend, SocketEvent, SocketType, ZmqError, ZmqResult};
 use crossbeam::queue::SegQueue;
 use dashmap::DashMap;
+use futures::channel::mpsc;
 use futures::SinkExt;
 use parking_lot::Mutex;
 use std::sync::Arc;
@@ -17,6 +18,7 @@ pub(crate) struct GenericSocketBackend {
     fair_queue_inner: Option<Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>>,
     pub(crate) round_robin: SegQueue<PeerIdentity>,
     socket_type: SocketType,
+    pub(crate) socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
 }
 
 impl GenericSocketBackend {
@@ -29,6 +31,7 @@ impl GenericSocketBackend {
             fair_queue_inner,
             round_robin: SegQueue::new(),
             socket_type,
+            socket_monitor: Mutex::new(None),
         }
     }
 
@@ -85,6 +88,10 @@ impl SocketBackend for GenericSocketBackend {
     fn shutdown(&self) {
         self.peers.clear();
     }
+
+    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
+        &self.socket_monitor
+    }
 }
 
 impl MultiPeerBackend for GenericSocketBackend {

diff --git a/src/dealer.rs b/src/dealer.rs
@@ -3,8 +3,12 @@ use crate::codec::{Message, ZmqFramedRead};
 use crate::fair_queue::FairQueue;
 use crate::transport::AcceptStopHandle;
 use crate::util::PeerIdentity;
-use crate::{Endpoint, MultiPeerBackend, Socket, SocketBackend, SocketType, ZmqMessage, ZmqResult};
+use crate::{
+    Endpoint, MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketType, ZmqMessage,
+    ZmqResult,
+};
 use async_trait::async_trait;
+use futures::channel::mpsc;
 use futures::StreamExt;
 use std::collections::hash_map::RandomState;
 use std::collections::HashMap;
@@ -42,6 +46,12 @@ impl Socket for DealerSocket {
     fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle, RandomState> {
         &mut self.binds
     }
+
+    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
+        let (sender, receiver) = mpsc::channel(1024);
+        self.backend.socket_monitor.lock().replace(sender);
+        receiver
+    }
 }
 
 impl DealerSocket {

diff --git a/src/lib.rs b/src/lib.rs
@@ -38,8 +38,10 @@ use util::PeerIdentity;
 extern crate enum_primitive_derive;
 
 use async_trait::async_trait;
+use futures::channel::mpsc;
 use futures_codec::FramedWrite;
 use num_traits::ToPrimitive;
+use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::convert::TryFrom;
 use std::fmt::{Debug, Display};
@@ -102,6 +104,19 @@ impl Display for SocketType {
     }
 }
 
+#[derive(Debug)]
+pub enum SocketEvent {
+    Connected(Endpoint, PeerIdentity),
+    ConnectDelayed,
+    ConnectRetried,
+    Listening(Endpoint),
+    Accepted(Endpoint, PeerIdentity),
+    AcceptFailed(ZmqError),
+    Closed,
+    CloseFailed,
+    Disconnected(PeerIdentity),
+}
+
 pub trait MultiPeerBackend: SocketBackend {
     /// This should not be public..
     /// Find a better way of doing this
@@ -113,6 +128,7 @@ pub trait MultiPeerBackend: SocketBackend {
 pub trait SocketBackend: Send + Sync {
     fn socket_type(&self) -> SocketType;
     fn shutdown(&self);
+    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>>;
 }
 
 #[async_trait]
@@ -148,9 +164,39 @@ pub trait Socket: Sized + Send {
         let endpoint = endpoint.try_into()?;
 
         let cloned_backend = self.backend();
-        let cback = move |result| util::peer_connected(result, cloned_backend.clone());
+        let cback = move |result| {
+            let cloned_backend = cloned_backend.clone();
+            async move {
+                let result = match result {
+                    Ok((socket, endpoint)) => {
+                        match util::peer_connected(socket, cloned_backend.clone()).await {
+                            Ok(peer_id) => Ok((endpoint, peer_id)),
+                            Err(e) => Err(e),
+                        }
+                    }
+                    Err(e) => Err(e),
+                };
+                match result {
+                    Ok((endpoint, peer_id)) => {
+                        if let Some(monitor) = cloned_backend.monitor().lock().as_mut() {
+                            let _ = monitor.try_send(SocketEvent::Accepted(endpoint, peer_id));
+                        }
+                    }
+                    Err(e) => {
+                        if let Some(monitor) = cloned_backend.monitor().lock().as_mut() {
+                            let _ = monitor.try_send(SocketEvent::AcceptFailed(e));
+                        }
+                    }
+                }
+            }
+        };
+
         let (endpoint, stop_handle) = transport::begin_accept(endpoint, cback).await?;
 
+        if let Some(monitor) = self.backend().monitor().lock().as_mut() {
+            let _ = monitor.try_send(SocketEvent::Listening(endpoint.clone()));
+        }
+
         self.binds().insert(endpoint.clone(), stop_handle);
         Ok(endpoint)
     }
@@ -190,11 +236,30 @@ pub trait Socket: Sized + Send {
         let backend = self.backend();
         let endpoint = endpoint.try_into()?;
 
-        let connect_result = transport::connect(endpoint).await;
-        util::peer_connected(connect_result, backend).await;
-        Ok(())
+        let result = match transport::connect(endpoint).await {
+            Ok((socket, endpoint)) => match util::peer_connected(socket, backend).await {
+                Ok(peer_id) => Ok((endpoint, peer_id)),
+                Err(e) => Err(e),
+            },
+            Err(e) => Err(e),
+        };
+        match result {
+            Ok((endpoint, peer_id)) => {
+                if let Some(monitor) = self.backend().monitor().lock().as_mut() {
+                    let _ = monitor.try_send(SocketEvent::Connected(endpoint, peer_id));
+                }
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
     }
 
+    /// Creates and setups new socket monitor
+    ///
+    /// Subsequent calls to this method each create a new monitor channel.
+    /// Sender side of previous one is dropped.
+    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent>;
+
     // TODO: async fn connections(&self) -> ?
 
     /// Disconnects from the given endpoint, blocking until finished.

diff --git a/src/pub.rs b/src/pub.rs
@@ -4,11 +4,14 @@ use crate::error::ZmqResult;
 use crate::message::*;
 use crate::transport::AcceptStopHandle;
 use crate::util::PeerIdentity;
-use crate::{BlockingSend, MultiPeerBackend, Socket, SocketBackend, SocketType, ZmqError};
-use futures::channel::oneshot;
+use crate::{
+    BlockingSend, MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketType, ZmqError,
+};
+use futures::channel::{mpsc, oneshot};
 
 use async_trait::async_trait;
 use dashmap::DashMap;
+use parking_lot::Mutex;
 use std::collections::HashMap;
 use std::io::ErrorKind;
 use std::pin::Pin;
@@ -22,6 +25,7 @@ pub(crate) struct Subscriber {
 
 pub(crate) struct PubSocketBackend {
     subscribers: DashMap<PeerIdentity, Subscriber>,
+    socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
 }
 
 impl PubSocketBackend {
@@ -81,6 +85,10 @@ impl SocketBackend for PubSocketBackend {
     fn shutdown(&self) {
         self.subscribers.clear();
     }
+
+    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
+        &self.socket_monitor
+    }
 }
 
 impl MultiPeerBackend for PubSocketBackend {
@@ -184,6 +192,7 @@ impl Socket for PubSocket {
         Self {
             backend: Arc::new(PubSocketBackend {
                 subscribers: DashMap::new(),
+                socket_monitor: Mutex::new(None),
             }),
             binds: HashMap::new(),
         }
@@ -196,6 +205,12 @@ impl Socket for PubSocket {
     fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
         &mut self.binds
     }
+
+    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
+        let (sender, receiver) = mpsc::channel(1024);
+        self.backend.socket_monitor.lock().replace(sender);
+        receiver
+    }
 }
 
 #[cfg(test)]

diff --git a/src/pull.rs b/src/pull.rs
@@ -3,8 +3,12 @@ use crate::codec::{Message, ZmqFramedRead};
 use crate::fair_queue::FairQueue;
 use crate::transport::AcceptStopHandle;
 use crate::util::PeerIdentity;
-use crate::{BlockingRecv, Endpoint, MultiPeerBackend, Socket, SocketType, ZmqMessage, ZmqResult};
+use crate::{
+    BlockingRecv, Endpoint, MultiPeerBackend, Socket, SocketEvent, SocketType, ZmqMessage,
+    ZmqResult,
+};
 use async_trait::async_trait;
+use futures::channel::mpsc;
 use futures::StreamExt;
 use std::collections::hash_map::RandomState;
 use std::collections::HashMap;
@@ -37,6 +41,12 @@ impl Socket for PullSocket {
     fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle, RandomState> {
         &mut self.binds
     }
+
+    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
+        let (sender, receiver) = mpsc::channel(1024);
+        self.backend.socket_monitor.lock().replace(sender);
+        receiver
+    }
 }
 
 #[async_trait]

diff --git a/src/push.rs b/src/push.rs
@@ -2,10 +2,11 @@ use crate::backend::GenericSocketBackend;
 use crate::codec::Message;
 use crate::transport::AcceptStopHandle;
 use crate::{
-    BlockingSend, Endpoint, MultiPeerBackend, Socket, SocketBackend, SocketType, ZmqMessage,
-    ZmqResult,
+    BlockingSend, Endpoint, MultiPeerBackend, Socket, SocketBackend, SocketEvent, SocketType,
+    ZmqMessage, ZmqResult,
 };
 use async_trait::async_trait;
+use futures::channel::mpsc;
 use std::collections::hash_map::RandomState;
 use std::collections::HashMap;
 use std::sync::Arc;
@@ -36,6 +37,12 @@ impl Socket for PushSocket {
     fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle, RandomState> {
         &mut self.binds
     }
+
+    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
+        let (sender, receiver) = mpsc::channel(1024);
+        self.backend.socket_monitor.lock().replace(sender);
+        receiver
+    }
 }
 
 #[async_trait]

diff --git a/src/rep.rs b/src/rep.rs
@@ -21,6 +21,7 @@ struct RepPeer {
 struct RepSocketBackend {
     pub(crate) peers: DashMap<PeerIdentity, RepPeer>,
     fair_queue_inner: Arc<Mutex<QueueInner<ZmqFramedRead, PeerIdentity>>>,
+    socket_monitor: Mutex<Option<mpsc::Sender<SocketEvent>>>,
 }
 
 pub struct RepSocket {
@@ -44,6 +45,7 @@ impl Socket for RepSocket {
             backend: Arc::new(RepSocketBackend {
                 peers: DashMap::new(),
                 fair_queue_inner: fair_queue.inner(),
+                socket_monitor: Mutex::new(None),
             }),
             current_request: None,
             fair_queue,
@@ -58,6 +60,12 @@ impl Socket for RepSocket {
     fn binds(&mut self) -> &mut HashMap<Endpoint, AcceptStopHandle> {
         &mut self.binds
     }
+
+    fn monitor(&mut self) -> mpsc::Receiver<SocketEvent> {
+        let (sender, receiver) = mpsc::channel(1024);
+        self.backend.socket_monitor.lock().replace(sender);
+        receiver
+    }
 }
 
 impl MultiPeerBackend for RepSocketBackend {
@@ -77,6 +85,9 @@ impl MultiPeerBackend for RepSocketBackend {
     }
 
     fn peer_disconnected(&self, peer_id: &PeerIdentity) {
+        if let Some(monitor) = self.monitor().lock().as_mut() {
+            let _ = monitor.try_send(SocketEvent::Disconnected(peer_id.clone()));
+        }
         self.peers.remove(peer_id);
     }
 }
@@ -89,6 +100,10 @@ impl SocketBackend for RepSocketBackend {
     fn shutdown(&self) {
         self.peers.clear();
     }
+
+    fn monitor(&self) -> &Mutex<Option<mpsc::Sender<SocketEvent>>> {
+        &self.socket_monitor
+    }
 }
 
 #[async_trait]