Add consumer flow strategy creation helpers

Author: acogoluegnes

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

@@ -29,11 +29,31 @@ interface Context {
     long messageCount();
   }
 
-  class DefaultConsumerFlowStrategy implements ConsumerFlowStrategy {
+  static ConsumerFlowStrategy creditOnChunkArrival() {
+    return creditOnChunkArrival(1);
+  }
+
+  static ConsumerFlowStrategy creditOnChunkArrival(int initialCredits) {
+    return new CreditOnChunkArrivalConsumerFlowStrategy(initialCredits);
+  }
+
+  static ConsumerFlowStrategy creditWhenHalfMessagesProcessed() {
+    return creditOnProcessedMessageCount(1, 0.5);
+  }
+
+  static ConsumerFlowStrategy creditWhenHalfMessagesProcessed(int initialCredits) {
+    return creditOnProcessedMessageCount(initialCredits, 0.5);
+  }
+
+  static ConsumerFlowStrategy creditOnProcessedMessageCount(int initialCredits, double ratio) {
+    return new MessageCountConsumerFlowStrategy(initialCredits, ratio);
+  }
+
+  class CreditOnChunkArrivalConsumerFlowStrategy implements ConsumerFlowStrategy {
 
     private final int initialCredits;
 
-    public DefaultConsumerFlowStrategy(int initialCredits) {
+    private CreditOnChunkArrivalConsumerFlowStrategy(int initialCredits) {
       this.initialCredits = initialCredits;
     }
 
@@ -52,9 +72,11 @@ public LongConsumer start(Context context) {
   class MessageCountConsumerFlowStrategy implements ConsumerFlowStrategy {
 
     private final int initialCredits;
+    private final double ratio;
 
-    public MessageCountConsumerFlowStrategy(int initialCredits) {
+    private MessageCountConsumerFlowStrategy(int initialCredits, double ratio) {
       this.initialCredits = initialCredits;
+      this.ratio = ratio;
     }
 
     @Override
@@ -64,8 +86,9 @@ public int initialCredits() {
 
     @Override
     public LongConsumer start(Context context) {
+      long l = (long) (context.messageCount() * ratio);
+      long limit = Math.max(1, l);
       AtomicLong processedMessages = new AtomicLong(0);
-      long limit = context.messageCount() == 1 ? 1 : context.messageCount() / 2;
       return messageOffset -> {
         if (processedMessages.incrementAndGet() == limit) {
           context.credits(1);

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

@@ -425,14 +425,14 @@ private DefaultFlowConfiguration(ConsumerBuilder consumerBuilder) {
       this.consumerBuilder = consumerBuilder;
     }
 
-    private ConsumerFlowStrategy strategy = new ConsumerFlowStrategy.DefaultConsumerFlowStrategy(1);
+    private ConsumerFlowStrategy strategy = ConsumerFlowStrategy.creditOnChunkArrival(1);
 
     @Override
     public FlowConfiguration initialCredits(int initialCredits) {
       if (initialCredits <= 0) {
         throw new IllegalArgumentException("Credits must be positive");
       }
-      this.strategy = new ConsumerFlowStrategy.DefaultConsumerFlowStrategy(initialCredits);
+      this.strategy = ConsumerFlowStrategy.creditOnChunkArrival(initialCredits);
       return this;
     }
 

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

@@ -1678,6 +1678,6 @@ private static Response responseOk() {
   }
 
   private static ConsumerFlowStrategy flowStrategy() {
-    return new ConsumerFlowStrategy.DefaultConsumerFlowStrategy(10);
+    return ConsumerFlowStrategy.creditOnChunkArrival(10);
   }
 }

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

@@ -13,6 +13,7 @@
 // [email protected].
 package com.rabbitmq.stream.impl;
 
+import static com.rabbitmq.stream.ConsumerFlowStrategy.creditOnProcessedMessageCount;
 import static java.util.stream.LongStream.range;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -26,8 +27,8 @@ public class MessageCountConsumerFlowStrategyTest {
   AtomicInteger requestedCredits = new AtomicInteger();
 
   @Test
-  void test() {
-    ConsumerFlowStrategy strategy = new ConsumerFlowStrategy.MessageCountConsumerFlowStrategy(1);
+  void shouldCreditOnceLimitIsReached() {
+    ConsumerFlowStrategy strategy = build(0.5);
     long messageCount = 1000;
     LongConsumer processedCallback = strategy.start(context(messageCount));
     range(0, messageCount / 2 - 1).forEach(ignored -> processedCallback.accept(42));
@@ -40,6 +41,23 @@ void test() {
     assertThat(requestedCredits).hasValue(1);
   }
 
+  @Test
+  void smallChunksAndSmallRatiosShouldCredit() {
+    ConsumerFlowStrategy strategy = build(0.5);
+    LongConsumer processedCallback = strategy.start(context(1));
+    processedCallback.accept(42);
+    assertThat(requestedCredits).hasValue(1);
+
+    strategy = build(0.05);
+    processedCallback = strategy.start(context(15));
+    processedCallback.accept(42);
+    assertThat(requestedCredits).hasValue(1);
+  }
+
+  ConsumerFlowStrategy build(double ratio) {
+    return creditOnProcessedMessageCount(1, ratio);
+  }
+
   ConsumerFlowStrategy.Context context(long messageCount) {
     requestedCredits.set(0);
     return new ConsumerFlowStrategy.Context() {

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

@@ -13,6 +13,7 @@
 // [email protected].
 package com.rabbitmq.stream.impl;
 
+import static com.rabbitmq.stream.ConsumerFlowStrategy.creditWhenHalfMessagesProcessed;
 import static com.rabbitmq.stream.impl.TestUtils.*;
 import static java.lang.String.format;
 import static java.util.Collections.synchronizedList;
@@ -199,7 +200,7 @@ void consumeWithAsyncConsumerFlowControl() throws Exception {
         environment.consumerBuilder().stream(stream)
             .offset(OffsetSpecification.first())
             .flow()
-            .strategy(new ConsumerFlowStrategy.MessageCountConsumerFlowStrategy(1))
+            .strategy(creditWhenHalfMessagesProcessed())
             .builder();
 
     List<MessageHandler.Context> messageContexts = synchronizedList(new ArrayList<>());