Support clean SAC change in super stream consumer

References rabbitmq/rabbitmq-server#3753

Conflicts:
	src/main/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinator.java
	src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java
	src/main/java/com/rabbitmq/stream/impl/Utils.java

Conflicts:
	src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java
	src/main/java/com/rabbitmq/stream/impl/StreamConsumer.java

Conflicts:
	src/main/java/com/rabbitmq/stream/MessageHandler.java
	src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java
	src/main/java/com/rabbitmq/stream/impl/Utils.java
	src/test/java/com/rabbitmq/stream/impl/OffsetTrackingCoordinatorTest.java

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/MessageHandler.java

@@ -70,6 +70,13 @@ interface Context {
      */
     long committedChunkId();
 
+    /**
+     * The stream the message comes from.
+     *
+     * @return the stream the message comes from
+     */
+    String stream();
+
     /**
      * The consumer that receives the message.
      *

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

@@ -95,6 +95,7 @@ Runnable subscribe(
       OffsetSpecification offsetSpecification,
       String trackingReference,
       SubscriptionListener subscriptionListener,
+      Runnable trackingClosingCallback,
       MessageHandler messageHandler,
       Map<String, String> subscriptionProperties) {
     // FIXME fail immediately if there's no locator (can provide a supplier that does not retry)
@@ -113,6 +114,7 @@ Runnable subscribe(
             offsetSpecification,
             trackingReference,
             subscriptionListener,
+            trackingClosingCallback,
             messageHandler,
             subscriptionProperties);
 
@@ -235,6 +237,7 @@ private static class SubscriptionTracker {
     private final MessageHandler messageHandler;
     private final StreamConsumer consumer;
     private final SubscriptionListener subscriptionListener;
+    private final Runnable trackingClosingCallback;
     private final Map<String, String> subscriptionProperties;
     private volatile long offset;
     private volatile boolean hasReceivedSomething = false;
@@ -247,13 +250,15 @@ private SubscriptionTracker(
         OffsetSpecification initialOffsetSpecification,
         String offsetTrackingReference,
         SubscriptionListener subscriptionListener,
+        Runnable trackingClosingCallback,
         MessageHandler messageHandler,
         Map<String, String> subscriptionProperties) {
       this.consumer = consumer;
       this.stream = stream;
       this.initialOffsetSpecification = initialOffsetSpecification;
       this.offsetTrackingReference = offsetTrackingReference;
       this.subscriptionListener = subscriptionListener;
+      this.trackingClosingCallback = trackingClosingCallback;
       this.messageHandler = messageHandler;
       if (this.offsetTrackingReference == null) {
         this.subscriptionProperties = subscriptionProperties;
@@ -267,6 +272,11 @@ private SubscriptionTracker(
     }
 
     synchronized void cancel() {
+      // the flow of messages in the user message handler should stop, we can call the tracking
+      // closing callback
+      // with automatic offset tracking, it will store the last dispatched offset
+      LOGGER.debug("Calling tracking consumer closing callback (may be no-op)");
+      this.trackingClosingCallback.run();
       if (this.manager != null) {
         LOGGER.debug("Removing consumer from manager " + this.consumer);
         this.manager.remove(this);
@@ -298,10 +308,10 @@ private static final class MessageHandlerContext implements Context {
     private final long offset;
     private final long timestamp;
     private final long committedOffset;
-    private final Consumer consumer;
+    private final StreamConsumer consumer;
 
     private MessageHandlerContext(
-        long offset, long timestamp, long committedOffset, Consumer consumer) {
+        long offset, long timestamp, long committedOffset, StreamConsumer consumer) {
       this.offset = offset;
       this.timestamp = timestamp;
       this.committedOffset = committedOffset;
@@ -328,6 +338,10 @@ public long committedChunkId() {
       return committedOffset;
     }
 
+    public String stream() {
+      return this.consumer.stream();
+    }
+
     @Override
     public Consumer consumer() {
       return this.consumer;
@@ -455,14 +469,14 @@ private ClientSubscriptionsManager(
           (subscriptionId, offset, chunkTimestamp, committedOffset, message) -> {
             SubscriptionTracker subscriptionTracker =
                 subscriptionTrackers.get(subscriptionId & 0xFF);
-            if (subscriptionTracker != null) {
+            if (subscriptionTracker != null && subscriptionTracker.consumer.isOpen()) {
               subscriptionTracker.offset = offset;
               subscriptionTracker.hasReceivedSomething = true;
               subscriptionTracker.messageHandler.handle(
                   new MessageHandlerContext(
                       offset, chunkTimestamp, committedOffset, subscriptionTracker.consumer),
                   message);
-              // FIXME set offset here as well, best effort to avoid duplicates
+              // FIXME set offset here as well, best effort to avoid duplicates?
             } else {
               LOGGER.debug("Could not find stream subscription {}", subscriptionId);
             }
@@ -741,6 +755,8 @@ void add(
           offsetSpecification =
               offsetSpecification == null ? DEFAULT_OFFSET_SPECIFICATION : offsetSpecification;
 
+          // TODO consider using/emulating ConsumerUpdateListener, to have only one API, not 2
+          // even when the consumer is not a SAC.
           SubscriptionContext subscriptionContext =
               new DefaultSubscriptionContext(offsetSpecification);
           subscriptionTracker.subscriptionListener.preSubscribe(subscriptionContext);

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

@@ -16,6 +16,7 @@
 import static com.rabbitmq.stream.impl.Utils.offsetBefore;
 
 import com.rabbitmq.stream.Constants;
+import com.rabbitmq.stream.ConsumerUpdateListener.Status;
 import com.rabbitmq.stream.MessageHandler.Context;
 import com.rabbitmq.stream.StreamException;
 import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration;
@@ -258,20 +259,34 @@ public LongConsumer trackingCallback() {
     @Override
     public Runnable closingCallback() {
       return () -> {
-        if (this.lastProcessedOffset == null) {
-          LOGGER.debug("Not storing anything as nothing has been processed.");
+        if (this.consumer.isSac() && this.consumer.sacStatus() != Status.ACTIVE) {
+          LOGGER.debug("Not storing offset on closing because consumer is a non-active SAC");
         } else {
-          try {
-            long lastStoredOffset = consumer.lastStoredOffset();
-            if (offsetBefore(lastStoredOffset, lastProcessedOffset.get())) {
-              LOGGER.debug("Storing {} offset before closing", this.lastProcessedOffset);
-              this.consumer.store(this.lastProcessedOffset.get());
-            }
-          } catch (StreamException e) {
-            if (e.getCode() == Constants.RESPONSE_CODE_NO_OFFSET) {
-              LOGGER.debug(
-                  "Nothing stored yet, storing {} offset before closing", this.lastProcessedOffset);
-              this.consumer.store(this.lastProcessedOffset.get());
+          if (this.lastProcessedOffset == null) {
+            LOGGER.debug("Not storing anything as nothing has been processed.");
+          } else {
+            Runnable storageOperation =
+                () -> {
+                  this.consumer.store(this.lastProcessedOffset.get());
+                  if (this.consumer.isSac()) {
+                    LOGGER.debug(
+                        "Consumer is SAC, making sure offset has been stored, in case another SAC takes over");
+                    this.consumer.waitForOffsetToBeStored(lastProcessedOffset.get());
+                  }
+                };
+            try {
+              long lastStoredOffset = consumer.lastStoredOffset();
+              if (offsetBefore(lastStoredOffset, lastProcessedOffset.get())) {
+                LOGGER.debug("Storing {} offset before closing", this.lastProcessedOffset);
+                storageOperation.run();
+              }
+            } catch (StreamException e) {
+              if (e.getCode() == Constants.RESPONSE_CODE_NO_OFFSET) {
+                LOGGER.debug(
+                    "Nothing stored yet, storing {} offset before closing",
+                    this.lastProcessedOffset);
+                storageOperation.run();
+              }
             }
           }
         }