Add on_reconnect callback on the async client

socketio/src/asynchronous/client/builder.rs

@@ -11,8 +11,10 @@ use url::Url;
 use crate::{error::Result, Event, Payload, TransportType};
 
 use super::{
-    callback::{Callback, DynAsyncAnyCallback, DynAsyncCallback},
-    client::Client,
+    callback::{
+        Callback, DynAsyncAnyCallback, DynAsyncCallback, DynAsyncReconnectSettingsCallback,
+    },
+    client::{Client, ReconnectSettings},
 };
 use crate::asynchronous::socket::Socket as InnerSocket;
 
@@ -21,9 +23,10 @@ use crate::asynchronous::socket::Socket as InnerSocket;
 /// namespace is specified, the default namespace `/` is taken. The `connect` method
 /// acts the `build` method and returns a connected [`Client`].
 pub struct ClientBuilder {
-    address: String,
+    pub(crate) address: String,
     pub(crate) on: HashMap<Event, Callback<DynAsyncCallback>>,
     pub(crate) on_any: Option<Callback<DynAsyncAnyCallback>>,
+    pub(crate) on_reconnect: Option<Callback<DynAsyncReconnectSettingsCallback>>,
     pub(crate) namespace: String,
     tls_config: Option<TlsConnector>,
     opening_headers: Option<HeaderMap>,
@@ -82,6 +85,7 @@ impl ClientBuilder {
             address: address.into(),
             on: HashMap::new(),
             on_any: None,
+            on_reconnect: None,
             namespace: "/".to_owned(),
             tls_config: None,
             opening_headers: None,
@@ -192,6 +196,42 @@ impl ClientBuilder {
         self
     }
 
+    /// Registers a callback for reconnect events. The event handler must return
+    /// a [ReconnectSettings] struct with the settings that should be updated.
+    ///
+    /// # Example
+    /// ```rust
+    /// use rust_socketio::{asynchronous::{ClientBuilder, ReconnectSettings}};
+    /// use futures_util::future::FutureExt;
+    /// use serde_json::json;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let client = ClientBuilder::new("http://localhost:4200/")
+    ///         .namespace("/admin")
+    ///         .on_reconnect(|| {
+    ///             async {
+    ///                 let mut settings = ReconnectSettings::new();
+    ///                 settings.address("http://server?test=123");
+    ///                 settings.auth(json!({ "token": "abc" }));
+    ///                 settings
+    ///             }.boxed()
+    ///         })
+    ///         .connect()
+    ///         .await;
+    /// }
+    /// ```
+    pub fn on_reconnect<F>(mut self, callback: F) -> Self
+    where
+        F: for<'a> std::ops::FnMut() -> BoxFuture<'static, ReconnectSettings>
+            + 'static
+            + Send
+            + Sync,
+    {
+        self.on_reconnect = Some(Callback::<DynAsyncReconnectSettingsCallback>::new(callback));
+        self
+    }
+
     /// Registers a Callback for all [`crate::event::Event::Custom`] and [`crate::event::Event::Message`].
     ///
     /// # Example

socketio/src/asynchronous/client/callback.rs

@@ -6,7 +6,7 @@ use std::{
 
 use crate::{Event, Payload};
 
-use super::client::Client;
+use super::client::{Client, ReconnectSettings};
 
 /// Internal type, provides a way to store futures and return them in a boxed manner.
 pub(crate) type DynAsyncCallback =
@@ -16,6 +16,9 @@ pub(crate) type DynAsyncAnyCallback = Box<
     dyn for<'a> FnMut(Event, Payload, Client) -> BoxFuture<'static, ()> + 'static + Send + Sync,
 >;
 
+pub(crate) type DynAsyncReconnectSettingsCallback =
+    Box<dyn for<'a> FnMut() -> BoxFuture<'static, ReconnectSettings> + 'static + Send + Sync>;
+
 pub(crate) struct Callback<T> {
     inner: T,
 }
@@ -77,3 +80,29 @@ impl Callback<DynAsyncAnyCallback> {
         }
     }
 }
+
+impl Deref for Callback<DynAsyncReconnectSettingsCallback> {
+    type Target =
+        dyn for<'a> FnMut() -> BoxFuture<'static, ReconnectSettings> + 'static + Sync + Send;
+
+    fn deref(&self) -> &Self::Target {
+        self.inner.as_ref()
+    }
+}
+
+impl DerefMut for Callback<DynAsyncReconnectSettingsCallback> {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.inner.as_mut()
+    }
+}
+
+impl Callback<DynAsyncReconnectSettingsCallback> {
+    pub(crate) fn new<T>(callback: T) -> Self
+    where
+        T: for<'a> FnMut() -> BoxFuture<'static, ReconnectSettings> + 'static + Sync + Send,
+    {
+        Callback {
+            inner: Box::new(callback),
+        }
+    }
+}

socketio/src/asynchronous/client/client.rs

@@ -33,6 +33,33 @@ enum DisconnectReason {
     Server,
 }
 
+/// Settings that can be updated before reconnecting to a server
+#[derive(Default)]
+pub struct ReconnectSettings {
+    address: Option<String>,
+    auth: Option<serde_json::Value>,
+}
+
+impl ReconnectSettings {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    /// Sets the URL that will be used when reconnecting to the server
+    pub fn address<T>(&mut self, address: T) -> &mut Self
+    where
+        T: Into<String>,
+    {
+        self.address = Some(address.into());
+        self
+    }
+
+    /// Sets the authentication data that will be send in the opening request
+    pub fn auth(&mut self, auth: serde_json::Value) {
+        self.auth = Some(auth);
+    }
+}
+
 /// A socket which handles communication with the server. It's initialized with
 /// a specific address as well as an optional namespace to connect to. If `None`
 /// is given the client will connect to the default namespace `"/"`.
@@ -81,7 +108,19 @@ impl Client {
     }
 
     pub(crate) async fn reconnect(&mut self) -> Result<()> {
-        let builder = self.builder.write().await;
+        let mut builder = self.builder.write().await;
+
+        if let Some(config) = builder.on_reconnect.as_mut() {
+            let reconnect_settings = config().await;
+            if let Some(address) = reconnect_settings.address {
+                builder.address = address;
+            }
+
+            if let Some(auth) = reconnect_settings.auth {
+                self.auth = Some(auth);
+            }
+        }
+
         let socket = builder.inner_create().await?;
 
         // New inner socket that can be connected
@@ -548,7 +587,10 @@ mod test {
     };
 
     use crate::{
-        asynchronous::client::{builder::ClientBuilder, client::Client},
+        asynchronous::{
+            client::{builder::ClientBuilder, client::Client},
+            ReconnectSettings,
+        },
         error::Result,
         packet::{Packet, PacketId},
         Payload, TransportType,
@@ -700,13 +742,28 @@ mod test {
     async fn socket_io_reconnect_integration() -> Result<()> {
         static CONNECT_NUM: AtomicUsize = AtomicUsize::new(0);
         static MESSAGE_NUM: AtomicUsize = AtomicUsize::new(0);
+        static ON_RECONNECT_CALLED: AtomicUsize = AtomicUsize::new(0);
 
         let url = crate::test::socket_io_restart_server();
 
-        let socket = ClientBuilder::new(url)
+        let socket = ClientBuilder::new(url.clone())
             .reconnect(true)
             .max_reconnect_attempts(100)
             .reconnect_delay(100, 100)
+            .on_reconnect(move || {
+                let url = url.clone();
+                async move {
+                    ON_RECONNECT_CALLED.fetch_add(1, Ordering::Release);
+
+                    let mut settings = ReconnectSettings::new();
+
+                    // Try setting the address to what we already have, just
+                    // to test. This is not strictly necessary in real usage.
+                    settings.address(url.to_string());
+                    settings
+                }
+                .boxed()
+            })
             .on("open", |_, socket| {
                 async move {
                     CONNECT_NUM.fetch_add(1, Ordering::Release);
@@ -753,6 +810,10 @@ mod test {
 
         assert_eq!(load(&CONNECT_NUM), 2, "should connect twice");
         assert_eq!(load(&MESSAGE_NUM), 2, "should receive two messages");
+        assert!(
+            load(&ON_RECONNECT_CALLED) > 1,
+            "should call on_reconnect at least once"
+        );
 
         socket.disconnect().await?;
         Ok(())

socketio/src/asynchronous/mod.rs

@@ -5,3 +5,4 @@ mod socket;
 #[cfg(feature = "async")]
 pub use client::builder::ClientBuilder;
 pub use client::client::Client;
+pub use client::client::ReconnectSettings;