remove locker from batcher

Author: zhongkechen

sdk/src/main/java/com/amazonaws/lambda/durable/execution/ApiRequestBatcher.java

@@ -2,32 +2,33 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.amazonaws.lambda.durable.execution;
 
-import com.amazonaws.lambda.durable.util.ExceptionHelper;
+import com.amazonaws.lambda.durable.exception.IllegalDurableOperationException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
 import java.util.function.Function;
 
 /**
- * Batches API requests to optimize throughput by grouping individual calls into batch operations.
- * Batches are flushed when full, when size limits are reached, or after a timeout.
+ * Batches API requests to optimize throughput by grouping individual calls into batch operations. Batches are flushed
+ * when full, when size limits are reached, or after a timeout.
  *
  * @param <T> Request type
  */
 public class ApiRequestBatcher<T> {
+    private static final Duration MAX_DELAY = Duration.ofMinutes(60);
 
-    /** Timeout before auto-flushing incomplete batch */
-    private final Duration flushDelay;
     /** Maximum items allowed in a single batch */
     private final int maxItemCount;
     /** Maximum bytes allowed in a single batch */
     private final int maxBatchBytes;
     /** Calculates byte size of each request */
     private final Function<T, Integer> calculateItemSize;
     /** Executes the batch operation */
-    private final Function<List<T>, CompletableFuture<Void>> executeBatch;
+    private final Consumer<List<T>> executeBatch;
 
     private record Item<T>(T request, CompletableFuture<Void> result) {}
 
@@ -38,15 +39,29 @@ private class Batch {
         /** Current batch size in bytes */
         private int totalBytes;
 
+        long expireTime;
+        /** Timer to auto-flush incomplete batch */
+        private final CompletableFuture<Void> flushTimer;
+
         Batch() {
             this.items = new ArrayList<>();
+            this.totalBytes = 0;
+            this.expireTime = System.nanoTime() + MAX_DELAY.toNanos();
+            this.flushTimer = new CompletableFuture<>();
+            this.flushTimer.thenRunAsync(this::execute, InternalExecutor.INSTANCE);
         }
 
         /** Adds request to batch and returns its result future */
-        CompletableFuture<Void> add(T request) {
+        CompletableFuture<Void> add(T request, Duration delay) {
             totalBytes += calculateItemSize.apply(request);
             CompletableFuture<Void> result = new CompletableFuture<>();
             items.add(new Item<>(request, result));
+            long newExpireTime = System.nanoTime() + delay.toNanos();
+            if (expireTime > newExpireTime) {
+                // the batch needs to be completed earlier than previously scheduled
+                expireTime = newExpireTime;
+                flushAfterDelay(delay.toNanos());
+            }
             return result;
         }
 
@@ -60,55 +75,65 @@ boolean isFull() {
             return items.size() >= maxItemCount;
         }
 
+        void flushAfterDelay(long delayInNanos) {
+            flushTimer.completeOnTimeout(null, delayInNanos, TimeUnit.NANOSECONDS);
+        }
+
+        void flushNow() {
+            flushAfterDelay(0);
+        }
+
+        void cancel() {
+            var ex = new IllegalDurableOperationException("Batch cancelled");
+            for (Item<T> item : items) {
+                item.result().completeExceptionally(ex);
+            }
+        }
+
         /** Executes batch and completes all item futures */
-        void execute() {
+        private void execute() {
+            // detach this from active batch if it's still active
+            detachActiveBatchAndCreateNew(this);
+
             List<T> requests = new ArrayList<>(items.size());
             for (Item<T> item : items) {
                 requests.add(item.request());
             }
 
-            executeBatch.apply(requests).whenComplete((v, ex) -> {
-                if (ex == null) {
-                    for (Item<T> item : items) {
-                        item.result().complete(null);
-                    }
-                } else {
-                    Throwable cause = ExceptionHelper.unwrapCompletableFuture(ex);
-                    for (Item<T> item : items) {
-                        item.result().completeExceptionally(cause);
-                    }
+            try {
+                executeBatch.accept(requests);
+                for (Item<T> item : items) {
+                    item.result().complete(null);
+                }
+            } catch (Throwable ex) {
+                for (Item<T> item : items) {
+                    item.result().completeExceptionally(ex);
                 }
-            });
+            }
         }
     }
 
-    /** Synchronizes batch state access */
-    private final Object lock = new Object();
     /** Current batch accepting requests */
-    private Batch activeBatch;
-    /** Timer to auto-flush incomplete batch */
-    private CompletableFuture<Void> flushTimer;
+    private final AtomicReference<Batch> activeBatchAtom;
 
     /**
      * Creates a new ApiRequestBatcher with the specified configuration.
      *
-     * @param flushDelay Maximum time to wait before flushing a batch
      * @param maxItemCount Maximum number of items per batch
      * @param maxBatchBytes Maximum total size in bytes for all items in a batch
      * @param calculateItemSize Function to calculate the size in bytes of each item
      * @param executeBatch Function to execute the batch action
      */
     public ApiRequestBatcher(
-            Duration flushDelay,
             int maxItemCount,
             int maxBatchBytes,
             Function<T, Integer> calculateItemSize,
-            Function<List<T>, CompletableFuture<Void>> executeBatch) {
-        this.flushDelay = flushDelay;
+            Consumer<List<T>> executeBatch) {
         this.maxItemCount = maxItemCount;
         this.maxBatchBytes = maxBatchBytes;
         this.calculateItemSize = calculateItemSize;
         this.executeBatch = executeBatch;
+        this.activeBatchAtom = new AtomicReference<>(new Batch());
     }
 
     /**
@@ -117,84 +142,51 @@ public ApiRequestBatcher(
      * @param request Request to batch
      * @return Future completed when batch executes
      */
-    public CompletableFuture<Void> submit(T request) {
-        CompletableFuture<Void> result;
-        Batch batchToFlush = null;
-        Batch fullBatch = null;
-
-        synchronized (lock) {
-            // Flush if request doesn't fit
-            if (activeBatch != null && !activeBatch.canFit(request)) {
-                batchToFlush = detachBatch();
-            }
-
-            // Create batch and start timer
-            if (activeBatch == null) {
-                activeBatch = new Batch();
-                if (flushTimer != null) {
-                    cancelTimer();
+    public CompletableFuture<Void> submit(T request, Duration flushDelay) {
+        // Flush the current batch if request doesn't fit
+        while (true) {
+            Batch activeBatch = activeBatchAtom.get();
+
+            if (activeBatch.isFull() || !activeBatch.canFit(request)) {
+                if (!flushActiveBatchAndCreateNew(activeBatch)) {
+                    // failed to flush due to a race condition.
+                    continue;
                 }
             }
 
-            result = activeBatch.add(request);
+            var result = activeBatch.add(request, flushDelay);
 
-            // Flush if batch full
+            // Flush early if batch is full
             if (activeBatch.isFull()) {
-                fullBatch = detachBatch();
+                flushActiveBatchAndCreateNew(activeBatch);
             }
-
-            // Start flush timer for new batch
-            if (activeBatch != null && flushTimer == null) {
-                flushTimer = new CompletableFuture<>();
-                flushTimer
-                        .completeOnTimeout(null, flushDelay.toMillis(), TimeUnit.MILLISECONDS)
-                        .thenRun(() -> {
-                            Batch toFlush;
-                            synchronized (lock) {
-                                if (activeBatch != null) {
-                                    toFlush = detachBatch();
-                                } else {
-                                    return;
-                                }
-                            }
-                            toFlush.execute();
-                        });
-            }
-
-            // Cancel timer if batch was flushed
-            if (activeBatch == null && flushTimer != null) {
-                cancelTimer();
-            }
-        }
-
-        // Execute outside lock to avoid blocking
-        if (batchToFlush != null) {
-            batchToFlush.execute();
+            return result;
         }
+    }
 
-        if (fullBatch != null) {
-            fullBatch.execute();
+    private Batch detachActiveBatchAndCreateNew(Batch oldBatch) {
+        if (activeBatchAtom.compareAndSet(oldBatch, new Batch())) {
+            return oldBatch;
         }
 
-        return result;
+        return null;
     }
 
-    /** Detaches and returns active batch */
-    private Batch detachBatch() {
-        if (activeBatch == null) {
-            throw new IllegalStateException("activeBatch must not be null");
+    /** flushes active batch and crate a new batch. Return true if successful */
+    private boolean flushActiveBatchAndCreateNew(Batch oldBatch) {
+        Batch activeBatch = detachActiveBatchAndCreateNew(oldBatch);
+        if (activeBatch != null) {
+            activeBatch.flushNow();
         }
-        Batch batch = activeBatch;
-        activeBatch = null;
-        return batch;
+        return activeBatch != null;
     }
 
-    /** Cancels and clears flush timer */
-    private void cancelTimer() {
-        if (flushTimer == null) {
-            throw new IllegalStateException("flushTimer must not be null");
+    public void shutdown() {
+        Batch activeBatch = activeBatchAtom.get();
+        while (!activeBatchAtom.compareAndSet(activeBatch, new Batch())) {
+            // try again
+            activeBatch = activeBatchAtom.get();
         }
-        flushTimer.cancel(false);
-        flushTimer = null;
+        activeBatchAtom.get().cancel();
     }
 }