Take QueryOffsetResponse in Client, not only offset

References rabbitmq/rabbitmq-server#3783

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/impl/Client.java

@@ -1077,7 +1077,7 @@ public void storeOffset(String reference, String stream, long offset) {
     channel.writeAndFlush(bb);
   }
 
-  public long queryOffset(String reference, String stream) {
+  public QueryOffsetResponse queryOffset(String reference, String stream) {
     if (reference == null || reference.isEmpty() || reference.length() > 256) {
       throw new IllegalArgumentException(
           "Reference must a non-empty string of less than 256 characters");
@@ -1103,10 +1103,7 @@ public long queryOffset(String reference, String stream) {
       channel.writeAndFlush(bb);
       request.block();
       QueryOffsetResponse response = request.response.get();
-      if (!response.isOk()) {
-        LOGGER.info("Query offset failed with code {}", formatConstant(response.getResponseCode()));
-      }
-      return response.getOffset();
+      return response;
     } catch (RuntimeException e) {
       outstandingRequests.remove(correlationId);
       throw new StreamException(e);
@@ -1775,7 +1772,7 @@ static class OpenResponse extends Response {
     }
   }
 
-  static class QueryOffsetResponse extends Response {
+  public static class QueryOffsetResponse extends Response {
 
     private final long offset;
 

src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java

@@ -29,6 +29,7 @@
 import com.rabbitmq.stream.impl.Client.CreditNotification;
 import com.rabbitmq.stream.impl.Client.MessageListener;
 import com.rabbitmq.stream.impl.Client.MetadataListener;
+import com.rabbitmq.stream.impl.Client.QueryOffsetResponse;
 import com.rabbitmq.stream.impl.Client.ShutdownListener;
 import com.rabbitmq.stream.impl.Utils.ClientFactory;
 import com.rabbitmq.stream.impl.Utils.ClientFactoryContext;
@@ -647,9 +648,9 @@ synchronized void add(
 
         String offsetTrackingReference = subscriptionTracker.offsetTrackingReference;
         if (offsetTrackingReference != null) {
-          long trackedOffset =
+          QueryOffsetResponse queryOffsetResponse =
               client.queryOffset(offsetTrackingReference, subscriptionTracker.stream);
-          if (trackedOffset != 0) {
+          if (queryOffsetResponse.isOk() && queryOffsetResponse.getOffset() != 0) {
             if (offsetSpecification != null && isInitialSubscription) {
               // subscription call (not recovery), so telling the user their offset specification is
               // ignored
@@ -660,11 +661,11 @@ synchronized void add(
             }
             LOGGER.debug(
                 "Using offset {} to start consuming from {} with consumer {} " + "(instead of {})",
-                trackedOffset,
+                queryOffsetResponse.getOffset(),
                 subscriptionTracker.stream,
                 offsetTrackingReference,
                 offsetSpecification);
-            offsetSpecification = OffsetSpecification.offset(trackedOffset + 1);
+            offsetSpecification = OffsetSpecification.offset(queryOffsetResponse.getOffset() + 1);
           }
         }
 

src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java

@@ -17,7 +17,9 @@
 import com.rabbitmq.stream.MessageHandler;
 import com.rabbitmq.stream.MessageHandler.Context;
 import com.rabbitmq.stream.OffsetSpecification;
+import com.rabbitmq.stream.StreamException;
 import com.rabbitmq.stream.SubscriptionListener;
+import com.rabbitmq.stream.impl.Client.QueryOffsetResponse;
 import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration;
 import com.rabbitmq.stream.impl.StreamEnvironment.TrackingConsumerRegistration;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -28,29 +30,18 @@
 class StreamConsumer implements Consumer {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(StreamConsumer.class);
-
-  private volatile Runnable closingCallback;
-
   private final Runnable closingTrackingCallback;
-
   private final AtomicBoolean closed = new AtomicBoolean(false);
-
   private final String name;
-
   private final String stream;
-
   private final StreamEnvironment environment;
-
+  private final LongConsumer trackingCallback;
+  private final Runnable initCallback;
+  private volatile Runnable closingCallback;
   private volatile Client trackingClient;
-
   private volatile Status status;
-
   private volatile long lastRequestedStoredOffset = 0;
 
-  private final LongConsumer trackingCallback;
-
-  private final Runnable initCallback;
-
   StreamConsumer(
       String stream,
       OffsetSpecification offsetSpecification,
@@ -195,15 +186,35 @@ void running() {
 
   long lastStoredOffset() {
     if (canTrack()) {
+      // the client can be null by now, so we catch any exception
+      QueryOffsetResponse response;
       try {
-        // the client can be null by now, but we catch the exception and return 0
-        // callers should know how to deal with a stored offset of 0
-        return this.trackingClient.queryOffset(this.name, this.stream);
+        response = this.trackingClient.queryOffset(this.name, this.stream);
       } catch (Exception e) {
-        return 0;
+        throw new IllegalStateException(
+            String.format(
+                "Not possible to query offset for consumer %s on stream %s for now",
+                this.name, this.stream),
+            e);
       }
+      if (response.isOk()) {
+        return response.getOffset();
+      } else {
+        throw new StreamException(
+            String.format(
+                "QueryOffset for consumer %s on stream %s returned an error",
+                this.name, this.stream),
+            response.getResponseCode());
+      }
+
+    } else if (this.name == null) {
+      throw new UnsupportedOperationException(
+          "Not possible to query stored offset for a consumer without a name");
     } else {
-      return 0;
+      throw new IllegalStateException(
+          String.format(
+              "Not possible to query offset for consumer %s on stream %s for now",
+              this.name, this.stream));
     }
   }
 

src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java

@@ -570,7 +570,9 @@ TrackingConsumerRegistration registerTrackingConsumer(
               offsetTrackingRegistration.closingCallback().run();
               LOGGER.debug("Offset tracking registration closing sequence executed");
             } catch (Exception e) {
-              LOGGER.warn("Error while executing offset tracking registration closing sequence");
+              LOGGER.warn(
+                  "Error while executing offset tracking registration closing sequence: {}",
+                  e.getMessage());
             }
             closingCallable.run();
           };

src/test/java/com/rabbitmq/stream/impl/AuthorisationTest.java

@@ -289,7 +289,8 @@ void storeQueryOffsetShouldSucceedOnAuthorisedStreamShouldFailOnUnauthorisedStre
       configurationClient.storeOffset("configuration", s, 10);
 
       Duration timeToCheckOffsetTracking =
-          waitAtMost(5, () -> configurationClient.queryOffset("configuration", s) == 10);
+          waitAtMost(
+              5, () -> configurationClient.queryOffset("configuration", s).getOffset() == 10);
 
       Client client = client();
 
@@ -300,7 +301,7 @@ void storeQueryOffsetShouldSucceedOnAuthorisedStreamShouldFailOnUnauthorisedStre
       assertThat(configurationClient.queryOffset("default-client", s)).isNotEqualTo(10);
 
       // querying is not even authorised for the default client, it should return 0
-      assertThat(client.queryOffset("configuration", s)).isZero();
+      assertThat(client.queryOffset("configuration", s).getOffset()).isZero();
 
     } finally {
       assertThat(configurationClient.delete(s).isOk()).isTrue();

src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java

@@ -36,6 +36,7 @@
 import com.rabbitmq.stream.SubscriptionListener;
 import com.rabbitmq.stream.codec.WrapperMessageBuilder;
 import com.rabbitmq.stream.impl.Client.MessageListener;
+import com.rabbitmq.stream.impl.Client.QueryOffsetResponse;
 import com.rabbitmq.stream.impl.MonitoringTestUtils.ConsumersPoolInfo;
 import com.rabbitmq.stream.impl.Utils.ClientFactory;
 import java.time.Duration;
@@ -1095,8 +1096,8 @@ void shouldUseStoredOffsetOnRecovery(Consumer<ConsumersCoordinatorTest> configur
     long lastStoredOffset = 5;
     long lastReceivedOffset = 10;
     when(client.queryOffset(consumerName, "stream"))
-        .thenReturn((long) 0)
-        .thenReturn(lastStoredOffset);
+        .thenReturn(new QueryOffsetResponse(Constants.RESPONSE_CODE_OK, 0L))
+        .thenReturn(new QueryOffsetResponse(Constants.RESPONSE_CODE_OK, lastStoredOffset));
 
     ArgumentCaptor<OffsetSpecification> offsetSpecificationArgumentCaptor =
         ArgumentCaptor.forClass(OffsetSpecification.class);

src/test/java/com/rabbitmq/stream/impl/OffsetTrackingTest.java

@@ -97,7 +97,7 @@ void trackAndQuery(
       client.storeOffset(trackingReference, s, storedOffset);
     }
     Thread.sleep(100L); // store offset is fire-and-forget
-    long offset = client.queryOffset(queryReference, s);
+    long offset = client.queryOffset(queryReference, s).getOffset();
     assertThat(offset).as(message).isEqualTo(expectedOffset);
   }
 
@@ -213,7 +213,8 @@ void consumeAndStore(BiConsumer<String, Client> streamCreator, TestInfo info) th
                 waitAtMost(
                     5,
                     () ->
-                        lastStoredOffset.get() == consumerReference.get().queryOffset(reference, s),
+                        lastStoredOffset.get()
+                            == consumerReference.get().queryOffset(reference, s).getOffset(),
                     () ->
                         "expecting last stored offset to be "
                             + lastStoredOffset
@@ -249,7 +250,7 @@ void consumeAndStore(BiConsumer<String, Client> streamCreator, TestInfo info) th
                                     client.credit(subscriptionId, 1))
                             .messageListener(messageListener));
 
-                long offsetToStartFrom = consumer.queryOffset(reference, s) + 1;
+                long offsetToStartFrom = consumer.queryOffset(reference, s).getOffset() + 1;
                 consumer.subscribe(b(0), s, OffsetSpecification.offset(offsetToStartFrom), 1);
 
                 assertThat(consumeLatchSecondWave.await(10, TimeUnit.SECONDS)).isTrue();
@@ -330,7 +331,7 @@ void storeOffsetAndThenAttachByTimestampShouldWork() throws Exception {
 
     IntStream.range(0, messageCount).forEach(i -> client.storeOffset("some reference", stream, i));
 
-    waitAtMost(() -> client.queryOffset("some reference", stream) == messageCount - 1);
+    waitAtMost(() -> client.queryOffset("some reference", stream).getOffset() == messageCount - 1);
 
     confirmLatch.set(new CountDownLatch(messageCount));
 

src/test/java/com/rabbitmq/stream/impl/StreamConsumerTest.java

@@ -502,7 +502,8 @@ void autoTrackingShouldStorePeriodicallyAndAfterInactivity() throws Exception {
     waitAtMost(5, () -> messageCount.get() == storeEvery * 2);
 
     Client client = cf.get();
-    waitAtMost(5, () -> client.queryOffset(reference, stream) == lastReceivedOffset.get());
+    waitAtMost(
+        5, () -> client.queryOffset(reference, stream).getOffset() == lastReceivedOffset.get());
 
     int extraMessages = storeEvery / 10;
     IntStream.range(0, extraMessages)
@@ -514,7 +515,8 @@ void autoTrackingShouldStorePeriodicallyAndAfterInactivity() throws Exception {
 
     waitAtMost(5, () -> messageCount.get() == storeEvery * 2 + extraMessages);
 
-    waitAtMost(5, () -> client.queryOffset(reference, stream) == lastReceivedOffset.get());
+    waitAtMost(
+        5, () -> client.queryOffset(reference, stream).getOffset() == lastReceivedOffset.get());
   }
 
   @Test
@@ -550,7 +552,8 @@ void autoTrackingShouldStoreAfterClosing() throws Exception {
     consumer.close();
 
     Client client = cf.get();
-    waitAtMost(5, () -> client.queryOffset(reference, stream) == lastReceivedOffset.get());
+    waitAtMost(
+        5, () -> client.queryOffset(reference, stream).getOffset() == lastReceivedOffset.get());
   }
 
   @Test

src/test/java/com/rabbitmq/stream/impl/SuperStreamConsumerTest.java

@@ -174,7 +174,9 @@ void manualOffsetTrackingShouldStoreOnAllPartitions() throws Exception {
     // offset near the end (the message count per partition minus a few messages)
     long almostLastOffset = messageCount / partitionCount - messageCount / (partitionCount * 10);
     partitions.forEach(
-        p -> assertThat(client.queryOffset(consumerName, p)).isGreaterThan(almostLastOffset));
+        p ->
+            assertThat(client.queryOffset(consumerName, p).getOffset())
+                .isGreaterThan(almostLastOffset));
     consumer.close();
   }
 
@@ -226,7 +228,9 @@ void autoOffsetTrackingShouldStoreOnAllPartitions() throws Exception {
     // offset near the end (the message count per partition minus a few messages)
     long almostLastOffset = messageCount / partitionCount - messageCount / (partitionCount * 10);
     partitions.forEach(
-        p -> assertThat(client.queryOffset(consumerName, p)).isGreaterThan(almostLastOffset));
+        p ->
+            assertThat(client.queryOffset(consumerName, p).getOffset())
+                .isGreaterThan(almostLastOffset));
     consumer.close();
   }
 }

src/test/java/com/rabbitmq/stream/perf/StreamPerfTestTest.java

@@ -184,9 +184,9 @@ void offsetShouldBeStoredWhenOptionIsEnabled() throws Exception {
     Future<?> run = run(builder().storeEvery(10).consumerNames("consumer-%2$d-on-stream-%1$s"));
     waitUntilStreamExists(s);
     String consumerName = "consumer-1-on-stream-" + s;
-    long offset = client.queryOffset(consumerName, s);
+    long offset = client.queryOffset(consumerName, s).getOffset();
     waitOneSecond();
-    waitAtMost(() -> client.queryOffset(consumerName, s) > offset);
+    waitAtMost(() -> client.queryOffset(consumerName, s).getOffset() > offset);
     run.cancel(true);
     waitRunEnds();
   }
@@ -196,9 +196,9 @@ void offsetShouldNotBeStoredWhenOptionIsNotEnabled() throws Exception {
     Future<?> run = run(builder());
     waitUntilStreamExists(s);
     String consumerName = s + "-0"; // default value when offset tracking is enabled
-    assertThat(client.queryOffset(consumerName, s)).isZero();
+    assertThat(client.queryOffset(consumerName, s).getOffset()).isZero();
     waitOneSecond();
-    assertThat(client.queryOffset(consumerName, s)).isZero();
+    assertThat(client.queryOffset(consumerName, s).getOffset()).isZero();
     run.cancel(true);
     waitRunEnds();
   }