Merge pull request #63 from rabbitmq/consumer-subscription-listener

Add subscription listener

Author: acogoluegnes

src/docs/asciidoc/usage.adoc

@@ -26,6 +26,26 @@ include::{test-examples}/Api.java[tag=connection-settings]
 <1> Use the `guest` user by default
 <2> Use the `admin` user for this connection
 
+=== Subscription Listener
+
+The client provides a `SubscriptionListener` interface callback to add behavior before a subscription is created.
+This callback is meant for stream consumers: it can be used to dynamically set the offset the consumer attaches to in the stream.
+It is called when the consumer is first created and when the client has to re-subscribe (e.g. after a disconnection).
+
+It is possible to use the callback to get the last processed offset from an external store.
+The following code snippet shows how this can be done (note the interaction with the external store is not detailed):
+
+.Using the subscription listener to attach to a stream
+[source,java,indent=0]
+--------
+include::{test-examples}/Api.java[tag=subscription-listener]
+--------
+<1> Set subscription listener
+<2> Get offset from external store
+<3> Set offset to use for the subscription
+<4> Get the message offset
+<5> Store the offset in the external store after processing
+
 === Metrics Collection
 
 The library provides the {javadoc-url}/com/rabbitmq/client/amqp/metrics/MetricsCollector.html[`MetricsCollector`] abstraction to collect metrics.

src/main/java/com/rabbitmq/client/amqp/ConsumerBuilder.java

@@ -74,6 +74,17 @@ public interface ConsumerBuilder {
    */
   StreamOptions stream();
 
+  /**
+   * Set a listener to customize the subscription before the consumer is created (or recovered).
+   *
+   * <p>This callback is available for stream consumers.
+   *
+   * @param subscriptionListener subscription listener
+   * @return this builder instance
+   * @see SubscriptionListener
+   */
+  ConsumerBuilder subscriptionListener(SubscriptionListener subscriptionListener);
+
   /**
    * Build the consumer.
    *
@@ -164,4 +175,41 @@ enum StreamOffsetSpecification {
     /** Very end of the stream (new chunks). */
     NEXT
   }
+
+  /**
+   * Callback to modify a consumer subscription before the link creation.
+   *
+   * <p>This allows looking up the last processed offset for a stream consumer and attaching to this
+   * offset.
+   */
+  interface SubscriptionListener {
+
+    /**
+     * Pre-subscription callback.
+     *
+     * <p>It is called before the link is created but also every time it recovers, e.g. after a
+     * connection failure.
+     *
+     * <p>Configuration set with {@link Context#streamOptions()} overrides the one set with {@link
+     * ConsumerBuilder#stream()}.
+     *
+     * @param context subscription context
+     */
+    void preSubscribe(Context context);
+
+    /** Subscription context. */
+    interface Context {
+
+      /**
+       * Stream options, to set the offset to start consuming from.
+       *
+       * <p>Only the {@link StreamOptions} are accessible, the {@link StreamOptions#builder()}
+       * method returns <code>null</code>
+       *
+       * @return the stream options
+       * @see StreamOptions
+       */
+      ConsumerBuilder.StreamOptions streamOptions();
+    }
+  }
 }

src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumer.java

@@ -18,15 +18,16 @@
 package com.rabbitmq.client.amqp.impl;
 
 import static com.rabbitmq.client.amqp.Resource.State.*;
+import static com.rabbitmq.client.amqp.impl.AmqpConsumerBuilder.*;
 import static java.time.Duration.ofSeconds;
+import static java.util.Optional.ofNullable;
 
 import com.rabbitmq.client.amqp.AmqpException;
 import com.rabbitmq.client.amqp.BackOffDelayPolicy;
 import com.rabbitmq.client.amqp.Consumer;
+import com.rabbitmq.client.amqp.ConsumerBuilder;
 import com.rabbitmq.client.amqp.metrics.MetricsCollector;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
+import java.util.*;
 import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
@@ -59,6 +60,7 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
   private final String queue;
   private final Map<String, Object> filters;
   private final Map<String, Object> linkProperties;
+  private final ConsumerBuilder.SubscriptionListener subscriptionListener;
   private final AmqpConnection connection;
   private final AtomicReference<PauseStatus> pauseStatus =
       new AtomicReference<>(PauseStatus.UNPAUSED);
@@ -89,11 +91,17 @@ final class AmqpConsumer extends ResourceBase implements Consumer {
     this.queue = builder.queue();
     this.filters = Map.copyOf(builder.filters());
     this.linkProperties = Map.copyOf(builder.properties());
+    this.subscriptionListener =
+        ofNullable(builder.subscriptionListener()).orElse(NO_OP_SUBSCRIPTION_LISTENER);
     this.connection = builder.connection();
     this.sessionHandler = this.connection.createSessionHandler();
     this.nativeReceiver =
         this.createNativeReceiver(
-            this.sessionHandler.session(), this.address, this.linkProperties, this.filters);
+            this.sessionHandler.session(),
+            this.address,
+            this.linkProperties,
+            this.filters,
+            this.subscriptionListener);
     this.initStateFromNativeReceiver(this.nativeReceiver);
     this.metricsCollector = this.connection.metricsCollector();
     this.startReceivingLoop();
@@ -153,8 +161,12 @@ private ClientReceiver createNativeReceiver(
       Session nativeSession,
       String address,
       Map<String, Object> properties,
-      Map<String, Object> filters) {
+      Map<String, Object> filters,
+      SubscriptionListener subscriptionListener) {
     try {
+      filters = new LinkedHashMap<>(filters);
+      StreamOptions streamOptions = AmqpConsumerBuilder.streamOptions(filters);
+      subscriptionListener.preSubscribe(() -> streamOptions);
       ReceiverOptions receiverOptions =
           new ReceiverOptions()
               .deliveryMode(DeliveryMode.AT_LEAST_ONCE)
@@ -221,7 +233,8 @@ void recoverAfterConnectionFailure() {
                     this.sessionHandler.sessionNoCheck(),
                     this.address,
                     this.linkProperties,
-                    this.filters),
+                    this.filters,
+                    this.subscriptionListener),
             e -> {
               boolean shouldRetry =
                   e instanceof AmqpException.AmqpResourceClosedException

src/main/java/com/rabbitmq/client/amqp/impl/AmqpConsumerBuilder.java

@@ -27,6 +27,8 @@
 
 class AmqpConsumerBuilder implements ConsumerBuilder {
 
+  static SubscriptionListener NO_OP_SUBSCRIPTION_LISTENER = ctx -> {};
+
   private final AmqpConnection connection;
   private String queue;
   private Consumer.MessageHandler messageHandler;
@@ -35,6 +37,7 @@ class AmqpConsumerBuilder implements ConsumerBuilder {
   private final Map<String, Object> filters = new LinkedHashMap<>();
   private final Map<String, Object> properties = new LinkedHashMap<>();
   private final StreamOptions streamOptions = new DefaultStreamOptions(this, this.filters);
+  private SubscriptionListener subscriptionListener = NO_OP_SUBSCRIPTION_LISTENER;
 
   AmqpConsumerBuilder(AmqpConnection connection) {
     this.connection = connection;
@@ -79,6 +82,16 @@ public StreamOptions stream() {
     return this.streamOptions;
   }
 
+  @Override
+  public ConsumerBuilder subscriptionListener(SubscriptionListener subscriptionListener) {
+    this.subscriptionListener = subscriptionListener;
+    return this;
+  }
+
+  SubscriptionListener subscriptionListener() {
+    return this.subscriptionListener;
+  }
+
   AmqpConnection connection() {
     return connection;
   }
@@ -186,4 +199,8 @@ private void offsetSpecification(Object value) {
       this.filters.put("rabbitmq:stream-offset-spec", value);
     }
   }
+
+  static StreamOptions streamOptions(Map<String, Object> filters) {
+    return new DefaultStreamOptions(null, filters);
+  }
 }