From 8a198eeb45592726736e21580330ece1bfeb5586 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arnaud=20Cogolu=C3=A8gnes?=
 <514737+acogoluegnes@users.noreply.github.com>
Date: Mon, 12 May 2025 14:54:47 +0200
Subject: [PATCH] Make dynamic batch pump more aggressively

Low value for maxUnconfirmedMessages combined with unfortunate timing
can make the dynamic batch flush only on timeout.

This commit makes the dynamic batch class "pump" for new items
(messages) more aggressively, which mitigates the problem.

References #750
---
 .../rabbitmq/stream/impl/DynamicBatch.java    | 26 +++++++++-----
 .../stream/impl/DynamicBatchTest.java         | 36 +++++++++++++++++++
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/src/main/java/com/rabbitmq/stream/impl/DynamicBatch.java b/src/main/java/com/rabbitmq/stream/impl/DynamicBatch.java
index 7e2d1a7369..c6038dfed6 100644
--- a/src/main/java/com/rabbitmq/stream/impl/DynamicBatch.java
+++ b/src/main/java/com/rabbitmq/stream/impl/DynamicBatch.java
@@ -69,15 +69,7 @@ private void loop() {
         if (state.items.size() >= state.batchSize) {
           this.maybeCompleteBatch(state, true);
         } else {
-          item = this.requests.poll();
-          if (item == null) {
-            this.maybeCompleteBatch(state, false);
-          } else {
-            state.items.add(item);
-            if (state.items.size() >= state.batchSize) {
-              this.maybeCompleteBatch(state, true);
-            }
-          }
+          pump(state, 2);
         }
       } else {
         this.maybeCompleteBatch(state, false);
@@ -85,6 +77,22 @@ private void loop() {
     }
   }
 
+  private void pump(State<T> state, int pumpCount) {
+    if (pumpCount <= 0) {
+      return;
+    }
+    T item = this.requests.poll();
+    if (item == null) {
+      this.maybeCompleteBatch(state, false);
+    } else {
+      state.items.add(item);
+      if (state.items.size() >= state.batchSize) {
+        this.maybeCompleteBatch(state, true);
+      }
+      this.pump(state, pumpCount - 1);
+    }
+  }
+
   private static final class State<T> {
 
     int batchSize;
diff --git a/src/test/java/com/rabbitmq/stream/impl/DynamicBatchTest.java b/src/test/java/com/rabbitmq/stream/impl/DynamicBatchTest.java
index dca9810762..50698320f8 100644
--- a/src/test/java/com/rabbitmq/stream/impl/DynamicBatchTest.java
+++ b/src/test/java/com/rabbitmq/stream/impl/DynamicBatchTest.java
@@ -23,8 +23,11 @@
 import com.rabbitmq.stream.impl.TestUtils.Sync;
 import java.util.Locale;
 import java.util.Random;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.IntStream;
 import org.junit.jupiter.api.Test;
 
@@ -118,4 +121,37 @@ void failedProcessingIsReplayed() throws Exception {
       waitAtMost(() -> collected.get() == itemCount);
     }
   }
+
+  @Test
+  void lowThrottlingValueShouldStillHighPublishingRate() throws Exception {
+    int batchSize = 10;
+    Semaphore semaphore = new Semaphore(batchSize);
+    DynamicBatch.BatchConsumer<Long> action =
+        items -> {
+          semaphore.release(items.size());
+          return true;
+        };
+
+    try (DynamicBatch<Long> batch = new DynamicBatch<>(action, batchSize)) {
+      MetricRegistry metrics = new MetricRegistry();
+      Meter rate = metrics.meter("publishing-rate");
+      AtomicBoolean keepGoing = new AtomicBoolean(true);
+      AtomicLong sequence = new AtomicLong();
+      new Thread(
+              () -> {
+                while (keepGoing.get() && !Thread.interrupted()) {
+                  long id = sequence.getAndIncrement();
+                  if (semaphore.tryAcquire()) {
+                    batch.add(id);
+                    rate.mark();
+                  }
+                }
+              })
+          .start();
+      long start = System.nanoTime();
+      waitAtMost(
+          () ->
+              System.nanoTime() - start > TimeUnit.SECONDS.toNanos(1) && rate.getMeanRate() > 1000);
+    }
+  }
 }