Merge pull request #906 from rabbitmq/improve-dynamic-batch

Improve producer adaptative batching

Author: acogoluegnes

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

@@ -30,39 +30,54 @@ final class DynamicBatch<T> implements AutoCloseable {
   private static final Logger LOGGER = LoggerFactory.getLogger(DynamicBatch.class);
   private static final int MIN_BATCH_SIZE = 16;
 
-  private final BlockingQueue<T> requests = new LinkedBlockingQueue<>();
+  // Additive Increase, Multiplicative Decrease (AIMD) parameters
+  private static final double MULTIPLICATIVE_DECREASE_FACTOR = 0.5;
+  private static final int ADDITIVE_INCREASE_STEP_DIVISOR = 10; // configuredBatchSize / 10
+
+  private static final long RETRY_DELAY_MS = 50;
+
+  private final BlockingQueue<T> requests;
   private final BatchConsumer<T> consumer;
   private final int configuredBatchSize, minBatchSize, maxBatchSize;
+  private final int additiveIncreaseStep;
   private final Thread thread;
+  private volatile boolean running = true;
 
   DynamicBatch(BatchConsumer<T> consumer, int batchSize, int maxUnconfirmed, String id) {
     this.consumer = consumer;
+    this.requests = new LinkedBlockingQueue<>(max(1, maxUnconfirmed));
     if (batchSize < maxUnconfirmed) {
       this.minBatchSize = min(MIN_BATCH_SIZE, batchSize / 2);
     } else {
-      this.minBatchSize = min(1, maxUnconfirmed / 2);
+      this.minBatchSize = max(1, maxUnconfirmed / 2);
     }
     this.configuredBatchSize = batchSize;
     this.maxBatchSize = batchSize * 2;
+
+    // Calculate additive increase step: 10% of configured size, minimum 1
+    this.additiveIncreaseStep = max(1, batchSize / ADDITIVE_INCREASE_STEP_DIVISOR);
+
     this.thread = ThreadUtils.newInternalThread(id, this::loop);
+    this.thread.setDaemon(true);
     this.thread.start();
   }
 
   void add(T item) {
     try {
       requests.put(item);
     } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
       throw new RuntimeException(e);
     }
   }
 
   private void loop() {
-    State<T> state = new State<>();
+    // Initial allocation based on maxBatchSize to avoid resizing if it grows
+    State<T> state = new State<>(new ArrayList<>(this.maxBatchSize));
     state.batchSize = this.configuredBatchSize;
-    state.items = new ArrayList<>(state.batchSize);
     Thread currentThread = Thread.currentThread();
     T item;
-    while (!currentThread.isInterrupted()) {
+    while (running && !currentThread.isInterrupted()) {
       try {
         item = this.requests.poll(100, TimeUnit.MILLISECONDS);
       } catch (InterruptedException e) {
@@ -71,59 +86,87 @@ private void loop() {
       }
       if (item != null) {
         state.items.add(item);
-        if (state.items.size() >= state.batchSize) {
-          this.maybeCompleteBatch(state, true);
-        } else {
-          pump(state, 2);
+        int remaining = state.batchSize - state.items.size();
+        if (remaining > 0) {
+          this.requests.drainTo(state.items, remaining);
         }
+        this.maybeCompleteBatch(state, state.items.size() >= state.batchSize);
       } 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;
-    List<T> items;
+    final ArrayList<T> items;
+
+    private State(ArrayList<T> items) {
+      this.items = items;
+    }
   }
 
   private void maybeCompleteBatch(State<T> state, boolean increaseIfCompleted) {
+    if (state.items.isEmpty()) {
+      return;
+    }
+
     try {
       boolean completed = this.consumer.process(state.items);
       if (completed) {
         if (increaseIfCompleted) {
-          state.batchSize = min(state.batchSize * 2, this.maxBatchSize);
+          // AIMD: Additive Increase
+          // Grow slowly and linearly when batch is full
+          state.batchSize = min(state.batchSize + this.additiveIncreaseStep, this.maxBatchSize);
         } else {
-          state.batchSize = max(state.batchSize / 2, this.minBatchSize);
+          // AIMD: Multiplicative Decrease
+          // React quickly to low utilization by cutting back
+          state.batchSize =
+              max((int) (state.batchSize * MULTIPLICATIVE_DECREASE_FACTOR), this.minBatchSize);
         }
-        state.items = new ArrayList<>(state.batchSize);
+        state.items.clear();
+        return;
       }
     } catch (Exception e) {
-      //      e.printStackTrace();
-      LOGGER.warn("Error during dynamic batch completion: {}", e.getMessage());
+      LOGGER.warn(
+          "Error during dynamic batch completion, batch size: {}, items: {}",
+          state.batchSize,
+          state.items.size(),
+          e);
+    }
+    try {
+      Thread.sleep(RETRY_DELAY_MS);
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
     }
   }
 
   @Override
   public void close() {
+    this.running = false;
     this.thread.interrupt();
+    try {
+      this.thread.join(TimeUnit.SECONDS.toMillis(5));
+      if (this.thread.isAlive()) {
+        LOGGER.warn("Dynamic batch thread did not terminate within timeout");
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      LOGGER.warn("Interrupted while waiting for dynamic batch thread to terminate");
+    }
+    // Process any remaining items in the queue
+    if (!this.requests.isEmpty()) {
+      List<T> remaining = new ArrayList<>();
+      this.requests.drainTo(remaining);
+      if (!remaining.isEmpty()) {
+        try {
+          this.consumer.process(remaining);
+        } catch (Exception e) {
+          LOGGER.warn("Error processing remaining {} items during shutdown", remaining.size(), e);
+        }
+      }
+    }
   }
 
   @FunctionalInterface

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

@@ -112,6 +112,7 @@ public StreamProducerBuilder batchPublishingDelay(Duration batchPublishingDelay)
   }
 
   @Override
+  @SuppressWarnings("deprecation")
   public ProducerBuilder dynamicBatch(boolean dynamicBatch) {
     this.dynamicBatch = dynamicBatch;
     return this;