From 2393b9e28adfdf25cf80886540f01dfebc079388 Mon Sep 17 00:00:00 2001
From: Colin Sullivan <colin@synadia.com>
Date: Wed, 13 May 2020 16:49:09 -0600
Subject: [PATCH] Track current server

Signed-off-by: Colin Sullivan <colin@synadia.com>
---
 .../io/nats/client/impl/NatsConnection.java    | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/src/main/java/io/nats/client/impl/NatsConnection.java b/src/main/java/io/nats/client/impl/NatsConnection.java
index eafe93b0d..cf42d2680 100644
--- a/src/main/java/io/nats/client/impl/NatsConnection.java
+++ b/src/main/java/io/nats/client/impl/NatsConnection.java
@@ -133,6 +133,8 @@ class NatsConnection implements Connection {
     private ExecutorService executor;
     private ExecutorService connectExecutor;
 
+    private String currentServer = null;
+
     NatsConnection(Options options) {
         boolean trace = options.isTraceConnection();
         timeTrace(trace, "creating connection object");
@@ -192,7 +194,7 @@ void connect(boolean reconnectOnConnect) throws InterruptedException, IOExceptio
 
         timeTrace(trace, "starting connect loop");
 
-        Collection<String> serversToTry = buildServerList(false);
+        Collection<String> serversToTry = buildServerList();
         for (String serverURI : serversToTry) {
             if (isClosed()) {
                 break;
@@ -206,6 +208,7 @@ void connect(boolean reconnectOnConnect) throws InterruptedException, IOExceptio
             tryToConnect(serverURI, System.nanoTime());
 
             if (isConnected()) {
+                this.currentServer = serverURI;
                 break;
             } else {
                 timeTrace(trace, "setting status to disconnected");
@@ -263,7 +266,7 @@ void reconnect() throws InterruptedException {
         boolean doubleAuthError = false;
 
         while (!isConnected() && !isClosed() && !this.isClosing()) {
-            Collection<String> serversToTry = buildServerList(true);
+            Collection<String> serversToTry = buildServerList();
 
             for (String server : serversToTry) {
                 if (isClosed()) {
@@ -291,6 +294,7 @@ void reconnect() throws InterruptedException {
                     break;
                 } else if (isConnected()) {
                     this.statistics.incrementReconnects();
+                    this.currentServer = server;
                     break;
                 } else {
                     String err = connectError.get();
@@ -1643,7 +1647,7 @@ void waitForReconnectTimeout() {
         this.reconnectWaiter.complete(Boolean.TRUE);
     }
 
-    Collection<String> buildServerList(boolean isReconnecting) {
+    Collection<String> buildServerList() {
         ArrayList<String> reconnectList = new ArrayList<>();
 
         reconnectList.addAll(getServers());
@@ -1652,17 +1656,17 @@ Collection<String> buildServerList(boolean isReconnecting) {
             return reconnectList;
         }
 
-        if (!isReconnecting) {
+        if (currentServer == null) {
             Collections.shuffle(reconnectList);
         } else {
-            // Remove the current (first) server from the list, shuffle if it makes sense,
+            // Remove the current server from the list, shuffle if it makes sense,
             // and then add it to the end of the list.  This prevents the client
             // from immediately reconnecting to a server it just lost connection with.
-            String s = reconnectList.remove(0);
+            reconnectList.remove(this.currentServer);
             if (reconnectList.size() > 1) {
                 Collections.shuffle(reconnectList);
             }
-            reconnectList.add(s);
+            reconnectList.add(this.currentServer);
         }
         return reconnectList;
     }