Add SAC test with manual offset tracking

acogoluegnes · acogoluegnes · commit b1861f6e193e · 2022-07-13T09:13:09.000+02:00
Requires to add Consumer#storedOffset() to let the user make sure the stored offset request made it to the broker before closing the consumer. References #46, rabbitmq/rabbitmq-server#3753 Conflicts: src/main/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinator.java
diff --git a/src/main/java/com/rabbitmq/stream/Consumer.java b/src/main/java/com/rabbitmq/stream/Consumer.java
@@ -33,4 +33,6 @@ public interface Consumer extends AutoCloseable {
   /** Close the consumer. */
   @Override
   void close();
+
+  long storedOffset();
 }
diff --git a/src/main/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinator.java b/src/main/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinator.java
@@ -228,7 +228,7 @@ public long flush() {
       } else {
         long result;
         try {
-          long lastStoredOffset = consumer.lastStoredOffset();
+          long lastStoredOffset = consumer.storedOffset();
           if (offsetBefore(lastStoredOffset, lastProcessedOffset.get())) {
             this.consumer.store(this.lastProcessedOffset.get());
           }
@@ -275,7 +275,7 @@ public Runnable closingCallback() {
                   }
                 };
             try {
-              long lastStoredOffset = consumer.lastStoredOffset();
+              long lastStoredOffset = consumer.storedOffset();
               if (offsetBefore(lastStoredOffset, lastProcessedOffset.get())) {
                 LOGGER.debug("Storing {} offset before closing", this.lastProcessedOffset);
                 storageOperation.run();
@@ -318,7 +318,7 @@ public Consumer<Context> postProcessingCallback() {
     public void flushIfNecessary() {
       if (this.clock.time() - this.lastTrackingActivity > this.checkIntervalInNs) {
         try {
-          long lastStoredOffset = consumer.lastStoredOffset();
+          long lastStoredOffset = consumer.storedOffset();
           if (offsetBefore(lastStoredOffset, lastRequestedOffset)) {
             this.consumer.store(this.lastRequestedOffset);
             this.lastTrackingActivity = clock.time();
diff --git a/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java
@@ -19,7 +19,6 @@
 import com.rabbitmq.stream.Constants;
 import com.rabbitmq.stream.Consumer;
 import com.rabbitmq.stream.ConsumerUpdateListener;
-import com.rabbitmq.stream.ConsumerUpdateListener.Status;
 import com.rabbitmq.stream.MessageHandler;
 import com.rabbitmq.stream.MessageHandler.Context;
 import com.rabbitmq.stream.OffsetSpecification;
@@ -138,7 +137,7 @@ class StreamConsumer implements Consumer {
                       && context.status() == ConsumerUpdateListener.Status.ACTIVE) {
                     LOGGER.debug("Looking up offset (stream {})", this.stream);
                     StreamConsumer consumer = (StreamConsumer) context.consumer();
-                    long offset = consumer.lastStoredOffset();
+                    long offset = consumer.storedOffset();
                     LOGGER.debug(
                         "Stored offset is {}, returning the value + 1 to the server", offset);
                     return OffsetSpecification.offset(offset + 1);
@@ -177,7 +176,7 @@ class StreamConsumer implements Consumer {
                       && context.status() == ConsumerUpdateListener.Status.ACTIVE) {
                     LOGGER.debug("Going from passive to active, looking up offset");
                     StreamConsumer consumer = (StreamConsumer) context.consumer();
-                    long offset = consumer.lastStoredOffset();
+                    long offset = consumer.storedOffset();
                     LOGGER.debug(
                         "Stored offset is {}, returning the value + 1 to the server", offset);
                     result = OffsetSpecification.offset(offset + 1);
@@ -238,7 +237,7 @@ void waitForOffsetToBeStored(long expectedStoredOffset) {
         AsyncRetry.asyncRetry(
                 () -> {
                   try {
-                    long lastStoredOffset = lastStoredOffset();
+                    long lastStoredOffset = storedOffset();
                     boolean stored = lastStoredOffset == expectedStoredOffset;
                     LOGGER.debug(
                         "Last stored offset from consumer {} on {} is {}, expecting {}",
@@ -448,7 +447,8 @@ void running() {
     this.status = Status.RUNNING;
   }
 
-  long lastStoredOffset() {
+  @Override
+  public long storedOffset() {
     if (canTrack()) {
       // the client can be null by now, so we catch any exception
       QueryOffsetResponse response;
diff --git a/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java b/src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java
@@ -157,6 +157,12 @@ public void store(long offset) {
         "Consumer#store(long) does not work for super streams, use MessageHandler.Context#storeOffset() instead");
   }
 
+  @Override
+  public long storedOffset() {
+    throw new UnsupportedOperationException(
+        "Consumer#storedOffset() does not work for super streams");
+  }
+
   @Override
   public void close() {
     for (Entry<String, Consumer> entry : consumers.entrySet()) {
diff --git a/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java b/src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java
@@ -130,7 +130,7 @@ void autoShouldStoreAfterSomeInactivity() {
   void autoShouldStoreLastProcessedOffsetOnClosing() {
     Duration checkInterval = Duration.ofMillis(100);
     OffsetTrackingCoordinator coordinator = new OffsetTrackingCoordinator(env, checkInterval);
-    when(consumer.lastStoredOffset()).thenReturn(5L);
+    when(consumer.storedOffset()).thenReturn(5L);
 
     Registration registration =
         coordinator.registerTrackingConsumer(
@@ -148,7 +148,7 @@ void autoShouldStoreLastProcessedOffsetOnClosing() {
   void autoShouldNotStoreLastProcessedOffsetOnClosingIfBehindStoredOffset() {
     Duration checkInterval = Duration.ofMillis(100);
     OffsetTrackingCoordinator coordinator = new OffsetTrackingCoordinator(env, checkInterval);
-    when(consumer.lastStoredOffset()).thenReturn(15L);
+    when(consumer.storedOffset()).thenReturn(15L);
 
     Registration registration =
         coordinator.registerTrackingConsumer(
@@ -200,7 +200,7 @@ void autoShouldNotStoreIfOffsetAlreadyStored() throws Exception {
 
     long storedOffset = 10;
 
-    when(consumer.lastStoredOffset()).thenReturn(storedOffset);
+    when(consumer.storedOffset()).thenReturn(storedOffset);
 
     Duration autoFlushInterval = Duration.ofMillis(checkInterval.toMillis() * 2);
     Consumer<Context> postProcessedMessageCallback =
@@ -223,7 +223,7 @@ void autoShouldNotFlushAfterInactivityIfLastStoreIsOnModulo() throws Exception {
 
     int storeEvery = 10;
 
-    when(consumer.lastStoredOffset()).thenReturn((long) storeEvery - 1);
+    when(consumer.storedOffset()).thenReturn((long) storeEvery - 1);
 
     Duration autoFlushInterval = Duration.ofMillis(checkInterval.toMillis() * 2);
     Consumer<Context> postProcessedMessageCallback =
@@ -252,7 +252,7 @@ void autoShouldStoreLastProcessedAfterInactivity() {
     int extraMessages = 3;
 
     long expectedLastStoredOffset = storeEvery + extraMessages - 1;
-    when(consumer.lastStoredOffset()).thenReturn((long) (storeEvery - 1));
+    when(consumer.storedOffset()).thenReturn((long) (storeEvery - 1));
 
     ArgumentCaptor<Long> lastStoredOffsetCaptor = ArgumentCaptor.forClass(Long.class);
     CountDownLatch flushLatch = new CountDownLatch(1);
@@ -285,7 +285,7 @@ void manualShouldNotStoreIfAlreadyUpToDate() throws Exception {
     OffsetTrackingCoordinator coordinator = new OffsetTrackingCoordinator(env, checkInterval);
 
     long lastStoredOffset = 50;
-    when(consumer.lastStoredOffset()).thenReturn(lastStoredOffset);
+    when(consumer.storedOffset()).thenReturn(lastStoredOffset);
 
     LongConsumer storeCallback =
         coordinator
@@ -309,7 +309,7 @@ void manualShouldStoreIfRequestedStoredOffsetIsBehind() {
 
     long lastRequestedOffset = 50;
     long lastStoredOffset = 40;
-    when(consumer.lastStoredOffset()).thenReturn(lastStoredOffset);
+    when(consumer.storedOffset()).thenReturn(lastStoredOffset);
 
     ArgumentCaptor<Long> lastStoredOffsetCaptor = ArgumentCaptor.forClass(Long.class);
     CountDownLatch storeLatch = new CountDownLatch(1);
diff --git a/src/test/java/com/rabbitmq/stream/impl/StreamConsumerSacTest.java b/src/test/java/com/rabbitmq/stream/impl/StreamConsumerSacTest.java
@@ -108,6 +108,72 @@ void autoTrackingSecondConsumerShouldTakeOverWhereTheFirstOneLeftOff() throws Ex
     consumer2.close();
   }
 
+  @Test
+  void manualTrackingSecondConsumerShouldTakeOverWhereTheFirstOneLeftOff() throws Exception {
+    int messageCount = 10000;
+    int storeEvery = 1000;
+    Map<Integer, AtomicInteger> receivedMessages = new ConcurrentHashMap<>();
+    receivedMessages.put(0, new AtomicInteger(0));
+    receivedMessages.put(1, new AtomicInteger(0));
+    AtomicLong lastReceivedOffset = new AtomicLong(0);
+    String consumerName = "foo";
+
+    Consumer consumer1 =
+        environment.consumerBuilder().stream(stream)
+            .name(consumerName)
+            .singleActiveConsumer()
+            .messageHandler(
+                (context, message) -> {
+                  lastReceivedOffset.set(context.offset());
+                  int count = receivedMessages.get(0).incrementAndGet();
+                  if (count % storeEvery == 0) {
+                    context.storeOffset();
+                  }
+                })
+            .offset(OffsetSpecification.first())
+            .manualTrackingStrategy()
+            .builder()
+            .build();
+
+    Consumer consumer2 =
+        environment.consumerBuilder().stream(stream)
+            .name(consumerName)
+            .singleActiveConsumer()
+            .messageHandler(
+                (context, message) -> {
+                  lastReceivedOffset.set(context.offset());
+                  int count = receivedMessages.get(1).incrementAndGet();
+                  if (count % storeEvery == 0) {
+                    context.storeOffset();
+                  }
+                })
+            .offset(OffsetSpecification.first())
+            .manualTrackingStrategy()
+            .builder()
+            .build();
+
+    publishAndWaitForConfirms(cf, messageCount, stream);
+    waitAtMost(() -> receivedMessages.getOrDefault(0, new AtomicInteger(0)).get() == messageCount);
+
+    assertThat(lastReceivedOffset).hasPositiveValue();
+    assertThat(receivedMessages.get(1)).hasValue(0);
+
+    long firstWaveLimit = lastReceivedOffset.get();
+
+    consumer1.store(firstWaveLimit);
+    waitAtMost(() -> consumer1.storedOffset() == firstWaveLimit);
+
+    consumer1.close();
+
+    publishAndWaitForConfirms(cf, messageCount, stream);
+
+    waitAtMost(() -> receivedMessages.getOrDefault(0, new AtomicInteger(1)).get() == messageCount);
+    assertThat(lastReceivedOffset).hasValueGreaterThan(firstWaveLimit);
+    assertThat(receivedMessages.get(0)).hasValue(messageCount);
+
+    consumer2.close();
+  }
+
   @Test
   void customTrackingSecondConsumerShouldTakeOverWhereTheFirstOneLeftOff() throws Exception {
     int messageCount = 10000;

Original file line number	Diff line number	Diff line change
`@@ -33,4 +33,6 @@ public interface Consumer extends AutoCloseable {`
`33`	`33`	`/** Close the consumer. */`
`34`	`34`	`@Override`
`35`	`35`	`void close();`
	`36`	`+`
	`37`	`+ long storedOffset();`
`36`	`38`	`}`