Add reconnect callbacks

Signed-off-by: Tomasz Pietrek <[email protected]>

Author: Jarema

async-nats/src/connector.rs

@@ -67,6 +67,9 @@ pub(crate) struct ConnectorOptions {
     pub(crate) reconnect_delay_callback: Box<dyn Fn(usize) -> Duration + Send + Sync + 'static>,
     pub(crate) auth_callback: Option<CallbackArg1<Vec<u8>, Result<Auth, AuthError>>>,
     pub(crate) max_reconnects: Option<usize>,
+    pub(crate) server_info_callback: Option<CallbackArg1<ServerInfo, ()>>,
+    pub(crate) reconnect_server_callback:
+        Option<CallbackArg1<(Vec<ServerAddr>, ServerInfo), ServerAddr>>,
 }
 
 /// Maintains a list of servers and establishes connections.
@@ -79,6 +82,7 @@ pub(crate) struct Connector {
     pub(crate) events_tx: tokio::sync::mpsc::Sender<Event>,
     pub(crate) state_tx: tokio::sync::watch::Sender<State>,
     pub(crate) max_payload: Arc<AtomicUsize>,
+    last_server_info: Option<ServerInfo>,
 }
 
 pub(crate) fn reconnect_delay_callback_default(attempts: usize) -> Duration {
@@ -110,6 +114,7 @@ impl Connector {
             state_tx,
             max_payload,
             connect_stats,
+            last_server_info: None,
         })
     }
 
@@ -141,7 +146,39 @@ impl Connector {
         let mut error = None;
 
         let mut servers = self.servers.clone();
-        if !self.options.retain_servers_order {
+
+        // If reconnect_server_callback is provided, let the user select the server
+        if let Some(ref callback) = self.options.reconnect_server_callback {
+            // Get the latest server info (use default if this is the first connection)
+            let server_info = self.last_server_info.clone().unwrap_or_default();
+
+            // Extract just the ServerAddr from the tuple
+            let server_addrs: Vec<ServerAddr> =
+                servers.iter().map(|(addr, _)| addr.clone()).collect();
+
+            // Call the user's callback
+            let selected_server = callback.call((server_addrs, server_info)).await;
+
+            // Find the selected server in our list and put it first
+            if let Some(pos) = servers
+                .iter()
+                .position(|(addr, _)| addr == &selected_server)
+            {
+                let selected = servers.remove(pos);
+                servers.insert(0, selected);
+                tracing::debug!(
+                    server = ?selected_server,
+                    "user selected server via reconnect_server_callback"
+                );
+            } else {
+                // Server not in our list, add it
+                tracing::debug!(
+                    server = ?selected_server,
+                    "user selected new server via reconnect_server_callback, adding to list"
+                );
+                servers.insert(0, (selected_server, 0));
+            }
+        } else if !self.options.retain_servers_order {
             servers.shuffle(&mut thread_rng());
             // sort_by is stable, meaning it will retain the order for equal elements.
             servers.sort_by(|a, b| a.1.cmp(&b.1));
@@ -187,6 +224,11 @@ impl Connector {
                     .await
                 {
                     Ok((server_info, mut connection)) => {
+                        // Call server_info_callback if provided
+                        if let Some(ref callback) = self.options.server_info_callback {
+                            callback.call(server_info.clone()).await;
+                        }
+
                         if !self.options.ignore_discovered_servers {
                             for url in &server_info.connect_urls {
                                 let server_addr = url.parse::<ServerAddr>().map_err(|err| {
@@ -327,6 +369,8 @@ impl Connector {
                                     server_info.max_payload,
                                     std::sync::atomic::Ordering::Relaxed,
                                 );
+                                // Save server info for next reconnect callback
+                                self.last_server_info = Some(server_info.clone());
                                 return Ok((server_info, connection));
                             }
                             None => {

async-nats/src/lib.rs

@@ -1025,6 +1025,8 @@ pub async fn connect_with_options<A: ToServerAddrs>(
             reconnect_delay_callback: options.reconnect_delay_callback,
             auth_callback: options.auth_callback,
             max_reconnects: options.max_reconnects,
+            server_info_callback: options.server_info_callback,
+            reconnect_server_callback: options.reconnect_server_callback,
         },
         events_tx,
         state_tx,

async-nats/src/options.rs

@@ -13,7 +13,7 @@
 
 use crate::auth::Auth;
 use crate::connector;
-use crate::{Client, ConnectError, Event, ToServerAddrs};
+use crate::{Client, ConnectError, Event, ServerAddr, ServerInfo, ToServerAddrs};
 use base64::engine::general_purpose::URL_SAFE_NO_PAD;
 use base64::engine::Engine;
 use futures_util::Future;
@@ -65,6 +65,9 @@ pub struct ConnectOptions {
     pub(crate) read_buffer_capacity: u16,
     pub(crate) reconnect_delay_callback: Box<dyn Fn(usize) -> Duration + Send + Sync + 'static>,
     pub(crate) auth_callback: Option<CallbackArg1<Vec<u8>, Result<Auth, AuthError>>>,
+    pub(crate) server_info_callback: Option<CallbackArg1<ServerInfo, ()>>,
+    pub(crate) reconnect_server_callback:
+        Option<CallbackArg1<(Vec<ServerAddr>, ServerInfo), ServerAddr>>,
 }
 
 impl fmt::Debug for ConnectOptions {
@@ -117,6 +120,8 @@ impl Default for ConnectOptions {
             }),
             auth: Default::default(),
             auth_callback: None,
+            server_info_callback: None,
+            reconnect_server_callback: None,
         }
     }
 }
@@ -909,6 +914,71 @@ impl ConnectOptions {
         self.read_buffer_capacity = size;
         self
     }
+
+    /// Registers a callback for when the server sends new INFO messages.
+    /// This typically happens when the server discovers new servers in the cluster.
+    ///
+    /// # Examples
+    /// ```no_run
+    /// # #[tokio::main]
+    /// # async fn main() -> Result<(), async_nats::ConnectError> {
+    /// async_nats::ConnectOptions::new()
+    ///     .server_info_callback(|server_info| async move {
+    ///         println!("Server info updated: {:?}", server_info);
+    ///     })
+    ///     .connect("demo.nats.io")
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn server_info_callback<F, Fut>(mut self, cb: F) -> ConnectOptions
+    where
+        F: Fn(ServerInfo) -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = ()> + 'static + Send + Sync,
+    {
+        self.server_info_callback = Some(CallbackArg1::<ServerInfo, ()>(Box::new(
+            move |server_info| Box::pin(cb(server_info)),
+        )));
+        self
+    }
+
+    /// Registers a callback for customizing server selection during reconnection.
+    /// The callback receives the list of available servers and the most recent ServerInfo,
+    /// and should return the ServerAddr to connect to.
+    ///
+    /// This allows fine-grained control over reconnection logic, such as:
+    /// - Custom server selection algorithms
+    /// - Geographic server preference
+    /// - Dynamic server list manipulation
+    ///
+    /// # Examples
+    /// ```no_run
+    /// # #[tokio::main]
+    /// # async fn main() -> Result<(), async_nats::ConnectError> {
+    /// async_nats::ConnectOptions::new()
+    ///     .reconnect_server_callback(|(servers, server_info)| async move {
+    ///         // Always prefer the first server in the list
+    ///         servers
+    ///             .first()
+    ///             .cloned()
+    ///             .unwrap_or_else(|| "nats://localhost:4222".parse().unwrap())
+    ///     })
+    ///     .connect("demo.nats.io")
+    ///     .await?;
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn reconnect_server_callback<F, Fut>(mut self, cb: F) -> ConnectOptions
+    where
+        F: Fn((Vec<ServerAddr>, ServerInfo)) -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = ServerAddr> + 'static + Send + Sync,
+    {
+        self.reconnect_server_callback = Some(CallbackArg1::<
+            (Vec<ServerAddr>, ServerInfo),
+            ServerAddr,
+        >(Box::new(move |args| Box::pin(cb(args)))));
+        self
+    }
 }
 
 pub(crate) type AsyncCallbackArg1<A, T> =

async-nats/tests/client_tests.rs

@@ -1181,4 +1181,197 @@ mod client {
             .await
             .expect("Expected to be able to create a new client");
     }
+
+    #[tokio::test]
+    async fn test_reconnect_server_callback() {
+        // This test validates that the reconnect_server_callback is invoked
+        // and that the client respects the server selection from the callback.
+        use std::sync::Arc;
+        use tokio::sync::Mutex;
+
+        let server = nats_server::run_basic_server();
+        let correct_addr = server.client_url();
+
+        // Parse the correct server address
+        let correct_server: ServerAddr = correct_addr.parse().unwrap();
+        let correct_server_for_callback = correct_server.clone();
+
+        // Create a fake/non-existent server address for initial attempt
+        let fake_server: ServerAddr = "nats://localhost:9999".parse().unwrap();
+
+        // Track callback invocations
+        let callback_invoked = Arc::new(Mutex::new(false));
+        let callback_invoked_clone = callback_invoked.clone();
+
+        // Track connection events
+        let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(10);
+
+        let client = ConnectOptions::new()
+            .reconnect_server_callback(move |(servers, _server_info)| {
+                let correct = correct_server_for_callback.clone();
+                let invoked = callback_invoked_clone.clone();
+                async move {
+                    *invoked.lock().await = true;
+
+                    // First time: return the correct server
+                    // (The fake server should fail, triggering reconnect)
+                    let has_fake_server = servers.iter().any(|s| {
+                        // Check if this server matches our fake one
+                        format!("{:?}", s).contains("9999")
+                    });
+
+                    if has_fake_server {
+                        // Return the correct server
+                        correct
+                    } else {
+                        // Return the first available server
+                        servers.first().cloned().unwrap_or(correct)
+                    }
+                }
+            })
+            .event_callback(move |event| {
+                let tx = event_tx.clone();
+                async move {
+                    if let Event::Connected = event {
+                        tx.send(event).await.ok();
+                    }
+                }
+            })
+            .retry_on_initial_connect()
+            .connect(vec![fake_server, correct_server.clone()])
+            .await
+            .unwrap();
+
+        // Wait for connection to be established
+        tokio::time::timeout(Duration::from_secs(10), async {
+            while let Some(event) = event_rx.recv().await {
+                if matches!(event, Event::Connected) {
+                    break;
+                }
+            }
+        })
+        .await
+        .expect("Client should connect within timeout");
+
+        // Verify callback was invoked
+        assert!(
+            *callback_invoked.lock().await,
+            "Callback should have been invoked"
+        );
+
+        // Verify we can publish and subscribe (connection is working)
+        let mut subscriber = client.subscribe("test").await.unwrap();
+        client.publish("test", "data".into()).await.unwrap();
+        client.flush().await.unwrap();
+
+        tokio::time::timeout(Duration::from_secs(5), subscriber.next())
+            .await
+            .expect("Should receive message within timeout")
+            .expect("Should receive a message");
+    }
+
+    #[tokio::test]
+    async fn test_server_info_callback() {
+        // This test validates that the server_info_callback is invoked
+        // when the server sends INFO messages.
+        use std::sync::Arc;
+        use tokio::sync::Mutex;
+
+        let server = nats_server::run_basic_server();
+
+        let callback_invoked = Arc::new(Mutex::new(false));
+        let callback_invoked_clone = callback_invoked.clone();
+
+        let _client = ConnectOptions::new()
+            .server_info_callback(move |server_info| {
+                let invoked = callback_invoked_clone.clone();
+                async move {
+                    *invoked.lock().await = true;
+                    // Verify we have some basic server info
+                    assert!(!server_info.server_id.is_empty() || server_info.port > 0);
+                }
+            })
+            .connect(server.client_url())
+            .await
+            .unwrap();
+
+        // Give callback time to be invoked
+        tokio::time::sleep(Duration::from_millis(500)).await;
+
+        // Verify callback was invoked
+        assert!(
+            *callback_invoked.lock().await,
+            "Server info callback should have been invoked"
+        );
+    }
+
+    #[tokio::test]
+    async fn test_reconnect_with_force_reconnect() {
+        // Test that reconnect_server_callback works correctly during force_reconnect
+        use std::sync::Arc;
+        use tokio::sync::Mutex;
+
+        let server = nats_server::run_basic_server();
+        let server_addr: ServerAddr = server.client_url().parse().unwrap();
+
+        let callback_count = Arc::new(Mutex::new(0));
+        let callback_count_clone = callback_count.clone();
+
+        let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(10);
+
+        let client = ConnectOptions::new()
+            .reconnect_server_callback(move |(servers, _)| {
+                let count = callback_count_clone.clone();
+                async move {
+                    *count.lock().await += 1;
+                    servers.first().cloned().unwrap()
+                }
+            })
+            .event_callback(move |event| {
+                let tx = event_tx.clone();
+                async move {
+                    tx.send(event).await.ok();
+                }
+            })
+            .connect(server_addr.clone())
+            .await
+            .unwrap();
+
+        // Wait for initial connection
+        tokio::time::timeout(Duration::from_secs(5), async {
+            while let Some(event) = event_rx.recv().await {
+                if matches!(event, Event::Connected) {
+                    break;
+                }
+            }
+        })
+        .await
+        .expect("Initial connection should succeed");
+
+        // Force a reconnect
+        client.force_reconnect().await.unwrap();
+
+        // Wait for reconnection
+        tokio::time::timeout(Duration::from_secs(5), async {
+            loop {
+                if let Some(event) = event_rx.recv().await {
+                    match event {
+                        Event::Disconnected => continue,
+                        Event::Connected => break,
+                        _ => continue,
+                    }
+                }
+            }
+        })
+        .await
+        .expect("Reconnection should succeed");
+
+        // Verify callback was invoked at least once (possibly twice - initial + reconnect)
+        let count = *callback_count.lock().await;
+        assert!(
+            count >= 1,
+            "Callback should have been invoked at least once, got: {}",
+            count
+        );
+    }
 }