Fix pending queue lookup race

Author: zileong.ong

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

@@ -83,15 +83,17 @@ public long getPendingByteLimit() {
      *         {@link #setPendingLimits(long, long) setPendingLimits}.
      */
     public long getPendingMessageCount() {
-        return this.getMessageQueue() != null ? this.getMessageQueue().length() : 0;
+        ConsumerMessageQueue queue = this.getMessageQueue();
+        return queue == null ? 0 : queue.length();
     }
 
     /**
      * @return the cumulative size of the messages waiting to be delivered/popped,
      *         {@link #setPendingLimits(long, long) setPendingLimits}.
      */
     public long getPendingByteCount() {
-        return this.getMessageQueue() != null ? this.getMessageQueue().sizeInBytes() : 0;
+        ConsumerMessageQueue queue = this.getMessageQueue();
+        return queue == null ? 0 : queue.sizeInBytes();
     }
 
     /**
@@ -249,4 +251,4 @@ public CompletableFuture<Boolean> drain(Duration timeout) throws InterruptedExce
      * Abstract method, called by the connection when the drain is complete.
      */
     abstract void cleanUpAfterDrain();
-}
+}

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

@@ -21,11 +21,18 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 public class SlowConsumerTests {
 
+    @Test
+    public void testPendingCountsUseSingleQueueLookup() {
+        assertEquals(0, new QueueInvalidatingConsumer().getPendingMessageCount());
+        assertEquals(0, new QueueInvalidatingConsumer().getPendingByteCount());
+    }
+
     @Test
     public void testDefaultPendingLimits() throws Exception {
         try (NatsTestServer ts = new NatsTestServer(false);
@@ -234,4 +241,30 @@ public void testSlowSubscriberNotification() throws Exception {
             assertEquals(sub, slow.get(0));
         }
     }
-}
+
+    private static class QueueInvalidatingConsumer extends NatsConsumer {
+        private final AtomicInteger queueLookups = new AtomicInteger();
+
+        QueueInvalidatingConsumer() {
+            super(null);
+        }
+
+        @Override
+        public boolean isActive() {
+            return true;
+        }
+
+        @Override
+        ConsumerMessageQueue getMessageQueue() {
+            return queueLookups.getAndIncrement() == 0 ? new ConsumerMessageQueue() : null;
+        }
+
+        @Override
+        void sendUnsubForDrain() {
+        }
+
+        @Override
+        void cleanUpAfterDrain() {
+        }
+    }
+}