feat(socketio): make rooms available in disconnect handler (#613)

Author: Mettwasser

crates/socketioxide/src/handler/disconnect.rs

@@ -52,6 +52,13 @@ use super::MakeErasedHandler;
 pub(crate) type BoxedDisconnectHandler<A> = Box<dyn ErasedDisconnectHandler<A>>;
 pub(crate) trait ErasedDisconnectHandler<A: Adapter>: Send + Sync + 'static {
     fn call(&self, s: Arc<Socket<A>>, reason: DisconnectReason);
+
+    fn call_with_defer(
+        &self,
+        s: Arc<Socket<A>>,
+        reason: DisconnectReason,
+        defer: fn(Arc<Socket<A>>),
+    );
 }
 
 impl<A: Adapter, T, H> MakeErasedHandler<H, A, T>
@@ -74,6 +81,17 @@ where
     fn call(&self, s: Arc<Socket<A>>, reason: DisconnectReason) {
         self.handler.call(s, reason);
     }
+
+    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self, s, defer), fields(id = ?s.id)))]
+    #[inline(always)]
+    fn call_with_defer(
+        &self,
+        s: Arc<Socket<A>>,
+        reason: DisconnectReason,
+        defer: fn(Arc<Socket<A>>),
+    ) {
+        self.handler.call_with_defer(s, reason, defer);
+    }
 }
 
 /// A trait used to extract the arguments from the disconnect event.
@@ -114,7 +132,17 @@ See `https://docs.rs/socketioxide/latest/socketioxide/extract/index.html` for de
 )]
 pub trait DisconnectHandler<A: Adapter, T>: Send + Sync + 'static {
     /// Call the handler with the given arguments.
-    fn call(&self, s: Arc<Socket<A>>, reason: DisconnectReason);
+    fn call(&self, s: Arc<Socket<A>>, reason: DisconnectReason) {
+        self.call_with_defer(s, reason, |_| ());
+    }
+
+    /// Call the handler and issue a deferred function that will be executed **after** the disconnect handler
+    fn call_with_defer(
+        &self,
+        s: Arc<Socket<A>>,
+        reason: DisconnectReason,
+        defer: fn(Arc<Socket<A>>),
+    );
 
     #[doc(hidden)]
     fn phantom(&self) -> std::marker::PhantomData<T> {
@@ -135,7 +163,7 @@ macro_rules! impl_handler_async {
             A: Adapter,
             $( $ty: FromDisconnectParts<A> + Send, )*
         {
-            fn call(&self, s: Arc<Socket<A>>, reason: DisconnectReason) {
+            fn call_with_defer(&self, s: Arc<Socket<A>>, reason: DisconnectReason, defer: fn(Arc<Socket<A>>)) {
                 $(
                     let $ty = match $ty::from_disconnect_parts(&s, reason) {
                         Ok(v) => v,
@@ -148,7 +176,11 @@ macro_rules! impl_handler_async {
                 )*
 
                 let fut = (self.clone())($($ty,)*);
-                tokio::spawn(fut);
+
+                tokio::spawn(async move {
+                    fut.await;
+                    defer(s);
+                });
             }
         }
     };

crates/socketioxide/src/socket.rs

@@ -445,6 +445,11 @@ impl<A: Adapter> Socket<A> {
     /// # Register a disconnect handler.
     /// You can register only one disconnect handler per socket. If you register multiple handlers, only the last one will be used.
     ///
+    /// This implementation is slightly different to the socket.io spec.
+    /// The difference being that [`rooms`](Self::rooms) are still available in this handler
+    /// and only cleaned up AFTER the execution of this handler.
+    /// Therefore you must not indefinitely stall/hang this handler, for example by entering an endless loop.
+    ///
     /// _It is recommended for code clarity to define your handler as top level function rather than closures._
     ///
     /// * See the [`disconnect`](crate::handler::disconnect) module doc for more details on disconnect handler.
@@ -770,15 +775,16 @@ impl<A: Adapter> Socket<A> {
     pub(crate) fn close(self: Arc<Self>, reason: DisconnectReason) {
         self.set_connected(false);
 
-        let handler = { self.disconnect_handler.lock().unwrap().take() };
-        if let Some(handler) = handler {
+        let disconnect_handler = { self.disconnect_handler.lock().unwrap().take() };
+
+        if let Some(handler) = disconnect_handler {
             #[cfg(feature = "tracing")]
             tracing::trace!(?reason, ?self.id, "spawning disconnect handler");
 
-            handler.call(self.clone(), reason);
+            handler.call_with_defer(self.clone(), reason, |s| s.ns.remove_socket(s.id));
+        } else {
+            self.ns.remove_socket(self.id);
         }
-
-        self.ns.remove_socket(self.id);
     }
 
     /// Receive data from client

crates/socketioxide/tests/disconnect_reason.rs

@@ -12,6 +12,7 @@ use std::time::Duration;
 
 use futures_util::{SinkExt, StreamExt};
 use socketioxide::{SocketIo, extract::SocketRef, socket::DisconnectReason};
+use socketioxide_core::adapter::Room;
 use tokio::sync::mpsc;
 
 mod fixture;
@@ -198,25 +199,38 @@ pub async fn client_ns_disconnect() {
 
 #[tokio::test]
 pub async fn server_ns_disconnect() {
-    let (tx, mut rx) = mpsc::channel::<DisconnectReason>(1);
+    let (tx, mut rx) = mpsc::channel::<(DisconnectReason, Vec<Room>)>(1);
     let (svc, io) = create_server().await;
+
     io.ns("/", async move |socket: SocketRef, io: SocketIo| {
+        socket.join("testRoom1");
+        socket.join("testRoom2");
+        socket.join("testRoom3");
+
         tokio::spawn(async move {
             tokio::time::sleep(Duration::from_millis(100)).await;
             let s = io.sockets().into_iter().next().unwrap();
             s.disconnect().unwrap();
         });
 
-        socket.on_disconnect(async move |reason: DisconnectReason| tx.try_send(reason).unwrap());
+        socket.on_disconnect(async move |s: SocketRef, reason: DisconnectReason| {
+            tx.try_send((reason, s.rooms())).unwrap();
+        });
     });
 
     let _stream = create_ws_connection(&svc).await;
 
-    let data = tokio::time::timeout(Duration::from_millis(100), rx.recv())
-        .await
-        .expect("timeout waiting for DisconnectReason::ServerNSDisconnect")
-        .unwrap();
-    assert_eq!(data, DisconnectReason::ServerNSDisconnect);
+    let (disconnect_reason, mut rooms) =
+        tokio::time::timeout(Duration::from_millis(100), rx.recv())
+            .await
+            .expect("timeout waiting for DisconnectReason::ServerNSDisconnect")
+            .unwrap();
+
+    assert_eq!(disconnect_reason, DisconnectReason::ServerNSDisconnect);
+
+    // Sort the rooms to guarantee order for the assertion
+    rooms.sort();
+    assert_eq!(rooms, vec!["testRoom1", "testRoom2", "testRoom3"]);
 }
 
 #[tokio::test]

e2e/engineioxide/engineioxide.rs

@@ -25,6 +25,7 @@ impl EngineIoHandler for MyHandler {
     fn on_connect(self: Arc<Self>, socket: Arc<Socket<Self::Data>>) {
         println!("socket connect {}", socket.id);
     }
+
     fn on_disconnect(&self, socket: Arc<Socket<Self::Data>>, reason: DisconnectReason) {
         println!("socket disconnect {}: {:?}", socket.id, reason);
     }