Better disconnection handling under simplification

nats-io · scottf · Nov 27, 2023 · Nov 16, 2023 · Nov 18, 2023 · Nov 18, 2023
commit 3ca2b8e8af8ecb0d417e6c4c0cf6d526b5a59f0d
diff --git a/src/main/java/io/nats/client/JetStreamManagement.java b/src/main/java/io/nats/client/JetStreamManagement.java
@@ -270,4 +270,12 @@ public interface JetStreamManagement {
      * @return true if the delete succeeded
      */
     boolean deleteMessage(String streamName, long seq, boolean erase) throws IOException, JetStreamApiException;
+
+    /**
+     * Gets a context for publishing and subscribing to subjects backed by Jetstream streams
+     * and consumers, using the same connection and JetStreamOptions as the management.
+     * @return a JetStream instance.
+     * @throws IOException various IO exception such as timeout or interruption
+     */
+    JetStream jetStream() throws IOException;
 }
diff --git a/src/main/java/io/nats/client/MessageConsumer.java b/src/main/java/io/nats/client/MessageConsumer.java
@@ -55,4 +55,10 @@ public interface MessageConsumer extends AutoCloseable {
      * @return the finished flag
      */
     boolean isFinished();
+
+    /**
+     * Geta a future
+     * @return
+     */
+//    CompletableFuture<Boolean> getStoppedFuture();
 }
diff --git a/src/main/java/io/nats/client/impl/MessageManager.java b/src/main/java/io/nats/client/impl/MessageManager.java
@@ -80,12 +80,6 @@ protected void startPullRequest(String pullSubject, PullRequestOptions pullReque
         // does nothing - only implemented for pulls, but in base class since instance is referenced as MessageManager, not subclass
     }
 
-    protected void messageReceived() {
-        synchronized (stateChangeLock) {
-            lastMsgReceived = System.currentTimeMillis();
-        }
-    }
-
     protected Boolean beforeQueueProcessorImpl(NatsMessage msg) {
         return true;
     }

diff --git a/src/main/java/io/nats/client/impl/NatsConsumerContext.java b/src/main/java/io/nats/client/impl/NatsConsumerContext.java
@@ -88,7 +88,7 @@ public NatsJetStreamPullSubscription subscribe(MessageHandler messageHandler, Di
             }
             ConsumerConfiguration cc = lastConsumer == null
                 ? originalOrderedCc
-                : streamCtx.js.nextOrderedConsumerConfiguration(originalOrderedCc, highestSeq, null);
+                : streamCtx.js.consumerConfigurationStartAfterLast(originalOrderedCc, highestSeq, null);
             pso = new OrderedPullSubscribeOptionsBuilder(streamCtx.streamName, cc).build();
         }
         else {

diff --git a/src/main/java/io/nats/client/impl/NatsJetStream.java b/src/main/java/io/nats/client/impl/NatsJetStream.java
@@ -35,6 +35,9 @@ public NatsJetStream(NatsConnection connection, JetStreamOptions jsOptions) thro
         super(connection, jsOptions);
     }
 
+    NatsJetStream(NatsJetStreamImpl impl) throws IOException {
+        super(impl);
+    }
     // ----------------------------------------------------------------------------------------------------
     // Publish
     // ----------------------------------------------------------------------------------------------------

diff --git a/src/main/java/io/nats/client/impl/NatsJetStreamImpl.java b/src/main/java/io/nats/client/impl/NatsJetStreamImpl.java
@@ -65,6 +65,13 @@ public CachedStreamInfo(StreamInfo si) {
         multipleSubjectFilter210Available = conn.getInfo().isNewerVersionThan("2.9.99");
     }
 
+    NatsJetStreamImpl(NatsJetStreamImpl impl) throws IOException {
+        conn = impl.conn;
+        jso = impl.jso;
+        consumerCreate290Available = impl.consumerCreate290Available;
+        multipleSubjectFilter210Available = impl.multipleSubjectFilter210Available;
+    }
+
     // ----------------------------------------------------------------------------------------------------
     // Management that is also needed by regular context
     // ----------------------------------------------------------------------------------------------------
@@ -173,7 +180,7 @@ String generateConsumerName() {
         return NUID.nextGlobalSequence();
     }
 
-    ConsumerConfiguration nextOrderedConsumerConfiguration(
+    ConsumerConfiguration consumerConfigurationStartAfterLast(
         ConsumerConfiguration originalCc,
         long lastStreamSeq,
         String newDeliverSubject)

diff --git a/src/main/java/io/nats/client/impl/NatsJetStreamManagement.java b/src/main/java/io/nats/client/impl/NatsJetStreamManagement.java
@@ -295,4 +295,12 @@ public boolean deleteMessage(String streamName, long seq, boolean erase) throws
         Message resp = makeRequestResponseRequired(subj, mdr.serialize(), jso.getRequestTimeout());
         return new SuccessApiResponse(resp).throwOnHasError().getSuccess();
     }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public JetStream jetStream() throws IOException {
+        return new NatsJetStream(this);
+    }
 }
diff --git a/src/main/java/io/nats/client/impl/NatsMessageConsumer.java b/src/main/java/io/nats/client/impl/NatsMessageConsumer.java
@@ -19,9 +19,12 @@
 import java.io.IOException;
 
 class NatsMessageConsumer extends NatsMessageConsumerBase implements PullManagerObserver {
-    protected final PullRequestOptions rePullPro;
+    protected final ConsumeOptions opts;
     protected final int thresholdMessages;
     protected final long thresholdBytes;
+    protected final SimplifiedSubscriptionMaker subscriptionMaker;
+    protected final Dispatcher userDispatcher;
+    protected final MessageHandler userMessageHandler;
 
     NatsMessageConsumer(SimplifiedSubscriptionMaker subscriptionMaker,
                         ConsumerInfo cachedConsumerInfo,
@@ -31,40 +34,67 @@ class NatsMessageConsumer extends NatsMessageConsumerBase implements PullManager
     {
         super(cachedConsumerInfo);
 
+        this.subscriptionMaker = subscriptionMaker;
+        this.opts = opts;
+        this.userDispatcher = userDispatcher;
+        this.userMessageHandler = userMessageHandler;
+
         int bm = opts.getBatchSize();
         long bb = opts.getBatchBytes();
-
         int rePullMessages = Math.max(1, bm * opts.getThresholdPercent() / 100);
         long rePullBytes = bb == 0 ? 0 : Math.max(1, bb * opts.getThresholdPercent() / 100);
-        rePullPro = PullRequestOptions.builder(rePullMessages)
-            .maxBytes(rePullBytes)
-            .expiresIn(opts.getExpiresInMillis())
-            .idleHeartbeat(opts.getIdleHeartbeat())
-            .build();
-
         thresholdMessages = bm - rePullMessages;
         thresholdBytes = bb == 0 ? Integer.MIN_VALUE : bb - rePullBytes;
 
+        doSub();
+    }
+
+    void doSub() throws JetStreamApiException, IOException {
         MessageHandler mh = userMessageHandler == null ? null : msg -> {
             userMessageHandler.onMessage(msg);
             if (stopped.get() && pmm.noMorePending()) {
                 finished.set(true);
             }
         };
-        initSub(subscriptionMaker.subscribe(mh, userDispatcher));
-        sub._pull(PullRequestOptions.builder(bm)
-                .maxBytes(bb)
-                .expiresIn(opts.getExpiresInMillis())
-                .idleHeartbeat(opts.getIdleHeartbeat())
-                .build(),
-            false, this);
+        super.initSub(subscriptionMaker.subscribe(mh, userDispatcher));
+        repull();
     }
 
     @Override
     public void pendingUpdated() {
         if (!stopped.get() && (pmm.pendingMessages <= thresholdMessages || (pmm.trackingBytes && pmm.pendingBytes <= thresholdBytes)))
         {
-            sub._pull(rePullPro, false, this);
+            repull();
+        }
+    }
+
+    boolean subMadeAfterHeartbeatError = false;
+
+    @Override
+    public void heartbeatError() {
+        try {
+            if (pmm.hasUnansweredPulls() && subMadeAfterHeartbeatError) {
+                // we went an entire heartbeat cycle without so much as
+                // this consumer is dead
+                lenientClose();
+                return;
+            }
+            subMadeAfterHeartbeatError = true;
+            doSub();
+        }
+        catch (JetStreamApiException | IOException e) {
+            // TODO FIGURE OUT WHAT TO DO HERE IF ANYTHING
         }
     }
+
+    private void repull() {
+        int rePullMessages = Math.max(1, opts.getBatchSize() - pmm.pendingMessages);
+        long rePullBytes = opts.getBatchBytes() == 0 ? 0 : opts.getBatchBytes() - pmm.pendingBytes;
+        PullRequestOptions pro = PullRequestOptions.builder(rePullMessages)
+            .maxBytes(rePullBytes)
+            .expiresIn(opts.getExpiresInMillis())
+            .idleHeartbeat(opts.getIdleHeartbeat())
+            .build();
+        sub._pull(pro, false, this);
+    }
 }
diff --git a/src/main/java/io/nats/client/impl/OrderedMessageManager.java b/src/main/java/io/nats/client/impl/OrderedMessageManager.java
@@ -83,7 +83,7 @@ private void handleErrorCondition() {
 
             // 3. make a new consumer using the same deliver subject but
             //    with a new starting point
-            ConsumerConfiguration userCC = js.nextOrderedConsumerConfiguration(originalCc, lastStreamSeq, newDeliverSubject);
+            ConsumerConfiguration userCC = js.consumerConfigurationStartAfterLast(originalCc, lastStreamSeq, newDeliverSubject);
             js._createConsumerUnsubscribeOnException(stream, userCC, sub);
 
             // 4. restart the manager.

diff --git a/src/main/java/io/nats/client/impl/OrderedPullMessageManager.java b/src/main/java/io/nats/client/impl/OrderedPullMessageManager.java
@@ -84,7 +84,7 @@ private void handleErrorCondition() {
 
             // 3. make a new consumer using the same deliver subject but
             //    with a new starting point
-            ConsumerConfiguration userCC = js.nextOrderedConsumerConfiguration(originalCc, lastStreamSeq, newDeliverSubject);
+            ConsumerConfiguration userCC = js.consumerConfigurationStartAfterLast(originalCc, lastStreamSeq, newDeliverSubject);
             js._createConsumerUnsubscribeOnException(stream, userCC, sub);
 
             // 4. restart the manager.

diff --git a/src/main/java/io/nats/client/impl/PullManagerObserver.java b/src/main/java/io/nats/client/impl/PullManagerObserver.java
@@ -15,4 +15,5 @@
 
 interface PullManagerObserver {
     void pendingUpdated();
+    void heartbeatError();
 }
diff --git a/src/main/java/io/nats/client/impl/PullMessageManager.java b/src/main/java/io/nats/client/impl/PullMessageManager.java
@@ -18,6 +18,9 @@
 import io.nats.client.SubscribeOptions;
 import io.nats.client.support.Status;
 
+import java.util.ArrayList;
+import java.util.List;
+
 import static io.nats.client.impl.MessageManager.ManageResult.*;
 import static io.nats.client.support.NatsJetStreamConstants.NATS_PENDING_BYTES;
 import static io.nats.client.support.NatsJetStreamConstants.NATS_PENDING_MESSAGES;
@@ -30,12 +33,14 @@ class PullMessageManager extends MessageManager {
     protected boolean trackingBytes;
     protected boolean raiseStatusWarnings;
     protected PullManagerObserver pullManagerObserver;
+    protected boolean initialized;
+    protected List<String> unansweredPulls;
 
     protected PullMessageManager(NatsConnection conn, SubscribeOptions so, boolean syncMode) {
         super(conn, so, syncMode);
-        trackingBytes = false;
-        pendingMessages = 0;
-        pendingBytes = 0;
+        initialized = false;
+        unansweredPulls = new ArrayList<>();
+        reset();
     }
 
     @Override
@@ -59,61 +64,98 @@ protected void startPullRequest(String pullSubject, PullRequestOptions pro, bool
             else {
                 shutdownHeartbeatTimer();
             }
+            unansweredPulls.add(pullSubject);
+        }
+    }
+
+    @Override
+    protected void handleHeartbeatError() {
+        super.handleHeartbeatError();
+        reset();
+        if (pullManagerObserver != null) {
+            pullManagerObserver.heartbeatError();
         }
     }
 
-    private void trackPending(int m, long b) {
+    private void trackIncoming(int m, long b, String pullsubject) {
         synchronized (stateChangeLock) {
-            pendingMessages -= m;
-            boolean zero = pendingMessages < 1;
-            if (trackingBytes) {
-                pendingBytes -= b;
-                zero |= pendingBytes < 1;
+            // message time used for heartbeat tracking
+            // subjects used to detect multiple failed heartbeats
+            lastMsgReceived = System.currentTimeMillis();
+
+            if (pullsubject == null) {
+                unansweredPulls.clear();
             }
-            if (zero) {
-                pendingMessages = 0;
-                pendingBytes = 0;
-                trackingBytes = false;
-                if (hb) {
-                    shutdownHeartbeatTimer();
-                }
+            else {
+                unansweredPulls.remove(pullsubject);
             }
-            if (pullManagerObserver != null) {
-                pullManagerObserver.pendingUpdated();
+
+            if (m != Integer.MIN_VALUE) {
+                pendingMessages -= m;
+                boolean zero = pendingMessages < 1;
+                if (trackingBytes) {
+                    pendingBytes -= b;
+                    zero |= pendingBytes < 1;
+                }
+                if (zero) {
+                    reset();
+                }
+
+                if (pullManagerObserver != null) {
+                    pullManagerObserver.pendingUpdated();
+                }
             }
         }
     }
 
+    protected void reset() {
+        pendingMessages = 0;
+        pendingBytes = 0;
+        trackingBytes = false;
+        if (initialized && hb) {
+            shutdownHeartbeatTimer();
+        }
+        initialized = true;
+    }
+
+    protected boolean hasUnansweredPulls() {
+        synchronized (stateChangeLock) {
+            return !unansweredPulls.isEmpty();
+        }
+    }
+
     @Override
     protected Boolean beforeQueueProcessorImpl(NatsMessage msg) {
-        messageReceived(); // record message time. Used for heartbeat tracking
-
         Status status = msg.getStatus();
 
         // normal js message
         if (status == null) {
-            trackPending(1, msg.consumeByteCount());
+            trackIncoming(1, msg.consumeByteCount(), null);
             return true;
         }
 
         // heartbeat just needed to be recorded
         if (status.isHeartbeat()) {
+            trackIncoming(Integer.MIN_VALUE, -1, msg.subject);
             return false;
         }
 
         Headers h = msg.getHeaders();
+        int m = Integer.MIN_VALUE;
+        long b = -1;
         if (h != null) {
             String s = h.getFirst(NATS_PENDING_MESSAGES);
             if (s != null) {
                 try {
-                    int m = Integer.parseInt(s);
-                    long b = Long.parseLong(h.getFirst(NATS_PENDING_BYTES));
-                    trackPending(m, b);
+                    m = Integer.parseInt(s);
+                    b = Long.parseLong(h.getFirst(NATS_PENDING_BYTES));
+                }
+                catch (NumberFormatException ignore) {
+                    m = Integer.MIN_VALUE; // shouldn't happen but don't fail; make sure don't track m/b
                 }
-                catch (NumberFormatException ignore) {} // shouldn't happen but don't fail
             }
         }
-
+        trackIncoming(m, b, msg.subject);
         return true;
     }
 

diff --git a/src/main/java/io/nats/client/impl/PushMessageManager.java b/src/main/java/io/nats/client/impl/PushMessageManager.java
@@ -73,7 +73,7 @@ protected void startup(NatsJetStreamSubscription sub) {
     @Override
     protected Boolean beforeQueueProcessorImpl(NatsMessage msg) {
         if (hb) {
-            messageReceived(); // only need to track when heartbeats are expected
+            lastMsgReceived = System.currentTimeMillis();  // only need to track when heartbeats are expected
             Status status = msg.getStatus();
             if (status != null) {
                 // only fc heartbeats get queued

diff --git a/src/test/java/io/nats/client/impl/JetStreamConsumerTests.java b/src/test/java/io/nats/client/impl/JetStreamConsumerTests.java
@@ -241,7 +241,7 @@ public void testHeartbeatError() throws Exception {
             sub = js.subscribe(tsc.subject(), d, m -> {}, false, pso);
             validate(sub, testHandler, latch, d);
 
-            latch = setupPulLFactory(js);
+            latch = setupPullFactory(js);
             sub = js.subscribe(tsc.subject(), PullSubscribeOptions.DEFAULT_PULL_OPTS);
             sub.pull(PullRequestOptions.builder(1).idleHeartbeat(100).expiresIn(2000).build());
             validate(sub, testHandler, latch, null);
@@ -278,8 +278,9 @@ private static CountDownLatch setupOrderedFactory(JetStream js) {
         return latch;
     }
 
-    private static CountDownLatch setupPulLFactory(JetStream js) {
-        CountDownLatch latch = new CountDownLatch(2);
+    private static CountDownLatch setupPullFactory(JetStream js) {
+        // expected latch count is 1 b/c pull is dead once there is a hb error
+        CountDownLatch latch = new CountDownLatch(1);
         ((NatsJetStream)js)._pullMessageManagerFactory =
             (conn, lJs, stream, so, serverCC, qmode, dispatcher) ->
                 new PullHeartbeatErrorSimulator(conn, false, latch);