fix race condition

Author: zhongkechen

sdk/src/main/java/com/amazonaws/lambda/durable/DurableConfig.java

@@ -78,7 +78,7 @@ private DurableConfig(Builder builder) {
         this.serDes = builder.serDes != null ? builder.serDes : new JacksonSerDes();
         this.executorService = builder.executorService != null ? builder.executorService : createDefaultExecutor();
         this.loggerConfig = builder.loggerConfig != null ? builder.loggerConfig : LoggerConfig.defaults();
-        this.pollingInterval = builder.pollingInterval != null ? builder.pollingInterval : Duration.ofSeconds(1);
+        this.pollingInterval = builder.pollingInterval != null ? builder.pollingInterval : Duration.ofMillis(1000);
         this.checkpointDelay = builder.checkpointDelay != null ? builder.checkpointDelay : Duration.ofSeconds(0);
     }
 

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

@@ -8,7 +8,6 @@
 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;
 
@@ -30,91 +29,22 @@ public class ApiRequestBatcher<T> {
     /** Executes the batch operation */
     private final Consumer<List<T>> executeBatch;
 
-    private record Item<T>(T request, CompletableFuture<Void> result) {}
-
-    /** Batch accumulator */
-    private class Batch {
-        /** Accumulated requests */
-        private final List<Item<T>> items;
-        /** 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, 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;
-        }
+    /** Accumulated requests */
+    private final List<Item<T>> items;
 
-        /** Returns true if request fits within byte limit */
-        boolean canFit(T request) {
-            return totalBytes + calculateItemSize.apply(request) <= maxBatchBytes;
-        }
-
-        /** Returns true if batch has reached item count limit */
-        boolean isFull() {
-            return items.size() >= maxItemCount;
-        }
+    /** Current batch size in bytes */
+    private volatile int totalBytes;
 
-        void flushAfterDelay(long delayInNanos) {
-            flushTimer.completeOnTimeout(null, delayInNanos, TimeUnit.NANOSECONDS);
-        }
+    /** Time when the current batch must be flushed */
+    private volatile long expireTime;
 
-        void flushNow() {
-            flushAfterDelay(0);
-        }
+    /** Timer to auto-flush incomplete batch */
+    private CompletableFuture<Void> flushTimer;
 
-        void cancel() {
-            var ex = new IllegalDurableOperationException("Batch cancelled");
-            for (Item<T> item : items) {
-                item.result().completeExceptionally(ex);
-            }
-        }
-
-        /** Executes batch and completes all item futures */
-        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());
-            }
-
-            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);
-                }
-            }
-        }
-    }
+    /** Future of flushing previous batch */
+    private CompletableFuture<Void> previousBatchFuture;
 
-    /** Current batch accepting requests */
-    private final AtomicReference<Batch> activeBatchAtom;
+    private record Item<T>(T request, CompletableFuture<Void> result) {}
 
     /**
      * Creates a new ApiRequestBatcher with the specified configuration.
@@ -133,7 +63,22 @@ public ApiRequestBatcher(
         this.maxBatchBytes = maxBatchBytes;
         this.calculateItemSize = calculateItemSize;
         this.executeBatch = executeBatch;
-        this.activeBatchAtom = new AtomicReference<>(new Batch());
+        this.previousBatchFuture = CompletableFuture.allOf();
+        this.items = new ArrayList<>();
+
+        initializeBatch();
+    }
+
+    private void initializeBatch() {
+        this.items.clear();
+        this.totalBytes = 0;
+        this.expireTime = System.nanoTime() + MAX_DELAY.toNanos();
+        this.flushTimer = new CompletableFuture<>();
+        this.flushTimer.thenRun(() -> {
+            synchronized (items) {
+                execute();
+            }
+        });
     }
 
     /**
@@ -144,49 +89,90 @@ public ApiRequestBatcher(
      */
     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;
-                }
+        synchronized (items) {
+            if (isFull() || !canFit(request)) {
+                flushNow();
             }
 
-            var result = activeBatch.add(request, flushDelay);
+            var future = add(request, flushDelay);
 
-            // Flush early if batch is full
-            if (activeBatch.isFull()) {
-                flushActiveBatchAndCreateNew(activeBatch);
+            if (isFull()) {
+                // Flush early if batch is full
+                flushNow();
             }
-            return result;
+            return future;
         }
     }
 
-    private Batch detachActiveBatchAndCreateNew(Batch oldBatch) {
-        if (activeBatchAtom.compareAndSet(oldBatch, new Batch())) {
-            return oldBatch;
+    /** Adds request to batch and returns its result future */
+    CompletableFuture<Void> add(T request, Duration delay) {
+        synchronized (items) {
+            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;
         }
+    }
 
-        return null;
+    /** Returns true if request fits within byte limit */
+    private boolean canFit(T request) {
+        return totalBytes + calculateItemSize.apply(request) <= maxBatchBytes;
     }
 
-    /** 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();
-        }
-        return activeBatch != null;
+    /** Returns true if batch has reached item count limit */
+    private boolean isFull() {
+        return items.size() >= maxItemCount;
+    }
+
+    private void flushAfterDelay(long delayInNanos) {
+        flushTimer.completeOnTimeout(null, delayInNanos, TimeUnit.NANOSECONDS);
+    }
+
+    private void flushNow() {
+        this.flushTimer.cancel(false);
+        // wait for new batch to be ready
+        execute();
     }
 
     public void shutdown() {
-        Batch activeBatch = activeBatchAtom.get();
-        while (!activeBatchAtom.compareAndSet(activeBatch, new Batch())) {
-            // try again
-            activeBatch = activeBatchAtom.get();
+        var ex = new IllegalDurableOperationException("Batch cancelled");
+        synchronized (items) {
+            for (Item<T> item : items) {
+                item.result().completeExceptionally(ex);
+            }
+            initializeBatch();
+        }
+    }
+
+    /** Executes batch and completes all item futures */
+    private void execute() {
+        if (items.isEmpty()) {
+            return;
         }
-        activeBatchAtom.get().cancel();
+        var copyItems = new ArrayList<>(items);
+        initializeBatch();
+
+        // append the current batch to the previous one
+        previousBatchFuture = previousBatchFuture.thenRunAsync(
+                () -> {
+                    try {
+                        var requests = copyItems.stream().map(Item::request).toList();
+                        executeBatch.accept(requests);
+                        for (Item<T> item : copyItems) {
+                            item.result().complete(null);
+                        }
+                    } catch (Throwable ex) {
+                        for (Item<T> item : copyItems) {
+                            item.result().completeExceptionally(ex);
+                        }
+                    }
+                },
+                InternalExecutor.INSTANCE);
     }
 }

sdk/src/main/java/com/amazonaws/lambda/durable/operation/BaseDurableOperation.java

@@ -242,9 +242,12 @@ protected ErrorObject serializeException(Throwable throwable) {
     }
 
     protected Throwable deserializeException(ErrorObject errorObject) {
+        Throwable original = null;
+        if (errorObject == null) {
+            return original;
+        }
         var errorType = errorObject.errorType();
         var errorData = errorObject.errorData();
-        Throwable original = null;
 
         if (errorType == null) {
             return original;

sdk/src/test/java/com/amazonaws/lambda/durable/execution/ApiRequestBatcherTest.java

@@ -4,21 +4,22 @@
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
-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.concurrent.TimeoutException;
 import java.util.function.Consumer;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -31,46 +32,47 @@ class ApiRequestBatcherTest {
     private static class Input {}
 
     private Input input;
-    private Clock fixedClock;
     private ApiRequestBatcher<Input> cut;
     private Consumer<List<Input>> doBatchAction;
 
     @BeforeEach
     void setUp() {
         input = mock(Input.class);
         doBatchAction = mock();
-        fixedClock = Clock.fixed(Instant.now(), ZoneOffset.UTC);
         cut = new ApiRequestBatcher<>(MAX_BATCH_SIZE, MAX_BATCH_BINARY_SIZE_IN_BYTES, item -> 0, doBatchAction);
     }
 
     @Test
     void whenSingleActionPerformed_anUncompletedFutureIsReturned() {
         CompletableFuture<Void> resultFuture = cut.submit(input, MAX_DELAY_MILLIS);
+        assertThrows(TimeoutException.class, () -> resultFuture.get(50, TimeUnit.MILLISECONDS));
 
         verify(doBatchAction, never()).accept(any());
         assertFalse(resultFuture.isDone());
     }
 
     @Test
-    void whenMultipleActionsPerformedBelowMaxBatchSize_anUncompletedFutureIsReturnedEachTime() {
+    void whenMultipleActionsPerformedBelowMaxBatchSize_anUncompletedFutureIsReturnedEachTime()
+            throws ExecutionException, InterruptedException, TimeoutException {
         List<CompletableFuture<Void>> resultFutures = new ArrayList<>();
         for (int i = 0; i < MAX_BATCH_SIZE - 1; i++) {
             resultFutures.add(cut.submit(input, MAX_DELAY_MILLIS));
         }
 
+        assertThrows(TimeoutException.class, () -> resultFutures.get(0).get(50, TimeUnit.MILLISECONDS));
+
         verify(doBatchAction, never()).accept(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++) {
+        for (int i = 0; i < MAX_BATCH_SIZE * 2; i++) {
             resultFutures.add(cut.submit(input, MAX_DELAY_MILLIS));
         }
 
         CompletableFuture.allOf(resultFutures.toArray(CompletableFuture[]::new)).join();
-        verify(doBatchAction).accept(any());
+        verify(doBatchAction, times(2)).accept(any());
     }
 
     @Test
@@ -98,7 +100,7 @@ void whenBatchInvokeThrows_allFuturesCompleteWithThatException() {
     }
 
     @Test
-    void whenBatchInvokeReturnsOutcome_allFuturesCompleteSuccessfully() {
+    void whenBatchInvokeReturns_allFuturesCompleteSuccessfully() {
         Input input1 = mock(Input.class);
         Input input2 = mock(Input.class);
         Input input3 = mock(Input.class);
@@ -118,7 +120,7 @@ void testSubmit_whenCannotAddItemDueToBinarySizeConstraint_thenFlushCurrentBatch
         var future1 = cut.submit(input, MAX_DELAY_MILLIS);
         var future2 = cut.submit(input, MAX_DELAY_MILLIS);
         CompletableFuture.allOf(future1, future2)
-                .thenAccept(v -> verify(doBatchAction).accept(any()))
+                .thenAccept(v -> verify(doBatchAction, times(2)).accept(any()))
                 .join();
     }
 
@@ -159,13 +161,4 @@ void whenBatchInvokeThrowsCompletionException_allFuturesCompleteWithUnwrappedCau
                 });
         CompletableFuture.allOf(resultFuture1, resultFuture2, resultFuture3).join();
     }
-
-    private Throwable getFutureCause(CompletableFuture<?> failedFuture) throws InterruptedException {
-        try {
-            failedFuture.get();
-            return null;
-        } catch (ExecutionException cause) {
-            return cause.getCause();
-        }
-    }
 }