From d8af2b2fa93b1b52d0415c9b4a0e9e9cf26e7520 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Th=C3=A9odore=20Pr=C3=A9vot?= <prevottheodore@gmail.com>
Date: Wed, 24 Jan 2024 11:52:58 +0100
Subject: [PATCH] fix(socketio/client): connect timeout panic under heavy
 traffic (#252)

---
 socketioxide/src/client.rs | 47 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/socketioxide/src/client.rs b/socketioxide/src/client.rs
index 57b932cc..ffb9fbb1 100644
--- a/socketioxide/src/client.rs
+++ b/socketioxide/src/client.rs
@@ -52,7 +52,7 @@ impl<A: Adapter> Client<A> {
 
             // cancel the connect timeout task for v5
             if let Some(tx) = esocket.data.connect_recv_tx.lock().unwrap().take() {
-                tx.send(()).unwrap();
+                tx.send(()).ok();
             }
 
             Ok(())
@@ -281,3 +281,48 @@ fn apply_payload_on_packet(data: Vec<u8>, socket: &EIoSocket<SocketData>) -> boo
         false
     }
 }
+
+#[cfg(test)]
+mod test {
+    use tokio::sync::mpsc;
+
+    use crate::adapter::LocalAdapter;
+    const CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(10);
+
+    fn create_client() -> super::Client<LocalAdapter> {
+        let config = crate::SocketIoConfig {
+            connect_timeout: CONNECT_TIMEOUT,
+            ..Default::default()
+        };
+        let client = Client::<LocalAdapter>::new(std::sync::Arc::new(config));
+        client.add_ns("/".into(), || {});
+        client
+    }
+
+    use super::*;
+    #[tokio::test]
+    async fn connect_timeout_fail() {
+        let client = create_client();
+        let (tx, mut rx) = mpsc::channel(1);
+        let close_fn = Box::new(move |_, _| tx.try_send(()).unwrap());
+        let sock = Arc::new(EIoSocket::new_dummy(Sid::new(), close_fn));
+        client.on_connect(sock.clone());
+        tokio::time::timeout(CONNECT_TIMEOUT * 2, rx.recv())
+            .await
+            .unwrap()
+            .unwrap();
+    }
+
+    #[tokio::test]
+    async fn connect_timeout() {
+        let client = create_client();
+        let (tx, mut rx) = mpsc::channel(1);
+        let close_fn = Box::new(move |_, _| tx.try_send(()).unwrap());
+        let sock = Arc::new(EIoSocket::new_dummy(Sid::new(), close_fn));
+        client.on_connect(sock.clone());
+        client.on_message("0".into(), sock.clone());
+        tokio::time::timeout(CONNECT_TIMEOUT * 2, rx.recv())
+            .await
+            .unwrap_err();
+    }
+}