poll again if the previous poll doesn't receive an update

Author: zhongkechen

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

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 package com.amazonaws.lambda.durable.execution;
 
-import com.amazonaws.lambda.durable.exception.IllegalDurableOperationException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
@@ -140,14 +139,14 @@ private void flushNow() {
         execute();
     }
 
+    /** Flushes pending batch and waits for completion */
     public void shutdown() {
-        var ex = new IllegalDurableOperationException("Batch cancelled");
         synchronized (items) {
-            for (Item<T> item : items) {
-                item.result().completeExceptionally(ex);
-            }
-            initializeBatch();
+            flushNow();
         }
+
+        // wait for previous batches to be flushed
+        previousBatchFuture.join();
     }
 
     /** Executes batch and completes all item futures */

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

@@ -3,7 +3,6 @@
 package com.amazonaws.lambda.durable.execution;
 
 import com.amazonaws.lambda.durable.DurableConfig;
-
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -59,10 +58,12 @@ CompletableFuture<Void> checkpoint(OperationUpdate update) {
         return checkpointApiRequestBatcher.submit(update, config.getCheckpointDelay());
     }
 
+    /** Poll for updates of the specified operation with preconfigured intervals */
     CompletableFuture<Operation> pollForUpdate(String operationId) {
         return pollForUpdate(operationId, config.getPollingInterval());
     }
 
+    /** Poll for updates of the specified operation with specified delay */
     CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
         logger.debug("Polling request received: operation id {}", operationId);
         var future = new CompletableFuture<Operation>();
@@ -72,14 +73,21 @@ CompletableFuture<Operation> pollForUpdate(String operationId, Duration delay) {
                     .computeIfAbsent(operationId, k -> Collections.synchronizedList(new ArrayList<>()))
                     .add(future);
         }
-        checkpointApiRequestBatcher.submit(null, delay);
+        checkpointApiRequestBatcher.submit(null, delay).thenCompose(v -> {
+            if (future.isDone()) {
+                return CompletableFuture.completedFuture(null);
+            }
+            return checkpointApiRequestBatcher.submit(null, delay);
+        });
         return future;
     }
 
     void shutdown() {
-        List<List<CompletableFuture<Operation>>> allFutures;
+        // wait for all checkpoint requests to complete
         checkpointApiRequestBatcher.shutdown();
 
+        // complete all polling futures with an exception
+        List<List<CompletableFuture<Operation>>> allFutures;
         synchronized (pollingFutures) {
             allFutures = new ArrayList<>(pollingFutures.values());
             pollingFutures.clear();
@@ -90,6 +98,7 @@ void shutdown() {
         }
     }
 
+    /** Calling GetExecutionState API to get all pages of operations given the nextMarker */
     public List<Operation> fetchAllPages(List<Operation> initialOperations, String nextMarker) {
         List<Operation> operations = new ArrayList<>();
         if (initialOperations != null) {

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

@@ -5,7 +5,6 @@
 import com.amazonaws.lambda.durable.DurableConfig;
 import com.amazonaws.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import com.amazonaws.lambda.durable.model.DurableExecutionInput.InitialExecutionState;
-
 import java.time.Duration;
 import java.util.Collections;
 import java.util.HashMap;

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

@@ -13,7 +13,6 @@
 import com.amazonaws.lambda.durable.execution.ThreadType;
 import com.amazonaws.lambda.durable.serde.SerDes;
 import com.amazonaws.lambda.durable.util.ExceptionHelper;
-
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
@@ -209,7 +208,7 @@ protected CompletableFuture<Operation> pollForOperationUpdates() {
         return executionManager.pollForOperationUpdates(operationId);
     }
 
-    protected  CompletableFuture<Operation> pollForOperationUpdates(Duration delay) {
+    protected CompletableFuture<Operation> pollForOperationUpdates(Duration delay) {
         return executionManager.pollForOperationUpdates(operationId, delay);
     }