Improve reconnect randomization

async-nats/Cargo.toml

@@ -41,6 +41,7 @@ lazy_static = "1.4"
 base64 = "0.13"
 tokio-retry = "0.3"
 ring = "0.16"
+rand = "0.8"
 
 [dev-dependencies]
 criterion =  { version = "0.3", features = ["async_tokio"]}

async-nats/src/connector.rs

@@ -30,6 +30,8 @@ use crate::ToServerAddrs;
 use crate::LANG;
 use crate::VERSION;
 use bytes::BytesMut;
+use rand::seq::SliceRandom;
+use rand::thread_rng;
 use std::cmp;
 use std::collections::HashMap;
 use std::io;
@@ -97,8 +99,13 @@ impl Connector {
 
     pub(crate) async fn try_connect(&mut self) -> Result<(ServerInfo, Connection), ConnectError> {
         let mut error = None;
+        let server_addrs = {
+            let mut rng = thread_rng();
+            let mut server_addrs: Vec<ServerAddr> = self.servers.keys().cloned().collect();
+            server_addrs.shuffle(&mut rng);
+            server_addrs
+        };
 
-        let server_addrs: Vec<ServerAddr> = self.servers.keys().cloned().collect();
         for server_addr in server_addrs {
             let server_attempts = self.servers.get_mut(&server_addr).unwrap();
             let duration = if *server_attempts == 0 {
@@ -109,7 +116,6 @@ impl Connector {
 
                 cmp::min(Duration::from_millis(2_u64.saturating_pow(exp)), max)
             };
-
             *server_attempts += 1;
             sleep(duration).await;
 
@@ -227,8 +233,8 @@ impl Connector {
                                 }
                             },
                             Some(_) => {
-                                self.events_tx.send(Event::Connected).await.ok();
                                 self.state_tx.send(State::Connected).ok();
+                                self.events_tx.send(Event::Connected).await.ok();
                                 return Ok((server_info, connection));
                             }
                             None => {

async-nats/src/lib.rs

@@ -644,8 +644,6 @@ impl ConnectionHandler {
                 .await
                 .unwrap();
         }
-        self.connector.events_tx.try_send(Event::Connected).ok();
-
         Ok(())
     }
 }

async-nats/tests/jwt_tests.rs

@@ -34,6 +34,48 @@ mod client {
             .await
             .expect("published");
     }
+
+    #[cfg(not(target_os = "windows"))]
+    #[tokio::test]
+    async fn jwt_lame_duck_reconnect() {
+        use async_nats::Event;
+        let path = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
+        let c = nats_server::run_cluster("tests/configs/jwt.conf");
+
+        let (tx_recconect, mut rx_reconnect) = tokio::sync::mpsc::channel(10);
+
+        let client = async_nats::ConnectOptions::with_credentials_file(
+            path.join("tests/configs/TestUser.creds"),
+        )
+        .await
+        .unwrap()
+        .event_callback({
+            let tx = tx_recconect.clone();
+            move |event| {
+                let tx = tx.clone();
+                async move {
+                    if event == Event::Connected {
+                        tx.send(()).await.unwrap();
+                    }
+                }
+            }
+        })
+        .connect(c.client_url())
+        .await
+        .unwrap();
+
+        let mut subscriber = client.subscribe("test".into()).await.unwrap();
+        for i in 0..2 {
+            rx_reconnect.recv().await;
+            let mut subscribe = client.subscribe("test".into()).await.unwrap();
+            client.publish("test".into(), "data".into()).await.unwrap();
+            subscribe.next().await.unwrap();
+            client.flush().await.unwrap();
+            assert!(subscriber.next().await.is_some());
+            nats_server::set_lame_duck_mode(&c.servers[i]);
+        }
+    }
+
     #[tokio::test]
     async fn jwt_reconnect() {
         use async_nats::ServerAddr;