feat(socketioxide): disconnect handler

socketioxide/src/client.rs

@@ -120,11 +120,14 @@ impl<A: Adapter> EngineIoHandler for Client<A> {
     }
     fn on_disconnect(&self, socket: &EIoSocket<Self>) {
         debug!("eio socket disconnect {}", socket.sid);
-        self.ns.values().for_each(|ns| {
-            if let Err(e) = ns.remove_socket(socket.sid) {
-                error!("Adapter error when disconnecting {}: {}, in a multiple server scenario it could leads to desyncronisation issues", socket.sid, e);
-            }
-        });
+        let data = self
+            .ns
+            .values()
+            .filter_map(|ns| ns.get_socket(socket.sid).ok())
+            .map(|s| s.close());
+        if let Err(e) = data.collect::<Result<Vec<_>, _>>() {
+            error!("Adapter error when disconnecting {}: {}, in a multiple server scenario it could leads to desyncronisation issues", socket.sid, e);
+        }
     }
 
     fn on_message(&self, msg: String, socket: &EIoSocket<Self>) {

socketioxide/src/ns.rs

@@ -3,12 +3,12 @@ use std::{
     sync::{Arc, RwLock},
 };
 
-use crate::errors::{AdapterError, SendError};
+use crate::errors::AdapterError;
 use crate::{
     adapter::{Adapter, LocalAdapter},
     errors::Error,
     handshake::Handshake,
-    packet::{Packet, PacketData},
+    packet::PacketData,
     socket::Socket,
     SocketIoConfig,
 };
@@ -67,15 +67,7 @@ impl<A: Adapter> Namespace<A> {
         socket
     }
 
-    pub fn disconnect(&self, sid: Sid) -> Result<(), SendError> {
-        if let Some(socket) = self.sockets.write().unwrap().remove(&sid) {
-            self.adapter
-                .del_all(sid)
-                .map_err(|err| AdapterError(Box::new(err)))?;
-            socket.send(Packet::disconnect(self.path.clone()))?;
-        }
-        Ok(())
-    }
+    /// Remove a socket from a namespace and propagate the event to the adapter
     pub fn remove_socket(&self, sid: Sid) -> Result<(), AdapterError> {
         self.sockets.write().unwrap().remove(&sid);
         self.adapter
@@ -87,19 +79,11 @@ impl<A: Adapter> Namespace<A> {
         self.sockets.read().unwrap().values().any(|s| s.sid == sid)
     }
 
-    /// Called when a namespace receive a particular packet that should be transmitted to the socket
-    pub fn socket_recv(&self, sid: Sid, packet: PacketData) -> Result<(), Error> {
-        self.get_socket(sid)?.recv(packet)
-    }
-
     pub fn recv(&self, sid: Sid, packet: PacketData) -> Result<(), Error> {
         match packet {
-            PacketData::Disconnect => self
-                .remove_socket(sid)
-                .map_err(|err| AdapterError(Box::new(err)).into()),
             PacketData::Connect(_) => unreachable!("connect packets should be handled before"),
             PacketData::ConnectError(_) => Ok(()),
-            packet => self.socket_recv(sid, packet),
+            packet => self.get_socket(sid)?.recv(packet),
         }
     }
     pub fn get_socket(&self, sid: Sid) -> Result<Arc<Socket<A>>, Error> {

socketioxide/src/socket.rs

@@ -3,18 +3,18 @@
     fmt::Debug,
     sync::{
         atomic::{AtomicI64, Ordering},
-        Arc, RwLock,
+        Arc, Mutex, RwLock,
     },
     time::Duration,
 };
 
 use engineioxide::{sid_generator::Sid, SendPacket as EnginePacket};
-use futures::Future;
+use futures::{future::BoxFuture, Future};
 use serde::{de::DeserializeOwned, Serialize};
 use serde_json::Value;
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot};
 
-use crate::errors::SendError;
+use crate::errors::{SendError, AdapterError};
 use crate::retryer::Retryer;
 use crate::{
     adapter::{Adapter, Room},
@@ -28,13 +28,17 @@
     SocketIoConfig,
 };
 
+pub type DisconnectCallback<A> =
+    Box<dyn FnOnce(Arc<Socket<A>>) -> BoxFuture<'static, ()> + Send + Sync + 'static>;
+
 /// A Socket represents a client connected to a namespace.
 /// It is used to send and receive messages from the client, join and leave rooms, etc.
 pub struct Socket<A: Adapter> {
     config: Arc<SocketIoConfig>,
     ns: Arc<Namespace<A>>,
     message_handlers: RwLock<HashMap<String, BoxedHandler<A>>>,
-    ack_message: RwLock<HashMap<i64, oneshot::Sender<AckResponse<Value>>>>,
+    disconnect_handler: Mutex<Option<DisconnectCallback<A>>>,
+    ack_message: Mutex<HashMap<i64, oneshot::Sender<AckResponse<Value>>>>,
     ack_counter: AtomicI64,
     tx: tokio::sync::mpsc::Sender<EnginePacket>,
     pub handshake: Handshake,
@@ -54,7 +58,8 @@
             tx,
             ns,
             message_handlers: RwLock::new(HashMap::new()),
-            ack_message: RwLock::new(HashMap::new()),
+            disconnect_handler: Mutex::new(None),
+            ack_message: Mutex::new(HashMap::new()),
             ack_counter: AtomicI64::new(0),
             handshake,
             sid,
@@ -123,11 +128,35 @@
     {
         let handler = Box::new(move |s, v, p, ack_fn| Box::pin(callback(s, v, p, ack_fn)) as _);
         self.message_handlers
-            .write()
+            .lock()
             .unwrap()
             .insert(event.into(), MessageHandler::boxed(handler));
     }
 
+    /// ### Register a disconnect handler.
+    /// The callback will be called when the socket is disconnected from the server or the client or when the underlying connection crashes.
+    /// ##### Example
+    /// ```
+    /// # use socketioxide::Namespace;
+    /// # use serde_json::Value;
+    /// Namespace::builder().add("/", |socket| async move {
+    ///     socket.on("test", |socket, data: Value, bin, _| async move {
+    ///         // Close the current socket
+    ///         socket.disconnect().ok();
+    ///     });
+    ///     socket.on_disconnect(|socket| async move {
+    ///         println!("Socket {} on ns {} disconnected", socket.sid, socket.ns());
+    ///     });
+    /// });
+    pub fn on_disconnect<C, F>(&self, callback: C)
+    where
+        C: Fn(Arc<Socket<A>>) -> F + Send + Sync + 'static,
+        F: Future<Output = ()> + Send + 'static,
+    {
+        let handler = Box::new(move |s| Box::pin(callback(s)) as _);
+        *self.disconnect_handler.lock().unwrap() = Some(handler);
+    }
+
     /// Emit a message to the client
     /// ##### Example
     /// ```
@@ -341,9 +370,13 @@
         Operators::new(self.ns.clone(), self.sid).broadcast()
     }
 
-    /// Disconnect the socket from the current namespace.
-    pub fn disconnect(&self) -> Result<(), SendError> {
-        self.ns.disconnect(self.sid)
+    /// Disconnect the socket from the current namespace,
+    ///
+    /// It will also call the disconnect handler if it is set.
+    pub fn disconnect(self: Arc<Self>) -> Result<(), SendError> {
+        self.send(Packet::disconnect(self.ns.path.clone()))?;
+        self.close()?;
+        Ok(())
     }
 
     /// Get the current namespace path.
@@ -371,22 +404,32 @@
     ) -> Result<AckResponse<V>, AckError> {
         let (tx, rx) = oneshot::channel();
         let ack = self.ack_counter.fetch_add(1, Ordering::SeqCst) + 1;
-        self.ack_message.write().unwrap().insert(ack, tx);
+        self.ack_message.lock().unwrap().insert(ack, tx);
         packet.inner.set_ack_id(ack);
         self.send(packet)?;
         let timeout = timeout.unwrap_or(self.config.ack_timeout);
         let v = tokio::time::timeout(timeout, rx).await??;
         Ok((serde_json::from_value(v.0)?, v.1))
     }
 
-    // Receive data from client:
+    /// Called when the socket is gracefully [`disconnect`](Socket::disconnect)ed from the server or the client
+    /// 
+    /// It maybe also closed when the underlying transport is closed or failed.
+    pub(crate) fn close(self: Arc<Self>) -> Result<(), AdapterError> {
+        if let Some(handler) = self.disconnect_handler.lock().unwrap().take() {
+            tokio::spawn(handler(self.clone()));
+        }
+        self.ns.remove_socket(self.sid)
+    }
 
+    // Receive data from client:
     pub(crate) fn recv(self: Arc<Self>, packet: PacketData) -> Result<(), Error> {
         match packet {
             PacketData::Event(e, data, ack) => self.recv_event(e, data, ack),
             PacketData::EventAck(data, ack_id) => self.recv_ack(data, ack_id),
             PacketData::BinaryEvent(e, packet, ack) => self.recv_bin_event(e, packet, ack),
             PacketData::BinaryAck(packet, ack) => self.recv_bin_ack(packet, ack),
+            PacketData::Disconnect => self.close().map_err(Error::from),
             _ => unreachable!(),
         }
     }
@@ -411,14 +454,14 @@
     }
 
     fn recv_ack(self: Arc<Self>, data: Value, ack: i64) -> Result<(), Error> {
-        if let Some(tx) = self.ack_message.write().unwrap().remove(&ack) {
+        if let Some(tx) = self.ack_message.lock().unwrap().remove(&ack) {
             tx.send((data, vec![])).ok();
         }
         Ok(())
     }
 
     fn recv_bin_ack(self: Arc<Self>, packet: BinaryPacket, ack: i64) -> Result<(), Error> {
-        if let Some(tx) = self.ack_message.write().unwrap().remove(&ack) {
+        if let Some(tx) = self.ack_message.lock().unwrap().remove(&ack) {
             tx.send((packet.data, packet.bin)).ok();
         }
         Ok(())