Augmented force reconnect (#1165)

Author: scottf

src/main/java/io/nats/client/Connection.java

@@ -547,11 +547,21 @@ enum Status {
     /**
      * Forces reconnect behavior. Stops the current connection including the reading and writing,
      * copies already queued outgoing messages, and then begins the reconnect logic.
+     * Does not flush. Does not force close the connection. See {@link #forceReconnect(ForceReconnectOptions)}.
      * @throws IOException the forceReconnect fails
      * @throws InterruptedException the connection is not connected
      */
     void forceReconnect() throws IOException, InterruptedException;
 
+    /**
+     * Forces reconnect behavior. Stops the current connection including the reading and writing,
+     * copies already queued outgoing messages, and then begins the reconnect logic.
+     * @param options options for how the forceReconnect works
+     * @throws IOException the forceReconnect fails
+     * @throws InterruptedException the connection is not connected
+     */
+    void forceReconnect(ForceReconnectOptions options) throws IOException, InterruptedException;
+
     /**
      * Calculates the round trip time between this client and the server.
      * @return the RTT as a duration

src/main/java/io/nats/client/ForceReconnectOptions.java

@@ -0,0 +1,106 @@
+// Copyright 2024 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package io.nats.client;
+
+import java.time.Duration;
+
+/**
+ * The PublishOptions class specifies the options for publishing with JetStream enabled servers.
+ * Options are created using a {@link ForceReconnectOptions.Builder Builder}.
+ */
+public class ForceReconnectOptions {
+
+    public static final ForceReconnectOptions FORCE_CLOSE_INSTANCE = ForceReconnectOptions.builder().forceClose().build();
+
+    private final boolean forceClose;
+    private final Duration flushWait;
+
+    private ForceReconnectOptions(Builder b) {
+        this.forceClose = b.forceClose;
+        this.flushWait = b.flushWait;
+    }
+
+    public boolean isForceClose() {
+        return forceClose;
+    }
+
+    public boolean isFlush() {
+        return flushWait != null;
+    }
+
+    public Duration getFlushWait() {
+        return flushWait;
+    }
+
+    /**
+     * Creates a builder for the options.
+     * @return the builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * ForceReconnectOptions are created using a Builder.
+     */
+    public static class Builder {
+        boolean forceClose = false;
+        Duration flushWait;
+
+        /**
+         * Constructs a new Builder with the default values.
+         */
+        public Builder() {}
+
+        public Builder forceClose() {
+            this.forceClose = true;
+            return this;
+        }
+
+        /**
+         * @param flushWait if supplied and at least 1 millisecond, the forceReconnect will try to
+         *                  flush before closing for the specified wait time. Flush happens before close
+         *                  so not affected by forceClose option
+         * @return the builder
+         */
+        public Builder flush(Duration flushWait) {
+            this.flushWait = flushWait == null || flushWait.toMillis() < 1 ? null : flushWait;
+            return this;
+        }
+
+        /**
+         * @param flushWaitMillis if supplied and at least 1 millisecond, the forceReconnect will try to
+         *                        flush before closing for the specified wait time. Flush happens before close
+         *                        so not affected by forceClose option
+         * @return the builder
+         */
+        public Builder flush(long flushWaitMillis) {
+            if (flushWaitMillis > 0) {
+                this.flushWait = Duration.ofMillis(flushWaitMillis);
+            }
+            else {
+                this.flushWait = null;
+            }
+            return this;
+        }
+
+        /**
+         * Builds the ForceReconnectOptions.
+         * @return ForceReconnectOptions
+         */
+        public ForceReconnectOptions build() {
+            return new ForceReconnectOptions(this);
+        }
+    }
+}

src/main/java/io/nats/client/Options.java

@@ -1567,7 +1567,7 @@ public Builder proxy(Proxy proxy) {
          * @return the Builder for chaining
          */
         public Builder dataPortType(String dataPortClassName) {
-            this.dataPortType = dataPortClassName;
+            this.dataPortType = dataPortClassName == null ? DEFAULT_DATA_PORT_TYPE : dataPortClassName;
             return this;
         }
 

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

@@ -271,18 +271,32 @@ else if (trace) {
 
     @Override
     public void forceReconnect() throws IOException, InterruptedException {
+        forceReconnect(null);
+    }
+
+    @Override
+    public void forceReconnect(ForceReconnectOptions options) throws IOException, InterruptedException {
         if (!tryingToConnect.get()) {
             try {
                 tryingToConnect.set(true);
-                forceReconnectImpl();
+                forceReconnectImpl(options);
             }
             finally {
                 tryingToConnect.set(false);
             }
         }
     }
 
-    void forceReconnectImpl() throws IOException, InterruptedException {
+    void forceReconnectImpl(ForceReconnectOptions options) throws InterruptedException {
+        if (options != null && options.getFlushWait() != null) {
+            try {
+                flush(options.getFlushWait());
+            }
+            catch (TimeoutException e) {
+                // ignore, don't care, too bad;
+            }
+        }
+
         closeSocketLock.lock();
         try {
             updateStatus(Status.DISCONNECTED);
@@ -299,7 +313,12 @@ void forceReconnectImpl() throws IOException, InterruptedException {
                 dataPort = null;
                 executor.submit(() -> {
                     try {
-                        closeMe.forceClose();
+                        if (options != null && options.isForceClose()) {
+                            closeMe.forceClose();
+                        }
+                        else {
+                            closeMe.close();
+                        }
                     }
                     catch (IOException ignore) {}
                 });

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

@@ -13,6 +13,7 @@
 
 package io.nats.client.impl;
 
+import io.nats.client.ForceReconnectOptions;
 import io.nats.client.Options;
 import io.nats.client.support.NatsUri;
 
@@ -39,10 +40,10 @@ public void run() {
                 writeWatcherTimer.cancel(); // we don't need to repeat this
                 connection.executeCallback((c, el) -> el.socketWriteTimeout(c));
                 try {
-                    connection.forceReconnect();
+                    connection.forceReconnect(ForceReconnectOptions.FORCE_CLOSE_INSTANCE);
                 }
                 catch (IOException e) {
-                    // retry maybe? forceReconnect
+                    // retry maybe?
                 }
                 catch (InterruptedException e) {
                     Thread.currentThread().interrupt();

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

@@ -0,0 +1,35 @@
+// Copyright 2024 The NATS Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at:
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package io.nats.client.impl;
+
+import java.io.IOException;
+
+public class ForceReconnectQueueCheckDataPort extends SocketDataPort {
+    public static String WRITE_CHECK;
+    public static long DELAY;
+
+    @Override
+    public void write(byte[] src, int toWrite) throws IOException {
+        String s = new String(src, 0, Math.min(7, toWrite));
+        if (s.startsWith(WRITE_CHECK)) {
+            try {
+                Thread.sleep(DELAY);
+            }
+            catch (InterruptedException e) {
+                throw new RuntimeException(e);
+            }
+        }
+        super.write(src, toWrite);
+    }
+}