Merge pull request #757 from rabbitmq/dynamic-batch-pumping

Make dynamic batch pump more aggressively

Author: acogoluegnes

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

@@ -69,22 +69,30 @@ 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);
       }
     }
   }
 
+  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;

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);
+    }
+  }
 }