Store offset on closing with automatic offset tracking

Fixes #40

Author: acogoluegnes

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

@@ -25,9 +25,13 @@
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.function.LongConsumer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 class OffsetTrackingCoordinator {
 
+  private static final Logger LOGGER = LoggerFactory.getLogger(OffsetTrackingCoordinator.class);
+
   private final StreamEnvironment streamEnvironment;
 
   private final AtomicBoolean started = new AtomicBoolean(false);
@@ -106,7 +110,7 @@ Registration registerTrackingConsumer(
                   TimeUnit.MILLISECONDS);
     }
 
-    return new Registration(tracker.postProcessingCallback(), tracker.trackingCallback());
+    return new Registration(tracker);
   }
 
   private ScheduledExecutorService executor() {
@@ -134,24 +138,28 @@ private interface Tracker {
     StreamConsumer consumer();
 
     LongConsumer trackingCallback();
+
+    Runnable closingCallback();
   }
 
   static class Registration {
 
-    private final java.util.function.Consumer<Context> postMessageProcessingCallback;
-    private final LongConsumer trackingCallback;
+    private final Tracker tracker;
 
-    Registration(Consumer<Context> postMessageProcessingCallback, LongConsumer trackingCallback) {
-      this.postMessageProcessingCallback = postMessageProcessingCallback;
-      this.trackingCallback = trackingCallback;
+    Registration(Tracker tracker) {
+      this.tracker = tracker;
     }
 
-    public Consumer<Context> postMessageProcessingCallback() {
-      return postMessageProcessingCallback;
+    Consumer<Context> postMessageProcessingCallback() {
+      return this.tracker.postProcessingCallback();
     }
 
-    public LongConsumer trackingCallback() {
-      return trackingCallback;
+    LongConsumer trackingCallback() {
+      return this.tracker.trackingCallback();
+    }
+
+    Runnable closingCallback() {
+      return this.tracker.closingCallback();
     }
   }
 
@@ -204,6 +212,22 @@ public StreamConsumer consumer() {
     public LongConsumer trackingCallback() {
       return Utils.NO_OP_LONG_CONSUMER;
     }
+
+    @Override
+    public Runnable closingCallback() {
+      return () -> {
+        long lastStoredOffset = consumer.lastStoredOffset();
+        if (lastStoredOffset < lastProcessedOffset) {
+          LOGGER.debug("Storing offset before closing");
+          this.consumer.store(this.lastProcessedOffset);
+        } else {
+          LOGGER.debug(
+              "Not storing offset before closing because last stored offset after last processed offset: {} > {}",
+              lastStoredOffset,
+              lastProcessedOffset);
+        }
+      };
+    }
   }
 
   private static final class ManualTrackingTracker implements Tracker {
@@ -249,6 +273,11 @@ public LongConsumer trackingCallback() {
         lastTrackingActivity = clock.time();
       };
     }
+
+    @Override
+    public Runnable closingCallback() {
+      return () -> {};
+    }
   }
 
   private static class LocalClock {

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

@@ -551,14 +551,33 @@ Client.ClientParameters clientParametersCopy() {
   TrackingConsumerRegistration registerTrackingConsumer(
       StreamConsumer streamConsumer, TrackingConfiguration configuration) {
     Runnable closingCallable = this.producersCoordinator.registerTrackingConsumer(streamConsumer);
-    Registration offsetTrackingRegistration = null;
+    Registration offsetTrackingRegistration;
     if (this.offsetTrackingCoordinator.needTrackingRegistration(configuration)) {
       offsetTrackingRegistration =
           this.offsetTrackingCoordinator.registerTrackingConsumer(streamConsumer, configuration);
+    } else {
+      offsetTrackingRegistration = null;
+    }
+
+    Runnable closingSequence;
+    if (offsetTrackingRegistration == null) {
+      closingSequence = closingCallable;
+    } else {
+      closingSequence =
+          () -> {
+            try {
+              LOGGER.debug("Executing offset tracking registration closing sequence");
+              offsetTrackingRegistration.closingCallback().run();
+              LOGGER.debug("Offset tracking registration closing sequence executed");
+            } catch (Exception e) {
+              LOGGER.warn("Error while executing offset tracking registration closing sequence");
+            }
+            closingCallable.run();
+          };
     }
 
     return new TrackingConsumerRegistration(
-        closingCallable,
+        closingSequence,
         offsetTrackingRegistration == null
             ? null
             : offsetTrackingRegistration.postMessageProcessingCallback(),

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

@@ -19,6 +19,7 @@
 import static org.mockito.Mockito.*;
 
 import com.rabbitmq.stream.MessageHandler.Context;
+import com.rabbitmq.stream.impl.OffsetTrackingCoordinator.Registration;
 import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration;
 import java.time.Duration;
 import java.util.concurrent.CountDownLatch;
@@ -123,6 +124,42 @@ void autoShouldStoreAfterSomeInactivity() {
     assertThat(latchAssert(flushLatch)).completes(5);
   }
 
+  @Test
+  void autoShouldStoreLastProcessedOffsetOnClosing() {
+    Duration checkInterval = Duration.ofMillis(100);
+    OffsetTrackingCoordinator coordinator = new OffsetTrackingCoordinator(env, checkInterval);
+    when(consumer.lastStoredOffset()).thenReturn(5L);
+
+    Registration registration =
+        coordinator.registerTrackingConsumer(
+            consumer, new TrackingConfiguration(true, true, 1, Duration.ofHours(1), Duration.ZERO));
+
+    long lastProcessedOffset = 10;
+    registration.postMessageProcessingCallback().accept(context(lastProcessedOffset, () -> {}));
+
+    registration.closingCallback().run();
+
+    verify(consumer, times(1)).store(lastProcessedOffset);
+  }
+
+  @Test
+  void autoShouldNotStoreLastProcessedOffsetOnClosingIfBehindStoredOffset() {
+    Duration checkInterval = Duration.ofMillis(100);
+    OffsetTrackingCoordinator coordinator = new OffsetTrackingCoordinator(env, checkInterval);
+    when(consumer.lastStoredOffset()).thenReturn(15L);
+
+    Registration registration =
+        coordinator.registerTrackingConsumer(
+            consumer, new TrackingConfiguration(true, true, 1, Duration.ofHours(1), Duration.ZERO));
+
+    long lastProcessedOffset = 10;
+    registration.postMessageProcessingCallback().accept(context(lastProcessedOffset, () -> {}));
+
+    registration.closingCallback().run();
+
+    verify(consumer, never()).store(anyLong());
+  }
+
   @Test
   void autoShouldStoreFixedMessageCountAndAutoTrackingAfterInactivity() {
     Duration checkInterval = Duration.ofMillis(100);

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

@@ -517,6 +517,42 @@ void autoTrackingShouldStorePeriodicallyAndAfterInactivity() throws Exception {
     waitAtMost(5, () -> client.queryOffset(reference, stream) == lastReceivedOffset.get());
   }
 
+  @Test
+  void autoTrackingShouldStoreAfterClosing() throws Exception {
+    int storeEvery = 10_000;
+    int messageCount = storeEvery * 5 - 100;
+    CountDownLatch consumeLatch = new CountDownLatch(messageCount);
+    String reference = "ref-1";
+    AtomicLong lastReceivedOffset = new AtomicLong(0);
+    Consumer consumer =
+        environment.consumerBuilder().name(reference).stream(stream)
+            .offset(OffsetSpecification.first())
+            .messageHandler(
+                (context, message) -> {
+                  lastReceivedOffset.set(context.offset());
+                  consumeLatch.countDown();
+                })
+            .autoTrackingStrategy()
+            .flushInterval(Duration.ofHours(1)) // long flush interval
+            .messageCountBeforeStorage(storeEvery)
+            .builder()
+            .build();
+
+    Producer producer = environment.producerBuilder().stream(stream).build();
+    IntStream.range(0, messageCount)
+        .forEach(
+            i ->
+                producer.send(
+                    producer.messageBuilder().addData("".getBytes()).build(),
+                    confirmationStatus -> {}));
+
+    latchAssert(consumeLatch).completes();
+    consumer.close();
+
+    Client client = cf.get();
+    waitAtMost(5, () -> client.queryOffset(reference, stream) == lastReceivedOffset.get());
+  }
+
   @Test
   @DisabledIfRabbitMqCtlNotSet
   void externalOffsetTrackingWithSubscriptionListener() throws Exception {
@@ -560,7 +596,8 @@ void externalOffsetTrackingWithSubscriptionListener() throws Exception {
     assertThat(offsetTracking.get()).isGreaterThanOrEqualTo(messageCount - 1);
 
     Host.killConnection("rabbitmq-stream-consumer");
-    waitAtMost(recoveryInitialDelay.multipliedBy(2), () -> subscriptionListenerCallCount.get() == 2);
+    waitAtMost(
+        recoveryInitialDelay.multipliedBy(2), () -> subscriptionListenerCallCount.get() == 2);
 
     publish.run();
     waitAtMost(5, () -> receivedMessages.get() == messageCount * 2);