add tests for batcher

zhongkechen · zhongkechen · commit 34c79a89e23f · 2026-02-10T14:06:08.000-08:00
diff --git a/sdk/src/main/java/com/amazonaws/lambda/durable/execution/ApiRequestBatcher.java b/sdk/src/main/java/com/amazonaws/lambda/durable/execution/ApiRequestBatcher.java
@@ -12,21 +12,20 @@
 import java.util.stream.Collectors;
 
 /**
- * This class simplifies automatic batching of async requests. Your code deals with individual items, the service you
- * are calling asynchronously has a cheaper batch API. You are willing to trade some latency by waiting for more calls
- * to arrive to group them in a single batch call into the service. 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.
+ * 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.
  *
  * @param <T> Input of every call
  */
-public class AsyncBatcher<T> {
+public class ApiRequestBatcher<T> {
 
     /** Maximum time to wait before flushing a batch */
     private final Duration maxDelay;
@@ -115,14 +114,14 @@ private Void failAllItems(Throwable wrappedCause) {
     private CompletableFuture<Void> batchFlushFuture;
 
     /**
-     * Creates a new AsyncBatcher with the specified configuration.
+     * 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
      */
-    public AsyncBatcher(
+    public ApiRequestBatcher(
             Duration maxDelay,
             int maxBatchSize,
             int maxBatchBinarySizeInBytes,
diff --git a/sdk/src/main/java/com/amazonaws/lambda/durable/execution/CheckpointBatcher.java b/sdk/src/main/java/com/amazonaws/lambda/durable/execution/CheckpointBatcher.java
@@ -31,6 +31,7 @@
  */
 class CheckpointBatcher {
     private static final int MAX_BATCH_SIZE_BYTES = 750 * 1024; // 750KB
+    private static final int MAX_BATCH_SIZE = 100; // max updates in one batch
     private static final Logger logger = LoggerFactory.getLogger(CheckpointBatcher.class);
 
     private final Consumer<List<Operation>> callback;
@@ -39,9 +40,8 @@ class CheckpointBatcher {
     private final AtomicBoolean isRunning = new AtomicBoolean(true);
     private final Duration pollingInterval;
     private final Map<String, List<CompletableFuture<Operation>>> pollingFutures = new ConcurrentHashMap<>();
-    private final AsyncBatcher<OperationUpdate> checkpointAsyncBatcher;
+    private final ApiRequestBatcher<OperationUpdate> checkpointApiRequestBatcher;
     private String checkpointToken;
-    private final Object batchCheckpointLock = new Object();
 
     CheckpointBatcher(
             DurableConfig config,
@@ -53,9 +53,9 @@ class CheckpointBatcher {
         this.callback = callback;
         this.checkpointToken = checkpointToken;
         this.pollingInterval = config.getPollingInterval();
-        this.checkpointAsyncBatcher = new AsyncBatcher<>(
+        this.checkpointApiRequestBatcher = new ApiRequestBatcher<>(
                 config.getCheckpointDelay(),
-                Integer.MAX_VALUE,
+                MAX_BATCH_SIZE,
                 MAX_BATCH_SIZE_BYTES,
                 CheckpointBatcher::estimateSize,
                 this::doBatchAction);
@@ -65,7 +65,7 @@ class CheckpointBatcher {
 
     CompletableFuture<Void> checkpoint(OperationUpdate update) {
         logger.debug("Checkpoint request received: Action {}", update.action());
-        return checkpointAsyncBatcher.doAction(update);
+        return checkpointApiRequestBatcher.doAction(update);
     }
 
     CompletableFuture<Operation> pollForUpdate(String operationId) {
@@ -76,6 +76,7 @@ CompletableFuture<Operation> pollForUpdate(String operationId) {
             pollingFutures
                     .computeIfAbsent(operationId, k -> Collections.synchronizedList(new ArrayList<>()))
                     .add(future);
+            pollingFutures.notifyAll();
         }
         return future;
     }
@@ -111,17 +112,30 @@ private void processQueue() {
         // background thread running
         while (isRunning.get()) {
             if (!pollingFutures.isEmpty()) {
+                // If pollers exist, poll for updates periodically.
+                // When wake from wait, sleep for an initial delay before calling the API
+                try {
+                    Thread.sleep(pollingInterval.toMillis());
+                } catch (InterruptedException ignored) {
+                    // ignored
+                }
                 doBatchAction(List.of()).join();
-            }
-            try {
-                Thread.sleep(pollingInterval.toMillis());
-            } catch (InterruptedException ignored) {
+            } else {
+                // if empty, waiting for new pollers
+                synchronized (pollingFutures) {
+                    try {
+                        pollingFutures.wait(pollingInterval.toMillis());
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        break;
+                    }
+                }
             }
         }
     }
 
     protected CompletableFuture<Void> doBatchAction(List<OperationUpdate> updates) {
-        // doBatchAction will be called from the polling thread and also from AsyncBatcher.
+        // doBatchAction will be called from the polling thread and also from ApiRequestBatcher.
         // Use synchronized here to make sure no concurrent checkpoint API calls
         synchronized (pollingFutures) {
             logger.debug("Calling durable API checkpointDurableExecution with {} updates", updates.size());
diff --git a/sdk/src/test/java/com/amazonaws/lambda/durable/execution/ApiRequestBatcherTest.java b/sdk/src/test/java/com/amazonaws/lambda/durable/execution/ApiRequestBatcherTest.java
@@ -0,0 +1,189 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package com.amazonaws.lambda.durable.execution;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.timeout;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import java.time.Clock;
+import java.time.Duration;
+import java.time.Instant;
+import java.time.ZoneOffset;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+class ApiRequestBatcherTest {
+    private static final Duration MAX_DELAY_MILLIS = Duration.ofMillis(100);
+    private static final int MAX_BATCH_SIZE = 3;
+    private static final int MAX_BATCH_BINARY_SIZE_IN_BYTES = 100;
+
+    private static class Input {}
+
+    private Input input;
+    private Clock fixedClock;
+    private ApiRequestBatcher<Input> cut;
+    private Function<List<Input>, CompletableFuture<Void>> doBatchAction;
+    private CompletableFuture<Void> batchResultFuture;
+
+    @BeforeEach
+    void setUp() {
+        input = mock(Input.class);
+        doBatchAction = mock();
+        batchResultFuture = new CompletableFuture<>();
+        fixedClock = Clock.fixed(Instant.now(), ZoneOffset.UTC);
+        cut = new ApiRequestBatcher<>(
+                MAX_DELAY_MILLIS, MAX_BATCH_SIZE, MAX_BATCH_BINARY_SIZE_IN_BYTES, item -> 0, doBatchAction);
+
+        when(doBatchAction.apply(any())).thenReturn(batchResultFuture);
+    }
+
+    @Test
+    void whenSingleActionPerformed_anUncompletedFutureIsReturned() {
+        CompletableFuture<Void> resultFuture = cut.doAction(input);
+
+        verify(doBatchAction, never()).apply(any());
+        assertFalse(resultFuture.isDone());
+    }
+
+    @Test
+    void whenMultipleActionsPerformedBelowMaxBatchSize_anUncompletedFutureIsReturnedEachTime() {
+        List<CompletableFuture<Void>> resultFutures = new ArrayList<>();
+        for (int i = 0; i < MAX_BATCH_SIZE - 1; i++) {
+            resultFutures.add(cut.doAction(input));
+        }
+
+        verify(doBatchAction, never()).apply(any());
+        assertTrue(resultFutures.stream().noneMatch(CompletableFuture::isDone));
+    }
+
+    @Test
+    void whenMultipleActionsPerformedMatchingMaxBatchSize_batchInvokeIsPerformed() {
+        List<CompletableFuture<Void>> resultFutures = new ArrayList<>();
+        for (int i = 0; i < MAX_BATCH_SIZE; i++) {
+            resultFutures.add(cut.doAction(input));
+        }
+
+        verify(doBatchAction).apply(any());
+        assertTrue(resultFutures.stream().noneMatch(CompletableFuture::isDone));
+    }
+
+    @Test
+    void whenBatchInvokeThrows_allFuturesCompleteWithThatException() throws InterruptedException {
+        CompletableFuture<Void> resultFuture1 = cut.doAction(input);
+        CompletableFuture<Void> resultFuture2 = cut.doAction(input);
+        CompletableFuture<Void> resultFuture3 = cut.doAction(input);
+
+        assertFalse(resultFuture1.isDone());
+        assertFalse(resultFuture2.isDone());
+        assertFalse(resultFuture3.isDone());
+
+        Throwable batchCause = mock(Throwable.class);
+        batchResultFuture.completeExceptionally(batchCause);
+
+        assertTrue(resultFuture1.isCompletedExceptionally());
+        assertTrue(resultFuture2.isCompletedExceptionally());
+        assertTrue(resultFuture3.isCompletedExceptionally());
+
+        assertEquals(batchCause, getFutureCause(resultFuture1));
+        assertEquals(batchCause, getFutureCause(resultFuture2));
+        assertEquals(batchCause, getFutureCause(resultFuture3));
+    }
+
+    @Test
+    void whenBatchInvokeReturnsOutcome_allFuturesCompleteSuccessfully() {
+        Input input1 = mock(Input.class);
+        Input input2 = mock(Input.class);
+        Input input3 = mock(Input.class);
+
+        CompletableFuture<Void> resultFuture1 = cut.doAction(input1);
+        CompletableFuture<Void> resultFuture2 = cut.doAction(input2);
+        CompletableFuture<Void> resultFuture3 = cut.doAction(input3);
+
+        assertFalse(resultFuture1.isDone());
+        assertFalse(resultFuture2.isDone());
+        assertFalse(resultFuture3.isDone());
+
+        batchResultFuture.complete(null);
+
+        assertTrue(resultFuture1.isDone());
+        assertTrue(resultFuture2.isDone());
+        assertTrue(resultFuture3.isDone());
+
+        assertFalse(resultFuture1.isCompletedExceptionally());
+        assertFalse(resultFuture2.isCompletedExceptionally());
+        assertFalse(resultFuture3.isCompletedExceptionally());
+    }
+
+    @Test
+    void testDoAction_whenCannotAddItemDueToBinarySizeConstraint_thenFlushCurrentBatchAndCreateNewOne() {
+        var cut = new ApiRequestBatcher<>(
+                MAX_DELAY_MILLIS,
+                MAX_BATCH_SIZE,
+                MAX_BATCH_BINARY_SIZE_IN_BYTES,
+                item -> MAX_BATCH_BINARY_SIZE_IN_BYTES,
+                doBatchAction);
+        List<CompletableFuture<Void>> resultFutures = new ArrayList<>();
+
+        resultFutures.add(cut.doAction(input));
+        resultFutures.add(cut.doAction(input));
+
+        verify(doBatchAction).apply(any());
+
+        assertTrue(resultFutures.stream().noneMatch(CompletableFuture::isDone));
+    }
+
+    @Test
+    void whenTimerFires_batchIsProcessed() {
+        var timerCut = new ApiRequestBatcher<>(
+                Duration.ofMillis(1), MAX_BATCH_SIZE, MAX_BATCH_BINARY_SIZE_IN_BYTES, item -> 0, doBatchAction);
+
+        CompletableFuture<Void> resultFuture = timerCut.doAction(input);
+
+        // Wait for the timeout to trigger
+        CompletableFuture.delayedExecutor(10, TimeUnit.MILLISECONDS).execute(() -> {});
+
+        verify(doBatchAction, timeout(50)).apply(any());
+        assertFalse(resultFuture.isDone());
+    }
+
+    @Test
+    void whenBatchInvokeThrowsCompletionException_allFuturesCompleteWithUnwrappedCause() throws InterruptedException {
+        CompletableFuture<Void> resultFuture1 = cut.doAction(input);
+        CompletableFuture<Void> resultFuture2 = cut.doAction(input);
+        CompletableFuture<Void> resultFuture3 = cut.doAction(input);
+
+        RuntimeException rootCause = new RuntimeException("Root cause");
+        batchResultFuture.completeExceptionally(rootCause);
+
+        assertTrue(resultFuture1.isCompletedExceptionally());
+        assertTrue(resultFuture2.isCompletedExceptionally());
+        assertTrue(resultFuture3.isCompletedExceptionally());
+
+        // Should get unwrapped root cause, not the CompletionException wrapper
+        assertEquals(rootCause, getFutureCause(resultFuture1));
+        assertEquals(rootCause, getFutureCause(resultFuture2));
+        assertEquals(rootCause, getFutureCause(resultFuture3));
+    }
+
+    private Throwable getFutureCause(CompletableFuture<?> failedFuture) throws InterruptedException {
+        try {
+            failedFuture.get();
+            return null;
+        } catch (ExecutionException cause) {
+            return cause.getCause();
+        }
+    }
+}
