[Bug] Poison pill comparison was failing (#1162)

Author: scottf

src/main/java/io/nats/client/impl/MessageQueue.java

@@ -13,6 +13,8 @@
 
 package io.nats.client.impl;
 
+import io.nats.client.Message;
+
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -30,6 +32,7 @@ class MessageQueue {
     protected static final int STOPPED = 0;
     protected static final int RUNNING = 1;
     protected static final int DRAINING = 2;
+    protected static final String POISON = "_poison";
 
     protected final AtomicLong length;
     protected final AtomicLong sizeInBytes;
@@ -44,9 +47,10 @@ class MessageQueue {
     protected final Duration requestCleanupInterval;
 
     // Poison pill is a graphic, but common term for an item that breaks loops or stop something.
-    // In this class the poisonPill is used to break out of timed waits on the blocking queue.
+    // In this class the poison pill is used to break out of timed waits on the blocking queue.
     // A simple == is used to check if any message in the queue is this message.
-    protected final NatsMessage poisonPill;
+    // /\ /\ /\ /\ which is why it is now a static. It's just a marker anyway.
+    protected static final NatsMessage POISON_PILL = new NatsMessage(POISON, null, EMPTY_BODY);
 
     MessageQueue(boolean singleReaderMode, Duration requestCleanupInterval) {
         this(singleReaderMode, -1, false, requestCleanupInterval, null);
@@ -80,9 +84,6 @@ class MessageQueue {
         this.offerLockMillis = requestCleanupInterval.toMillis();
         this.offerTimeoutMillis = Math.max(1, requestCleanupInterval.toMillis() * 95 / 100);
 
-        // The poisonPill is used to stop poll and accumulate when the queue is stopped
-        this.poisonPill = new NatsMessage("_poison", null, EMPTY_BODY);
-
         editLock = new ReentrantLock();
 
         this.singleReaderMode = singleReaderMode;
@@ -195,7 +196,7 @@ boolean push(NatsMessage msg, boolean internal) {
      */
     void poisonTheQueue() {
         try {
-            this.queue.add(this.poisonPill);
+            this.queue.add(POISON_PILL);
         } catch (IllegalStateException ie) { // queue was full, so we don't really need poison pill
             // ok to ignore this
         }
@@ -222,11 +223,11 @@ NatsMessage poll(Duration timeout) throws InterruptedException {
             }
         }
 
-        if (msg == poisonPill) {
-            return null;
-        }
+        return msg == null || isPoison(msg) ? null : msg;
+    }
 
-        return msg;
+    private boolean isPoison(Message msg) {
+        return msg == POISON_PILL;
     }
 
     NatsMessage pop(Duration timeout) throws InterruptedException {
@@ -286,7 +287,7 @@ NatsMessage accumulate(long maxSize, long maxMessages, Duration timeout)
 
         while (cursor != null) {
             NatsMessage next = this.queue.peek();
-            if (next != null && next != this.poisonPill) {
+            if (next != null && !isPoison(next)) {
                 long s = next.getSizeInBytes();
 
                 if (maxSize<0 || (size + s) < maxSize) { // keep going

src/main/java/io/nats/client/impl/NatsConnection.java

@@ -213,7 +213,7 @@ else if (cur.equals(first)) {
 
             // let server pool resolve hostnames, then loop through resolved
             List<NatsUri> resolvedList = resolveHost(cur);
-            while (!resolvedList.isEmpty()) {
+            for (NatsUri resolved : resolvedList) {
                 if (isClosed()) {
                     keepGoing = false;
                     break;
@@ -224,7 +224,6 @@ else if (cur.equals(first)) {
                 updateStatus(Status.CONNECTING);
 
                 timeTraceLogger.trace("trying to connect to %s", cur);
-                NatsUri resolved = resolvedList.remove(0);
                 tryToConnect(cur, resolved, System.nanoTime());
 
                 if (isConnected()) {
@@ -321,8 +320,6 @@ void forceReconnectImpl() throws IOException, InterruptedException {
         try {
             // calling connect just starts like a new connection versus reconnect
             // but we have to manually resubscribe like reconnect once it is connected
-            // also, lets assume we never want to try the currently connected server
-            serverPool.connectFailed(currentServer);       // we don't want to connect to the same server
             reconnectImpl();
             writer.setReconnectMode(false);
         }
@@ -374,40 +371,36 @@ else if (first.equals(cur)) {
                 // let server list provider resolve hostnames
                 // then loop through resolved
                 List<NatsUri> resolvedList = resolveHost(cur);
-                while (keepGoing && !resolvedList.isEmpty()) {
+                for (NatsUri resolved : resolvedList) {
                     if (isClosed()) {
                         keepGoing = false;
+                        break;
                     }
-                    else {
-                        connectError.set(""); // reset on each loop
-                        if (isDisconnectingOrClosed() || this.isClosing()) {
-                            keepGoing = false;
-                        }
-                        else {
-                            updateStatus(Status.RECONNECTING);
+                    connectError.set(""); // reset on each loop
+                    if (isDisconnectingOrClosed() || this.isClosing()) {
+                        keepGoing = false;
+                        break;
+                    }
+                    updateStatus(Status.RECONNECTING);
 
-                            timeTraceLogger.trace("reconnecting to server %s", cur);
-                            NatsUri resolved = resolvedList.remove(0);
-                            tryToConnect(cur, resolved, System.nanoTime());
+                    timeTraceLogger.trace("reconnecting to server %s", cur);
+                    tryToConnect(cur, resolved, System.nanoTime());
 
-                            if (isConnected()) {
-                                serverPool.connectSucceeded(cur);
-                                statistics.incrementReconnects();
-                                keepGoing = false;
-                            }
-                            else {
-                                serverPool.connectFailed(cur);
-                                String err = connectError.get();
-                                if (this.isAuthenticationError(err)) {
-                                    if (err.equals(this.serverAuthErrors.get(resolved))) {
-                                        keepGoing = false; // double auth error
-                                    }
-                                    else {
-                                        serverAuthErrors.put(resolved, err);
-                                    }
-                                }
-                            }
+                    if (isConnected()) {
+                        serverPool.connectSucceeded(cur);
+                        statistics.incrementReconnects();
+                        keepGoing = false;
+                        break;
+                    }
+
+                    serverPool.connectFailed(cur);
+                    String err = connectError.get();
+                    if (this.isAuthenticationError(err)) {
+                        if (err.equals(this.serverAuthErrors.get(resolved))) {
+                            keepGoing = false; // double auth error
+                            break;
                         }
+                        serverAuthErrors.put(resolved, err);
                     }
                 }
             }

src/test/java/io/nats/client/ConnectTests.java

@@ -15,6 +15,7 @@
 
 import io.nats.client.ConnectionListener.Events;
 import io.nats.client.NatsServerProtocolMock.ExitAt;
+import io.nats.client.api.ServerInfo;
 import io.nats.client.impl.ListenerForTesting;
 import io.nats.client.impl.SimulateSocketDataPortException;
 import org.junit.jupiter.api.Test;
@@ -559,4 +560,47 @@ public void exceptionOccurred(Connection conn, Exception exp) {
             }
         }
     }
-}
+
+    @Test
+    public void testRunInJsCluster() throws Exception {
+        ListenerForTesting[] listeners = new ListenerForTesting[3];
+        listeners[0] = new ListenerForTesting();
+        listeners[1] = new ListenerForTesting();
+        listeners[2] = new ListenerForTesting();
+
+        ThreeServerTestOptionsAppender appender = (ix, builder) ->
+            builder.connectionListener(listeners[ix]).errorListener(listeners[ix]);
+
+        runInJsCluster(ConnectTests::validateRunInJsCluster);
+
+        listeners[0] = new ListenerForTesting();
+        listeners[1] = new ListenerForTesting();
+        listeners[2] = new ListenerForTesting();
+
+        runInJsCluster(true, appender, ConnectTests::validateRunInJsCluster);
+    }
+
+    private static void validateRunInJsCluster(Connection nc1, Connection nc2, Connection nc3) throws InterruptedException {
+        Thread.sleep(200);
+        ServerInfo si1 = nc1.getServerInfo();
+        ServerInfo si2 = nc2.getServerInfo();
+        ServerInfo si3 = nc3.getServerInfo();
+        assertEquals(si1.getCluster(), si2.getCluster());
+        assertEquals(si1.getCluster(), si3.getCluster());
+        String port1 = "" + si1.getPort();
+        String port2 = "" + si2.getPort();
+        String port3 = "" + si3.getPort();
+        String urls1 = String.join(",", si1.getConnectURLs());
+        String urls2 = String.join(",", si2.getConnectURLs());
+        String urls3 = String.join(",", si3.getConnectURLs());
+        assertTrue(urls1.contains(port1));
+        assertTrue(urls1.contains(port2));
+        assertTrue(urls1.contains(port3));
+        assertTrue(urls2.contains(port1));
+        assertTrue(urls2.contains(port2));
+        assertTrue(urls2.contains(port3));
+        assertTrue(urls3.contains(port1));
+        assertTrue(urls3.contains(port2));
+        assertTrue(urls3.contains(port3));
+    }
+}

src/test/java/io/nats/client/impl/ReconnectTests.java

@@ -15,6 +15,7 @@
 
 import io.nats.client.*;
 import io.nats.client.ConnectionListener.Events;
+import io.nats.client.api.ServerInfo;
 import nats.io.NatsRunnerUtils;
 import org.junit.jupiter.api.Test;
 
@@ -716,30 +717,46 @@ private static Thread getReconnectOnConnectTestThread(AtomicReference<Connection
     @Test
     public void testForceReconnect() throws Exception {
         ListenerForTesting listener = new ListenerForTesting();
-        Options.Builder builder = Options.builder()
-            .connectionListener(listener)
-            .errorListener(listener);
+        ThreeServerTestOptionsAppender appender = makeOptionsAppender(listener);
 
-        runInJsServer(nc1 -> runInServer(nc2 -> {
-            int port1 = nc1.getServerInfo().getPort();
-            int port2 = nc2.getServerInfo().getPort();
+        runInJsCluster(appender, (nc0, nc1, nc2) -> {
+            _testForceReconnect(nc0, listener);
+        });
+    }
 
-            String[] servers = new String[]{
-                NatsRunnerUtils.getNatsLocalhostUri(port1),
-                NatsRunnerUtils.getNatsLocalhostUri(port2)
-            };
-            //noinspection resource
-            Connection nc = standardConnection(builder.servers(servers).build());
-            int connectedPort = nc.getServerInfo().getPort();
-            nc.forceReconnect();
-            Thread.sleep(3000);
-            assertNotEquals(connectedPort, nc.getServerInfo().getPort());
-        }));
+    @Test
+    public void testForceReconnectWithAccount() throws Exception {
+        ListenerForTesting listener = new ListenerForTesting();
+        ThreeServerTestOptionsAppender appender = makeOptionsAppender(listener);
+        runInJsCluster(true, appender, (nc0, nc1, nc2) -> {
+            _testForceReconnect(nc0, listener);
+        });
+
+    }
 
+    private static ThreeServerTestOptionsAppender makeOptionsAppender(ListenerForTesting listener) {
+        ThreeServerTestOptionsAppender appender = (ix, builder) -> {
+            if (ix == 0) {
+                builder.connectionListener(listener).ignoreDiscoveredServers().noRandomize();
+            }
+        };
+        return appender;
+    }
+
+    private static void _testForceReconnect(Connection nc0, ListenerForTesting listener) throws IOException, InterruptedException {
+        ServerInfo si = nc0.getServerInfo();
+        String connectedServer = si.getServerId();
+
+        nc0.forceReconnect();
+        standardConnectionWait(nc0);
+
+        si = nc0.getServerInfo();
+        assertNotEquals(connectedServer, si.getServerId());
         assertTrue(listener.getConnectionEvents().contains(Events.DISCONNECTED));
         assertTrue(listener.getConnectionEvents().contains(Events.RECONNECTED));
     }
 
+
     @Test
     public void testSocketDataPortTimeout() throws Exception {
         ListenerForTesting listener = new ListenerForTesting();