rename methods/fields and update comments

Author: zhongkechen

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

@@ -9,218 +9,192 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
-import java.util.stream.Collectors;
 
 /**
- * This class simplifies automatic batching of api requests. The individual request items will be grouped if the service
- * has a cheaper batch API, and we want to trade some latency by waiting for more calls to arrive. The batch call will
- * be made when either a full batch is built, too much time has passed, or size limits are reached. This class builds a
- * single batch at a time with thread-safe synchronization: - There is no batch yet. - First call arrives. Create a
- * batch with one item in it, start a timer. No call to service is made yet. - More calls arrive. They get added to the
- * same batch if size limits allow. - Either the batch is full, the timer has elapsed, or size limits are reached. Send
- * the batch request. Now a new batch can now be built. - If entire batch call fails, each call will fail. - If batch
- * call succeeded, outcome is analyzed one by one to complete results of each call. When you extend this class, you are
- * expected to implement the actual batch operation and to expose a public method to perform a single action. The
- * batcher includes comprehensive metrics tracking for performance monitoring.
+ * 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> Input of every call
+ * @param <T> Request type
  */
 public class ApiRequestBatcher<T> {
 
-    /** Maximum time to wait before flushing a batch */
-    private final Duration maxDelay;
-    /** Maximum number of items per batch */
-    private final int maxBatchSize;
-    /** Maximum total size in bytes for all items in a batch */
-    private final int maxBatchBinarySizeInBytes;
-    /** Function to calculate the size in bytes of each item */
-    private final Function<T, Integer> itemSizeInBytesProvider;
+    /** 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 Function<List<T>, CompletableFuture<Void>> doBatchAction;
+    private record Item<T>(T request, CompletableFuture<Void> result) {}
 
-    private record BatchItem<T>(T input, CompletableFuture<Void> outputFuture) {}
-
-    /** Represents a collection of items to be processed together as a batch. */
+    /** Batch accumulator */
     private class Batch {
-        /** List of items in this batch */
-        private final List<BatchItem<T>> batchItems;
-        /** Total size in bytes of all items in this batch */
-        private int batchSizeInBytes;
+        /** Accumulated requests */
+        private final List<Item<T>> items;
+        /** Current batch size in bytes */
+        private int totalBytes;
 
         Batch() {
-            this.batchItems = new ArrayList<>();
-        }
-
-        /**
-         * Adds an item to this batch and returns a future for the result.
-         *
-         * @param input The item to add to the batch
-         * @return A CompletableFuture that will be completed with the result
-         */
-        CompletableFuture<Void> addItem(T input) {
-            final int itemSize = itemSizeInBytesProvider.apply(input);
-            batchSizeInBytes += itemSize;
-
-            CompletableFuture<Void> resultFuture = new CompletableFuture<>();
-            batchItems.add(new BatchItem<>(input, resultFuture));
-            return resultFuture;
-        }
-
-        /** Checks if this batch can accept the given item without exceeding size limits. */
-        boolean canAcceptItem(T input) {
-            return batchSizeInBytes + itemSizeInBytesProvider.apply(input) <= maxBatchBinarySizeInBytes;
+            this.items = new ArrayList<>();
         }
 
-        /** Checks if this batch can accept more items without exceeding count limits. */
-        boolean canAcceptMore() {
-            return batchItems.size() < maxBatchSize;
+        /** Adds request to batch and returns its result future */
+        CompletableFuture<Void> add(T request) {
+            totalBytes += calculateItemSize.apply(request);
+            CompletableFuture<Void> result = new CompletableFuture<>();
+            items.add(new Item<>(request, result));
+            return result;
         }
 
-        /** Processes this batch by executing the batch action and handling results. */
-        void processBatch() {
-            List<T> inputs = extractInputs();
-
-            CompletableFuture<Void> batchFuture = doBatchAction.apply(inputs);
-
-            batchFuture.thenAccept(this::completeItems).exceptionally(this::failAllItems);
+        /** Returns true if request fits within byte limit */
+        boolean canFit(T request) {
+            return totalBytes + calculateItemSize.apply(request) <= maxBatchBytes;
         }
 
-        private List<T> extractInputs() {
-            return batchItems.stream().map(BatchItem::input).collect(Collectors.toList());
+        /** Returns true if batch has reached item count limit */
+        boolean isFull() {
+            return items.size() >= maxItemCount;
         }
 
-        /** Completes individual item futures with their corresponding results */
-        private void completeItems(Void v) {
-            for (BatchItem<T> batchItem : batchItems) {
-                batchItem.outputFuture().complete(null);
+        /** Executes batch and completes all item futures */
+        void execute() {
+            List<T> requests = new ArrayList<>(items.size());
+            for (Item<T> item : items) {
+                requests.add(item.request());
             }
-        }
 
-        /** Fails all item futures with the given exception */
-        private Void failAllItems(Throwable wrappedCause) {
-            Throwable cause = ExceptionHelper.unwrapCompletableFuture(wrappedCause);
-            for (BatchItem<T> batchItem : batchItems) {
-                batchItem.outputFuture().completeExceptionally(cause);
-            }
-            return null;
+            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);
+                    }
+                }
+            });
         }
     }
 
-    /** Lock for synchronizing access to current batch state */
-    private final Object currentBatchLock = new Object();
-    /** The current batch being filled with items */
-    private Batch currentBatch;
-    /** Future that completes when the current batch should be flushed due to timeout */
-    private CompletableFuture<Void> batchFlushFuture;
+    /** 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;
 
     /**
      * Creates a new ApiRequestBatcher with the specified configuration.
      *
-     * @param maxDelay Maximum time to wait before flushing a batch
-     * @param maxBatchSize Maximum number of items per batch
-     * @param maxBatchBinarySizeInBytes Maximum total size in bytes for all items in a batch
-     * @param itemSizeInBytesProvider Function to calculate the size in bytes of each item
+     * @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 maxDelay,
-            int maxBatchSize,
-            int maxBatchBinarySizeInBytes,
-            Function<T, Integer> itemSizeInBytesProvider,
-            Function<List<T>, CompletableFuture<Void>> doBatchAction) {
-        this.maxDelay = maxDelay;
-        this.maxBatchSize = maxBatchSize;
-        this.maxBatchBinarySizeInBytes = maxBatchBinarySizeInBytes;
-        this.itemSizeInBytesProvider = itemSizeInBytesProvider;
-        this.doBatchAction = doBatchAction;
-        this.currentBatch = null;
-        this.batchFlushFuture = null;
+            Duration flushDelay,
+            int maxItemCount,
+            int maxBatchBytes,
+            Function<T, Integer> calculateItemSize,
+            Function<List<T>, CompletableFuture<Void>> executeBatch) {
+        this.flushDelay = flushDelay;
+        this.maxItemCount = maxItemCount;
+        this.maxBatchBytes = maxBatchBytes;
+        this.calculateItemSize = calculateItemSize;
+        this.executeBatch = executeBatch;
     }
 
     /**
-     * Submits an item for batch processing. The item will be added to the current batch or trigger batch processing if
-     * limits are reached.
+     * Submits request for batched execution.
      *
-     * @param input The item to process
-     * @return A CompletableFuture that will be completed with the processing result
+     * @param request Request to batch
+     * @return Future completed when batch executes
      */
-    public CompletableFuture<Void> doAction(T input) {
-        CompletableFuture<Void> outputFuture;
-        Batch previousBatchToProcess = null;
-        Batch newBatchToProcess = null;
-
-        synchronized (currentBatchLock) {
-            // If current batch can't fit this item, flush it first
-            if (currentBatch != null && !currentBatch.canAcceptItem(input)) {
-                previousBatchToProcess = getAndClearCurrentBatch();
+    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 new batch if needed
-            if (currentBatch == null) {
-                currentBatch = new Batch();
-                if (batchFlushFuture != null) {
-                    cancelAndClearCurrentFlusher();
+            // Create batch and start timer
+            if (activeBatch == null) {
+                activeBatch = new Batch();
+                if (flushTimer != null) {
+                    cancelTimer();
                 }
             }
 
-            outputFuture = currentBatch.addItem(input);
+            result = activeBatch.add(request);
 
-            // If batch is full, process it immediately
-            if (!currentBatch.canAcceptMore()) {
-                newBatchToProcess = getAndClearCurrentBatch();
+            // Flush if batch full
+            if (activeBatch.isFull()) {
+                fullBatch = detachBatch();
             }
 
-            // Set up timeout-based flushing for non-full batches
-            if (currentBatch != null && batchFlushFuture == null) {
-                batchFlushFuture = new CompletableFuture<>();
-                batchFlushFuture
-                        .completeOnTimeout(null, maxDelay.toMillis(), TimeUnit.MILLISECONDS)
+            // 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 (currentBatchLock) {
-                                if (currentBatch != null) {
-                                    toFlush = getAndClearCurrentBatch();
+                            synchronized (lock) {
+                                if (activeBatch != null) {
+                                    toFlush = detachBatch();
                                 } else {
                                     return;
                                 }
                             }
-                            toFlush.processBatch();
+                            toFlush.execute();
                         });
             }
 
-            // Clean up flush future if no current batch
-            if (currentBatch == null && batchFlushFuture != null) {
-                cancelAndClearCurrentFlusher();
+            // Cancel timer if batch was flushed
+            if (activeBatch == null && flushTimer != null) {
+                cancelTimer();
             }
         }
 
-        // Process batches outside of synchronized block to avoid blocking
-        if (previousBatchToProcess != null) {
-            previousBatchToProcess.processBatch();
+        // Execute outside lock to avoid blocking
+        if (batchToFlush != null) {
+            batchToFlush.execute();
         }
 
-        if (newBatchToProcess != null) {
-            newBatchToProcess.processBatch();
+        if (fullBatch != null) {
+            fullBatch.execute();
         }
 
-        return outputFuture;
+        return result;
     }
 
-    /** Gets the current batch and clears it, ensuring it's not null */
-    private Batch getAndClearCurrentBatch() {
-        if (currentBatch == null) {
-            throw new IllegalStateException("currentBatch must not be null");
+    /** Detaches and returns active batch */
+    private Batch detachBatch() {
+        if (activeBatch == null) {
+            throw new IllegalStateException("activeBatch must not be null");
         }
-        final Batch batchToProcess = currentBatch;
-        currentBatch = null;
-        return batchToProcess;
+        Batch batch = activeBatch;
+        activeBatch = null;
+        return batch;
     }
 
-    /** Cancels the current flush future and clears it, ensuring it's not null */
-    private void cancelAndClearCurrentFlusher() {
-        if (batchFlushFuture == null) {
-            throw new IllegalStateException("batchFlushFuture must not be null");
+    /** Cancels and clears flush timer */
+    private void cancelTimer() {
+        if (flushTimer == null) {
+            throw new IllegalStateException("flushTimer must not be null");
         }
-        batchFlushFuture.cancel(false);
-        batchFlushFuture = null;
+        flushTimer.cancel(false);
+        flushTimer = null;
     }
 }