Merge pull request #779 from rabbitmq/environment-store-offset

Add Environment#storeOffset(String,String,long)

Author: acogoluegnes

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

@@ -83,6 +83,22 @@ static EnvironmentBuilder builder() {
    */
   StreamStats queryStreamStats(String stream);
 
+  /**
+   * Store the offset for a given reference on the given stream.
+   *
+   * <p>This method is useful to store a given offset before a consumer is created.
+   *
+   * <p>Prefer the {@link Consumer#store(long)} or {@link MessageHandler.Context#storeOffset()}
+   * methods to store offsets while consuming messages.
+   *
+   * @see Consumer#store(long)
+   * @see MessageHandler.Context#storeOffset()
+   * @param reference the reference to store the offset for, e.g. a consumer name
+   * @param stream the stream
+   * @param offset the offset to store
+   */
+  void storeOffset(String reference, String stream, long offset);
+
   /**
    * Return whether a stream exists or not.
    *

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

@@ -0,0 +1,125 @@
+// Copyright (c) 2025 Broadcom. All Rights Reserved.
+// The term "Broadcom" refers to Broadcom Inc. and/or its subsidiaries.
+//
+// This software, the RabbitMQ Stream Java client library, is dual-licensed under the
+// Mozilla Public License 2.0 ("MPL"), and the Apache License version 2 ("ASL").
+// For the MPL, please see LICENSE-MPL-RabbitMQ. For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+package com.rabbitmq.stream.impl;
+
+import static com.rabbitmq.stream.BackOffDelayPolicy.fixedWithInitialDelay;
+import static com.rabbitmq.stream.impl.AsyncRetry.asyncRetry;
+import static java.lang.String.format;
+import static java.time.Duration.ofMillis;
+
+import com.rabbitmq.stream.Constants;
+import com.rabbitmq.stream.NoOffsetException;
+import com.rabbitmq.stream.StreamException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.LongSupplier;
+import java.util.function.Supplier;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+class OffsetTrackingUtils {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(OffsetTrackingUtils.class);
+
+  private OffsetTrackingUtils() {}
+
+  static long storedOffset(Supplier<Client> clientSupplier, String name, String stream) {
+    // the client can be null, so we catch any exception
+    Client.QueryOffsetResponse response;
+    try {
+      response = clientSupplier.get().queryOffset(name, stream);
+    } catch (Exception e) {
+      throw new IllegalStateException(
+          format(
+              "Not possible to query offset for name %s on stream %s for now: %s",
+              name, stream, e.getMessage()),
+          e);
+    }
+    if (response.isOk()) {
+      return response.getOffset();
+    } else if (response.getResponseCode() == Constants.RESPONSE_CODE_NO_OFFSET) {
+      throw new NoOffsetException(
+          format(
+              "No offset stored for name %s on stream %s (%s)",
+              name, stream, Utils.formatConstant(response.getResponseCode())));
+    } else {
+      throw new StreamException(
+          format(
+              "QueryOffset for name %s on stream %s returned an error (%s)",
+              name, stream, Utils.formatConstant(response.getResponseCode())),
+          response.getResponseCode());
+    }
+  }
+
+  static void waitForOffsetToBeStored(
+      String caller,
+      ScheduledExecutorService scheduledExecutorService,
+      LongSupplier offsetSupplier,
+      String name,
+      String stream,
+      long expectedStoredOffset) {
+    String reference = String.format("{stream=%s/name=%s}", stream, name);
+    CompletableFuture<Boolean> storedTask =
+        asyncRetry(
+                () -> {
+                  try {
+                    long lastStoredOffset = offsetSupplier.getAsLong();
+                    boolean stored = lastStoredOffset == expectedStoredOffset;
+                    LOGGER.debug(
+                        "Last stored offset from {} on {} is {}, expecting {}",
+                        caller,
+                        reference,
+                        lastStoredOffset,
+                        expectedStoredOffset);
+                    if (!stored) {
+                      throw new IllegalStateException();
+                    } else {
+                      return true;
+                    }
+                  } catch (StreamException e) {
+                    if (e.getCode() == Constants.RESPONSE_CODE_NO_OFFSET) {
+                      LOGGER.debug(
+                          "No stored offset for {} on {}, expecting {}",
+                          caller,
+                          reference,
+                          expectedStoredOffset);
+                      throw new IllegalStateException();
+                    } else {
+                      throw e;
+                    }
+                  }
+                })
+            .description(
+                "Last stored offset for %s on %s must be %d",
+                caller, reference, expectedStoredOffset)
+            .delayPolicy(fixedWithInitialDelay(ofMillis(200), ofMillis(200)))
+            .retry(exception -> exception instanceof IllegalStateException)
+            .scheduler(scheduledExecutorService)
+            .build();
+
+    try {
+      storedTask.get(10, TimeUnit.SECONDS);
+      LOGGER.debug("Offset {} stored ({}, {})", expectedStoredOffset, caller, reference);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    } catch (ExecutionException | TimeoutException e) {
+      LOGGER.warn("Error while checking offset has been stored", e);
+      storedTask.cancel(true);
+    }
+  }
+}

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

@@ -18,11 +18,9 @@
 import static com.rabbitmq.stream.impl.AsyncRetry.asyncRetry;
 import static com.rabbitmq.stream.impl.Utils.offsetBefore;
 import static java.lang.String.format;
-import static java.time.Duration.ofMillis;
 
 import com.rabbitmq.stream.*;
 import com.rabbitmq.stream.MessageHandler.Context;
-import com.rabbitmq.stream.impl.Client.QueryOffsetResponse;
 import com.rabbitmq.stream.impl.StreamConsumerBuilder.TrackingConfiguration;
 import com.rabbitmq.stream.impl.StreamEnvironment.LocatorNotAvailableException;
 import com.rabbitmq.stream.impl.StreamEnvironment.TrackingConsumerRegistration;
@@ -329,53 +327,13 @@ static long getStoredOffsetSafely(StreamConsumer consumer, StreamEnvironment env
   }
 
   void waitForOffsetToBeStored(long expectedStoredOffset) {
-    CompletableFuture<Boolean> storedTask =
-        asyncRetry(
-                () -> {
-                  try {
-                    long lastStoredOffset = storedOffset();
-                    boolean stored = lastStoredOffset == expectedStoredOffset;
-                    LOGGER.debug(
-                        "Last stored offset from consumer {} on {} is {}, expecting {}",
-                        this.id,
-                        this.stream,
-                        lastStoredOffset,
-                        expectedStoredOffset);
-                    if (!stored) {
-                      throw new IllegalStateException();
-                    } else {
-                      return true;
-                    }
-                  } catch (StreamException e) {
-                    if (e.getCode() == Constants.RESPONSE_CODE_NO_OFFSET) {
-                      LOGGER.debug(
-                          "No stored offset for consumer {} on {}, expecting {}",
-                          this.id,
-                          this.stream,
-                          expectedStoredOffset);
-                      throw new IllegalStateException();
-                    } else {
-                      throw e;
-                    }
-                  }
-                })
-            .description(
-                "Last stored offset for consumer %s on stream %s must be %d",
-                this.name, this.stream, expectedStoredOffset)
-            .delayPolicy(fixedWithInitialDelay(ofMillis(200), ofMillis(200)))
-            .retry(exception -> exception instanceof IllegalStateException)
-            .scheduler(environment.scheduledExecutorService())
-            .build();
-
-    try {
-      storedTask.get(10, TimeUnit.SECONDS);
-      LOGGER.debug(
-          "Offset {} stored (consumer {}, stream {})", expectedStoredOffset, this.id, this.stream);
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-    } catch (ExecutionException | TimeoutException e) {
-      LOGGER.warn("Error while checking offset has been stored", e);
-    }
+    OffsetTrackingUtils.waitForOffsetToBeStored(
+        "consumer " + this.id,
+        this.environment.scheduledExecutorService(),
+        this::storedOffset,
+        this.name,
+        this.stream,
+        expectedStoredOffset);
   }
 
   void start() {
@@ -563,31 +521,7 @@ void running() {
   long storedOffset(Supplier<Client> clientSupplier) {
     checkNotClosed();
     if (canTrack()) {
-      // the client can be null by now, so we catch any exception
-      QueryOffsetResponse response;
-      try {
-        response = clientSupplier.get().queryOffset(this.name, this.stream);
-      } catch (Exception e) {
-        throw new IllegalStateException(
-            format(
-                "Not possible to query offset for consumer %s on stream %s for now: %s",
-                this.name, this.stream, e.getMessage()),
-            e);
-      }
-      if (response.isOk()) {
-        return response.getOffset();
-      } else if (response.getResponseCode() == Constants.RESPONSE_CODE_NO_OFFSET) {
-        throw new NoOffsetException(
-            format(
-                "No offset stored for consumer %s on stream %s (%s)",
-                this.name, this.stream, Utils.formatConstant(response.getResponseCode())));
-      } else {
-        throw new StreamException(
-            format(
-                "QueryOffset for consumer %s on stream %s returned an error (%s)",
-                this.name, this.stream, Utils.formatConstant(response.getResponseCode())),
-            response.getResponseCode());
-      }
+      return OffsetTrackingUtils.storedOffset(clientSupplier, this.name, this.stream);
     } else if (this.name == null) {
       throw new UnsupportedOperationException(
           "Not possible to query stored offset for a consumer without a name");

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

@@ -556,6 +556,29 @@ public StreamStats queryStreamStats(String stream) {
     }
   }
 
+  @Override
+  public void storeOffset(String reference, String stream, long offset) {
+    checkNotClosed();
+    this.maybeInitializeLocator();
+    locatorOperation(
+        Utils.namedFunction(
+            l -> {
+              l.storeOffset(reference, stream, offset);
+              return null;
+            },
+            "Store offset %d for stream '%s' with reference '%s'",
+            offset,
+            stream,
+            reference));
+    OffsetTrackingUtils.waitForOffsetToBeStored(
+        "env-store-offset",
+        this.scheduledExecutorService,
+        () -> OffsetTrackingUtils.storedOffset(() -> locator().client(), reference, stream),
+        reference,
+        stream,
+        offset);
+  }
+
   @Override
   public boolean streamExists(String stream) {
     checkNotClosed();

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

@@ -16,17 +16,58 @@
 
 import static org.assertj.core.api.Assertions.fail;
 
+import com.rabbitmq.stream.Constants;
 import java.time.Duration;
 import org.assertj.core.api.AbstractObjectAssert;
 
 final class Assertions {
 
   private Assertions() {}
 
+  static ResponseAssert assertThat(Client.Response response) {
+    return new ResponseAssert(response);
+  }
+
   static SyncAssert assertThat(TestUtils.Sync sync) {
     return new SyncAssert(sync);
   }
 
+  static class ResponseAssert extends AbstractObjectAssert<ResponseAssert, Client.Response> {
+
+    public ResponseAssert(Client.Response response) {
+      super(response, ResponseAssert.class);
+    }
+
+    ResponseAssert isOk() {
+      if (!actual.isOk()) {
+        fail(
+            "Response should be successful but was not, response code is: %s",
+            Utils.formatConstant(actual.getResponseCode()));
+      }
+      return this;
+    }
+
+    ResponseAssert isNotOk() {
+      if (actual.isOk()) {
+        fail("Response should not be successful but was, response code is: %s", actual);
+      }
+      return this;
+    }
+
+    ResponseAssert hasCode(short responseCode) {
+      if (actual.getResponseCode() != responseCode) {
+        fail(
+            "Response code should be %s but was %s",
+            Utils.formatConstant(responseCode), Utils.formatConstant(actual.getResponseCode()));
+      }
+      return this;
+    }
+
+    ResponseAssert hasCodeNoOffset() {
+      return hasCode(Constants.RESPONSE_CODE_NO_OFFSET);
+    }
+  }
+
   static class SyncAssert extends AbstractObjectAssert<SyncAssert, TestUtils.Sync> {
 
     private SyncAssert(TestUtils.Sync sync) {