Use message context in flow strategy callback

Author: acogoluegnes

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

@@ -14,13 +14,12 @@
 package com.rabbitmq.stream;
 
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.function.LongConsumer;
 
 public interface ConsumerFlowStrategy {
 
   int initialCredits();
 
-  LongConsumer start(Context context);
+  MessageProcessedCallback start(Context context);
 
   interface Context {
 
@@ -29,6 +28,12 @@ interface Context {
     long messageCount();
   }
 
+  @FunctionalInterface
+  interface MessageProcessedCallback {
+
+    void processed(MessageHandler.Context messageContext);
+  }
+
   static ConsumerFlowStrategy creditOnChunkArrival() {
     return creditOnChunkArrival(1);
   }
@@ -63,7 +68,7 @@ public int initialCredits() {
     }
 
     @Override
-    public LongConsumer start(Context context) {
+    public MessageProcessedCallback start(Context context) {
       context.credits(1);
       return value -> {};
     }
@@ -85,7 +90,7 @@ public int initialCredits() {
     }
 
     @Override
-    public LongConsumer start(Context context) {
+    public MessageProcessedCallback start(Context context) {
       long l = (long) (context.messageCount() * ratio);
       long limit = Math.max(1, l);
       AtomicLong processedMessages = new AtomicLong(0);

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

@@ -484,14 +484,14 @@ private static final class MessageHandlerContext implements Context {
     private final long timestamp;
     private final long committedOffset;
     private final StreamConsumer consumer;
-    private final LongConsumer processedCallback;
+    private final ConsumerFlowStrategy.MessageProcessedCallback processedCallback;
 
     private MessageHandlerContext(
         long offset,
         long timestamp,
         long committedOffset,
         StreamConsumer consumer,
-        LongConsumer processedCallback) {
+        ConsumerFlowStrategy.MessageProcessedCallback processedCallback) {
       this.offset = offset;
       this.timestamp = timestamp;
       this.committedOffset = committedOffset;
@@ -530,7 +530,7 @@ public Consumer consumer() {
 
     @Override
     public void processed() {
-      this.processedCallback.accept(this.offset);
+      this.processedCallback.processed(this);
     }
   }
 
@@ -567,7 +567,7 @@ private ClientSubscriptionsManager(Broker node, Client.ClientParameters clientPa
           (client, subscriptionId, offset, messageCount, dataSize) -> {
             SubscriptionTracker subscriptionTracker =
                 subscriptionTrackers.get(subscriptionId & 0xFF);
-            LongConsumer processCallback;
+            ConsumerFlowStrategy.MessageProcessedCallback processCallback;
             if (subscriptionTracker != null && subscriptionTracker.consumer.isOpen()) {
               ConsumerFlowStrategy.Context chunkContext =
                   new ConsumerFlowStrategy.Context() {
@@ -616,7 +616,7 @@ public long messageCount() {
                       chunkTimestamp,
                       committedOffset,
                       subscriptionTracker.consumer,
-                      (LongConsumer) chunkContext),
+                      (ConsumerFlowStrategy.MessageProcessedCallback) chunkContext),
                   message);
             } else {
               LOGGER.debug(

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

@@ -71,8 +71,6 @@ public class ConsumersCoordinatorTest {
 
   private static final SubscriptionListener NO_OP_SUBSCRIPTION_LISTENER = subscriptionContext -> {};
   private static final Runnable NO_OP_TRACKING_CLOSING_CALLBACK = () -> {};
-  private int initialCredits = 10;
-  private int additionalCredits = 1;
 
   @Mock StreamEnvironment environment;
   @Mock StreamConsumer consumer;

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

@@ -19,7 +19,6 @@
 
 import com.rabbitmq.stream.ConsumerFlowStrategy;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.LongConsumer;
 import org.junit.jupiter.api.Test;
 
 public class MessageCountConsumerFlowStrategyTest {
@@ -30,27 +29,27 @@ public class MessageCountConsumerFlowStrategyTest {
   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));
+    ConsumerFlowStrategy.MessageProcessedCallback callback = strategy.start(context(messageCount));
+    range(0, messageCount / 2 - 1).forEach(ignored -> callback.processed(null));
     assertThat(requestedCredits).hasValue(0);
-    processedCallback.accept(42);
+    callback.processed(null);
     assertThat(requestedCredits).hasValue(1);
-    processedCallback.accept(42);
+    callback.processed(null);
     assertThat(requestedCredits).hasValue(1);
-    range(0, messageCount).forEach(ignored -> processedCallback.accept(42));
+    range(0, messageCount).forEach(ignored -> callback.processed(null));
     assertThat(requestedCredits).hasValue(1);
   }
 
   @Test
   void smallChunksAndSmallRatiosShouldCredit() {
     ConsumerFlowStrategy strategy = build(0.5);
-    LongConsumer processedCallback = strategy.start(context(1));
-    processedCallback.accept(42);
+    ConsumerFlowStrategy.MessageProcessedCallback callback = strategy.start(context(1));
+    callback.processed(null);
     assertThat(requestedCredits).hasValue(1);
 
     strategy = build(0.05);
-    processedCallback = strategy.start(context(15));
-    processedCallback.accept(42);
+    callback = strategy.start(context(15));
+    callback.processed(null);
     assertThat(requestedCredits).hasValue(1);
   }