Consumer flow strategy (take two)

pom.xml

@@ -383,6 +383,9 @@
           <excludes>
             <exclude>**/*TestSuite.java</exclude>
           </excludes>
+          <systemProperties>
+            <rabbitmqctl.bin>DOCKER:rabbitmq</rabbitmqctl.bin>
+          </systemProperties>
         </configuration>
       </plugin>
 
@@ -875,6 +878,32 @@
 
     </profile>
 
+    <profile>
+      <!-- this avoids a compiler warning on Java 9+ -->
+      <!-- the compiler setting is not available on Java 8 -->
+      <id>use-release-compiler-argument-on-java-9-or-more</id>
+      <activation>
+        <jdk>[9,)</jdk>
+      </activation>
+      <build>
+        <plugins>
+          <plugin>
+            <artifactId>maven-compiler-plugin</artifactId>
+            <version>${maven.compiler.plugin.version}</version>
+            <configuration>
+              <source>1.8</source>
+              <target>1.8</target>
+              <release>8</release>
+              <compilerArgs>
+                <arg>-Xlint:deprecation</arg>
+                <arg>-Xlint:unchecked</arg>
+              </compilerArgs>
+            </configuration>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
+
   </profiles>
 
 </project>

src/docs/asciidoc/api.adoc

@@ -851,6 +851,10 @@ Useful when using an external store for offset tracking.
 |Number of credits when the subscription is created.
 Increase for higher throughput at the expense of memory usage.
 |1
+
+|`flow#strategy`
+|The `ConsumerFlowStrategy` to use.
+|`ConsumerFlowStrategy#creditOnChunkArrival(1)`
 |===
 
 [NOTE]
@@ -1099,6 +1103,48 @@ When a glitch happens and triggers the re-subscription, the server-side stored o
 Using this server-side stored offset can lead to duplicates, whereas using the in-memory, application-specific offset tracking variable is more accurate.
 A custom `SubscriptionListener` lets the application developer uses what's best for the application if the computed value is not optimal.
 
+===== Flow Control
+
+This section covers how a consumer can tell the broker when to send more messages.
+
+By default, the broker keeps sending messages as long as messages are processed and the `MessageHandler#handle(Context, Message)` method returns.
+This strategy works fine if message processing is fast enough.
+If message processing takes longer, one can be tempted to process messages in parallel with an `ExecutorService`.
+This will make the `handle` method return immediately and the broker will keep sending messages, potentially overflowing the consumer.
+
+What we miss in the parallel processing case is a way to tell the library we are done processing a message and that we are ready at some point to handle more messages.
+This is the goal of the `MessageHandler.Context#processed()` method.
+
+This method is by default a no-op because the default flow control strategy keeps asking for more messages as soon as message processing is done.
+This method gets some real behavior to control the flow of messages when an appropriate `ConsumerFlowStrategy` is set `ConsumerBuilder#flow()`.
+The following code snippet shows how to set a handy consumer flow strategy:
+
+.Setting a consumer flow control strategy
+[source,java,indent=0]
+--------
+include::{test-examples}/ConsumerUsage.java[tag=flow-control]
+--------
+<1> Set the flow control strategy
+<2> Make sure to call `Context#processed()`
+
+In the example we set up the `creditWhenHalfMessagesProcessed` strategy which asks for more messages once half of the current messages have been marked as processed.
+The broker does not send messages one by one, it sends <<chunk-definition,chunks>> of messages.
+A chunk of messages can contain 1 to several thousands of messages.
+So with the strategy set above, once `processed()` has been called for half of the messages of the current chunk, the library will ask the broker for another one (it will provide a _credit_ for the subscription).
+By doing this, the next chunk should arrive by the time we are done with the other half of the current chunk.
+This way the consumer is neither overwhelmed nor idle.
+
+The `ConsumerFlowStrategy` interface provides some static helpers to configure the appropriate strategy.
+
+Additional notes on consumer flow control:
+
+* Make sure to **call the `processed()` method** once you set up a `ConsumerFlowStrategy`.
+The method is a no-op by default, but it is essential to call it with count-based strategies like `creditWhenHalfMessagesProcessed` or `creditOnProcessedMessageCount`.
+No calling it will stop the dispatching of messages.
+* Make sure to call `processed()` only once.
+Whether the method is idempotent depends on the flow strategy implementation.
+Apart from the default one, the implementations the library provides does not make `processed()` idempotent.
+
 [[single-active-consumer]]
 ===== Single Active Consumer
 

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

@@ -166,6 +166,8 @@ public interface ConsumerBuilder {
    *
    * @return the flow configuration
    * @since 0.11.0
+   * @see ConsumerFlowStrategy#creditOnChunkArrival()
+   * @see MessageHandler.Context#processed()
    */
   FlowConfiguration flow();
 
@@ -231,7 +233,11 @@ interface AutoTrackingStrategy {
   /**
    * Message flow configuration.
    *
+   * <p>The default configuration uses {@link ConsumerFlowStrategy#creditOnChunkArrival()}.
+   *
    * @since 0.11.0
+   * @see ConsumerFlowStrategy#creditOnChunkArrival()
+   * @see MessageHandler.Context#processed()
    */
   interface FlowConfiguration {
 
@@ -240,11 +246,29 @@ interface FlowConfiguration {
      *
      * <p>Default is 1.
      *
+     * <p>This calls uses {@link ConsumerFlowStrategy#creditOnChunkArrival(int)}.
+     *
      * @param initialCredits the number of initial credits
      * @return this configuration instance
+     * @see ConsumerFlowStrategy#creditOnChunkArrival(int)
      */
     FlowConfiguration initialCredits(int initialCredits);
 
+    /**
+     * Flow strategy to use
+     *
+     * @param strategy the strategy to use
+     * @return this configuration instance
+     * @since 0.12.0
+     * @see ConsumerFlowStrategy
+     * @see ConsumerFlowStrategy#creditOnChunkArrival()
+     * @see ConsumerFlowStrategy#creditOnChunkArrival(int)
+     * @see ConsumerFlowStrategy#creditWhenHalfMessagesProcessed()
+     * @see ConsumerFlowStrategy#creditWhenHalfMessagesProcessed(int)
+     * @see ConsumerFlowStrategy#creditOnProcessedMessageCount(int, double)
+     */
+    FlowConfiguration strategy(ConsumerFlowStrategy strategy);
+
     /**
      * Go back to the builder.
      *

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

@@ -0,0 +1,227 @@
+// Copyright (c) 2023 VMware, Inc. or its affiliates.  All rights reserved.
+//
+// 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;
+
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Contract to determine when a subscription provides credits to get more messages.
+ *
+ * <p>The broker delivers "chunks" of messages to consumers. A chunk can contain from 1 to several
+ * thousands of messages. The broker send chunks as long as the subscription has <em>credits</em>. A
+ * client connection can provide credits for a given subscription and the broker will send the
+ * corresponding number of chunks (1 credit = 1 chunk).
+ *
+ * <p>This credit mechanism avoids overwhelming a consumer with messages. A consumer does not want
+ * to provide a credit only when it is done with messages of a chunk, because it will be idle
+ * between its credit request and the arrival of the next chunk. The idea is to keep consumers busy
+ * as much as possible, without accumulating an in-memory backlog on the client side. There is no
+ * ideal solution, it depends on the use cases and several parameters (processing time, network,
+ * etc).
+ *
+ * <p>This is an experimental API, subject to change.
+ *
+ * @since 0.12.0
+ * @see MessageHandler.Context#processed()
+ * @see ConsumerBuilder#flow()
+ */
+public interface ConsumerFlowStrategy {
+
+  /**
+   * The initial number of credits for a subscription.
+   *
+   * <p>It must be greater than 0. Values are usually between 1 and 10.
+   *
+   * @return initial number of credits
+   */
+  int initialCredits();
+
+  /**
+   * Return the behavior for {@link MessageHandler.Context#processed()} calls.
+   *
+   * <p>This method is called for each chunk of messages. Implementations return a callback that
+   * will be called when applications consider a message dealt with and call {@link
+   * MessageHandler.Context#processed()}. The callback can count messages and provide credits
+   * accordingly.
+   *
+   * @param context chunk context
+   * @return the message processed callback
+   */
+  MessageProcessedCallback start(Context context);
+
+  /** Chunk context. */
+  interface Context {
+
+    /**
+     * Provide credits for the subscription.
+     *
+     * <p>{@link ConsumerFlowStrategy} implementation should always provide 1 credit a given chunk.
+     *
+     * @param credits the number of credits provided, usually 1
+     */
+    void credits(int credits);
+
+    /**
+     * The number of messages in the chunk.
+     *
+     * @return number of messages in the chunk
+     */
+    long messageCount();
+  }
+
+  /** Behavior for {@link MessageHandler.Context#processed()} calls. */
+  @FunctionalInterface
+  interface MessageProcessedCallback {
+
+    /**
+     * Method called when {@link MessageHandler.Context#processed()} is called.
+     *
+     * <p>There is one instance of this class for a given chunk and it is called for the <code>
+     * processed()</code> calls of the message of this chunk.
+     *
+     * <p>Implementations can count messages and call {@link Context#credits(int)} when appropriate.
+     *
+     * <p>Note calls to {@link MessageHandler.Context#processed()} are not idempotent: an
+     * application can call the method several times for the same message and implementations must
+     * deal with these multiple calls if they impact their logic.
+     *
+     * @param messageContext context of the message
+     */
+    void processed(MessageHandler.Context messageContext);
+  }
+
+  /**
+   * Strategy that provides 1 initial credit and a credit on each new chunk.
+   *
+   * <p>Calls to {@link MessageHandler.Context#processed()} are ignored.
+   *
+   * @return flow strategy
+   */
+  static ConsumerFlowStrategy creditOnChunkArrival() {
+    return creditOnChunkArrival(1);
+  }
+
+  /**
+   * Strategy that provides the specified number of initial credits and a credit on each new chunk.
+   *
+   * <p>Calls to {@link MessageHandler.Context#processed()} are ignored.
+   *
+   * @param initialCredits number of initial credits
+   * @return flow strategy
+   */
+  static ConsumerFlowStrategy creditOnChunkArrival(int initialCredits) {
+    return new CreditOnChunkArrivalConsumerFlowStrategy(initialCredits);
+  }
+
+  /**
+   * Strategy that provides 1 initial credit and a credit when half of the chunk messages are
+   * processed.
+   *
+   * <p>Make sure to call {@link MessageHandler.Context#processed()} on every message when using
+   * this strategy, otherwise the broker may stop sending messages to the consumer.
+   *
+   * @return flow strategy
+   */
+  static ConsumerFlowStrategy creditWhenHalfMessagesProcessed() {
+    return creditOnProcessedMessageCount(1, 0.5);
+  }
+
+  /**
+   * Strategy that provides the specified number of initial credits and a credit when half of the
+   * chunk messages are processed.
+   *
+   * <p>Make sure to call {@link MessageHandler.Context#processed()} on every message when using
+   * this strategy, otherwise the broker may stop sending messages to the consumer.
+   *
+   * @param initialCredits number of initial credits
+   * @return flow strategy
+   */
+  static ConsumerFlowStrategy creditWhenHalfMessagesProcessed(int initialCredits) {
+    return creditOnProcessedMessageCount(initialCredits, 0.5);
+  }
+
+  /**
+   * Strategy that provides the specified number of initial credits and a credit when the specified
+   * ratio of the chunk messages are processed.
+   *
+   * <p>Make sure to call {@link MessageHandler.Context#processed()} on every message when using
+   * this strategy, otherwise the broker may stop sending messages to the consumer.
+   *
+   * @param initialCredits number of initial credits
+   * @return flow strategy
+   */
+  static ConsumerFlowStrategy creditOnProcessedMessageCount(int initialCredits, double ratio) {
+    return new MessageCountConsumerFlowStrategy(initialCredits, ratio);
+  }
+
+  /**
+   * Strategy that provides the specified number of initial credits and a credit on each new chunk.
+   *
+   * <p>Calls to {@link MessageHandler.Context#processed()} are ignored.
+   */
+  class CreditOnChunkArrivalConsumerFlowStrategy implements ConsumerFlowStrategy {
+
+    private final int initialCredits;
+
+    private CreditOnChunkArrivalConsumerFlowStrategy(int initialCredits) {
+      this.initialCredits = initialCredits;
+    }
+
+    @Override
+    public int initialCredits() {
+      return this.initialCredits;
+    }
+
+    @Override
+    public MessageProcessedCallback start(Context context) {
+      context.credits(1);
+      return value -> {};
+    }
+  }
+
+  /**
+   * Strategy that provides the specified number of initial credits and a credit when the specified
+   * ratio of the chunk messages are processed.
+   *
+   * <p>Make sure to call {@link MessageHandler.Context#processed()} on every message when using
+   * this strategy, otherwise the broker may stop sending messages to the consumer.
+   */
+  class MessageCountConsumerFlowStrategy implements ConsumerFlowStrategy {
+
+    private final int initialCredits;
+    private final double ratio;
+
+    private MessageCountConsumerFlowStrategy(int initialCredits, double ratio) {
+      this.initialCredits = initialCredits;
+      this.ratio = ratio;
+    }
+
+    @Override
+    public int initialCredits() {
+      return this.initialCredits;
+    }
+
+    @Override
+    public MessageProcessedCallback start(Context context) {
+      long l = (long) (context.messageCount() * ratio);
+      long limit = Math.max(1, l);
+      AtomicLong processedMessages = new AtomicLong(0);
+      return messageOffset -> {
+        if (processedMessages.incrementAndGet() == limit) {
+          context.credits(1);
+        }
+      };
+    }
+  }
+}