Introduce replay argument in dynamic batch support class

Author: acogoluegnes

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

@@ -22,7 +22,7 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
-import java.util.function.Predicate;
+import java.util.function.BiPredicate;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -33,11 +33,11 @@ final class DynamicBatch<T> implements AutoCloseable {
   private static final int MAX_BATCH_SIZE = 8192;
 
   private final BlockingQueue<T> requests = new LinkedBlockingQueue<>();
-  private final Predicate<List<T>> consumer;
+  private final BiPredicate<List<T>, Boolean> consumer;
   private final int configuredBatchSize;
   private final Thread thread;
 
-  DynamicBatch(Predicate<List<T>> consumer, int batchSize) {
+  DynamicBatch(BiPredicate<List<T>, Boolean> consumer, int batchSize) {
     this.consumer = consumer;
     this.configuredBatchSize = min(max(batchSize, MIN_BATCH_SIZE), MAX_BATCH_SIZE);
     this.thread = ConcurrencyUtils.defaultThreadFactory().newThread(this::loop);
@@ -53,8 +53,10 @@ void add(T item) {
   }
 
   private void loop() {
-    int batchSize = this.configuredBatchSize;
-    List<T> batch = new ArrayList<>(batchSize);
+    State<T> state = new State<>();
+    state.batchSize = this.configuredBatchSize;
+    state.items = new ArrayList<>(state.batchSize);
+    state.retry = false;
     Thread currentThread = Thread.currentThread();
     T item;
     while (!currentThread.isInterrupted()) {
@@ -65,44 +67,50 @@ private void loop() {
         return;
       }
       if (item != null) {
-        batch.add(item);
-        if (batch.size() >= batchSize) {
-          if (this.completeBatch(batch)) {
-            batchSize = min(batchSize * 2, MAX_BATCH_SIZE);
-            batch = new ArrayList<>(batchSize);
-          }
+        state.items.add(item);
+        if (state.items.size() >= state.batchSize) {
+          this.maybeCompleteBatch(state, true);
         } else {
           item = this.requests.poll();
           if (item == null) {
-            if (this.completeBatch(batch)) {
-              batchSize = max(batchSize / 2, MIN_BATCH_SIZE);
-              batch = new ArrayList<>(batchSize);
-            }
+            this.maybeCompleteBatch(state, false);
           } else {
-            batch.add(item);
-            if (batch.size() >= batchSize) {
-              if (this.completeBatch(batch)) {
-                batchSize = min(batchSize * 2, MAX_BATCH_SIZE);
-                batch = new ArrayList<>(batchSize);
-              }
+            state.items.add(item);
+            if (state.items.size() >= state.batchSize) {
+              this.maybeCompleteBatch(state, true);
             }
           }
         }
       } else {
-        if (this.completeBatch(batch)) {
-          batchSize = min(batchSize * 2, MAX_BATCH_SIZE);
-          batch = new ArrayList<>(batchSize);
-        }
+        this.maybeCompleteBatch(state, false);
       }
     }
   }
 
-  private boolean completeBatch(List<T> items) {
+  private static final class State<T> {
+
+    int batchSize;
+    List<T> items;
+    boolean retry;
+  }
+
+  private void maybeCompleteBatch(State<T> state, boolean increaseIfCompleted) {
     try {
-      return this.consumer.test(items);
+      boolean completed = this.consumer.test(state.items, state.retry);
+      if (completed) {
+        if (increaseIfCompleted) {
+          state.batchSize = min(state.batchSize * 2, MAX_BATCH_SIZE);
+        } else {
+          state.batchSize = max(state.batchSize / 2, MIN_BATCH_SIZE);
+        }
+        state.items = new ArrayList<>(state.batchSize);
+        state.retry = false;
+      } else {
+        state.retry = true;
+      }
     } catch (Exception e) {
       LOGGER.warn("Error during dynamic batch completion: {}", e.getMessage());
-      return false;
+      state.retry = true;
     }
   }
 

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

@@ -62,67 +62,61 @@ final class DynamicBatchMessageAccumulator implements MessageAccumulator {
     if (subEntrySize <= 1) {
       this.dynamicBatch =
           new DynamicBatch<>(
-              items -> {
-                // TODO add a "replay" flag to DynamicBatch to avoid checking the producer status
-                // the status check helps to avoid collecting the observation another time
-                if (producer.canSend()) {
+              (items, replay) -> {
+                if (!replay) {
                   items.forEach(
                       i -> {
                         AccumulatedEntity entity = (AccumulatedEntity) i;
                         this.observationCollector.published(
                             entity.observationContext(), entity.confirmationCallback().message());
                       });
-                  return this.publish(items);
-                } else {
-                  return false;
                 }
+                return this.publish(items);
               },
               batchSize);
     } else {
       byte compressionCode =
           compressionCodec == null ? Compression.NONE.code() : compressionCodec.code();
       this.dynamicBatch =
           new DynamicBatch<>(
-              items -> {
-                if (this.producer.canSend()) {
-                  List<Object> subBatches = new ArrayList<>();
-                  int count = 0;
-                  ProducerUtils.Batch batch =
-                      this.helper.batch(
-                          byteBufAllocator, compressionCode, compressionCodec, subEntrySize);
-                  AccumulatedEntity lastMessageInBatch = null;
-                  for (Object msg : items) {
-                    AccumulatedEntity message = (AccumulatedEntity) msg;
+              (items, replay) -> {
+                List<Object> subBatches = new ArrayList<>();
+                int count = 0;
+                ProducerUtils.Batch batch =
+                    this.helper.batch(
+                        byteBufAllocator, compressionCode, compressionCodec, subEntrySize);
+                AccumulatedEntity lastMessageInBatch = null;
+                for (Object msg : items) {
+                  AccumulatedEntity message = (AccumulatedEntity) msg;
+                  if (!replay) {
                     this.observationCollector.published(
                         message.observationContext(), message.confirmationCallback().message());
-                    lastMessageInBatch = message;
-                    batch.add(
-                        (Codec.EncodedMessage) message.encodedEntity(),
-                        message.confirmationCallback());
-                    count++;
-                    if (count == subEntrySize) {
-                      batch.time = lastMessageInBatch.time();
-                      batch.publishingId = lastMessageInBatch.publishingId();
-                      batch.encodedMessageBatch.close();
-                      subBatches.add(batch);
-                      lastMessageInBatch = null;
-                      batch =
-                          this.helper.batch(
-                              byteBufAllocator, compressionCode, compressionCodec, subEntrySize);
-                      count = 0;
-                    }
                   }
-
-                  if (!batch.isEmpty() && count < subEntrySize) {
+                  lastMessageInBatch = message;
+                  batch.add(
+                      (Codec.EncodedMessage) message.encodedEntity(),
+                      message.confirmationCallback());
+                  count++;
+                  if (count == subEntrySize) {
                     batch.time = lastMessageInBatch.time();
                     batch.publishingId = lastMessageInBatch.publishingId();
                     batch.encodedMessageBatch.close();
                     subBatches.add(batch);
+                    lastMessageInBatch = null;
+                    batch =
+                        this.helper.batch(
+                            byteBufAllocator, compressionCode, compressionCodec, subEntrySize);
+                    count = 0;
                   }
-                  return this.publish(subBatches);
-                } else {
-                  return false;
                 }
+
+                if (!batch.isEmpty() && count < subEntrySize) {
+                  batch.time = lastMessageInBatch.time();
+                  batch.publishingId = lastMessageInBatch.publishingId();
+                  batch.encodedMessageBatch.close();
+                  subBatches.add(batch);
+                }
+                return this.publish(subBatches);
               },
               batchSize * subEntrySize);
     }

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

@@ -40,28 +40,32 @@ void test() {
     Random random = new Random();
     DynamicBatch<String> batch =
         new DynamicBatch<>(
-            items -> {
+            (items, replay) -> {
               batchSizeMetrics.update(items.size());
-              sync.down(items.size());
               try {
                 Thread.sleep(random.nextInt(10) + 1);
               } catch (InterruptedException e) {
                 throw new RuntimeException(e);
               }
+              sync.down(items.size());
               return true;
             },
             100);
-    RateLimiter rateLimiter = RateLimiter.create(3000);
-    long start = System.nanoTime();
-    IntStream.range(0, itemCount)
-        .forEach(
-            i -> {
-              rateLimiter.acquire();
-              batch.add(String.valueOf(i));
-            });
-    Assertions.assertThat(sync).completes();
-    long end = System.nanoTime();
-    //    System.out.println("Done in " + Duration.ofNanos(end - start));
-    //    reporter.report();
+    try {
+      RateLimiter rateLimiter = RateLimiter.create(3000);
+      long start = System.nanoTime();
+      IntStream.range(0, itemCount)
+          .forEach(
+              i -> {
+                rateLimiter.acquire();
+                batch.add(String.valueOf(i));
+              });
+      Assertions.assertThat(sync).completes();
+      long end = System.nanoTime();
+      //    System.out.println("Done in " + Duration.ofNanos(end - start));
+      //    reporter.report();
+    } finally {
+      batch.close();
+    }
   }
 }