fix the pagination issue in checkpointBatcher (#54)

Author: zhongkechen

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

@@ -5,13 +5,14 @@
 import com.amazonaws.lambda.durable.client.DurableExecutionClient;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Supplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 
 /**
@@ -30,31 +31,31 @@ class CheckpointBatcher {
     private static final Logger logger = LoggerFactory.getLogger(CheckpointBatcher.class);
 
     private final CheckpointCallback callback;
-    private final Supplier<String> tokenSupplier;
     private final String durableExecutionArn;
     private final DurableExecutionClient client;
     private final BlockingQueue<CheckpointRequest> queue = new LinkedBlockingQueue<>();
     private final AtomicBoolean isProcessing = new AtomicBoolean(false);
+    private String checkpointToken;
 
     record CheckpointRequest(OperationUpdate update, CompletableFuture<Void> completion) {}
 
     CheckpointBatcher(
             DurableExecutionClient client,
             String durableExecutionArn,
-            Supplier<String> tokenSupplier,
+            String checkpointToken,
             CheckpointCallback callback) {
         this.client = client;
         this.durableExecutionArn = durableExecutionArn;
-        this.tokenSupplier = tokenSupplier;
         this.callback = callback;
+        this.checkpointToken = checkpointToken;
     }
 
     CompletableFuture<Void> checkpoint(OperationUpdate update) {
         logger.debug(
                 "Checkpoint request received: Action {}",
                 update != null ? update.action() : "NULL (Checkpoint request)");
         var future = new CompletableFuture<Void>();
-        queue.offer(new CheckpointRequest(update, future));
+        queue.add(new CheckpointRequest(update, future));
 
         if (isProcessing.compareAndSet(false, true)) {
             InternalExecutor.INSTANCE.execute(this::processQueue);
@@ -70,32 +71,46 @@ void shutdown() {
                 req -> req.completion().completeExceptionally(new IllegalStateException("CheckpointManager shutdown")));
     }
 
+    public List<Operation> fetchAllPages(List<Operation> initialOperations, String nextMarker) {
+        List<Operation> operations = new ArrayList<>();
+        if (initialOperations != null) {
+            operations.addAll(initialOperations);
+        }
+        while (nextMarker != null && !nextMarker.isEmpty()) {
+            var response = client.getExecutionState(durableExecutionArn, checkpointToken, nextMarker);
+            logger.debug("DAR getExecutionState called: {}.", response);
+            operations.addAll(response.operations());
+            nextMarker = response.nextMarker();
+        }
+        return operations;
+    }
+
     private void processQueue() {
         try {
             var batch = collectBatch();
             if (!batch.isEmpty()) {
                 // Filter out null updates (empty checkpoints for polling)
                 var updates = batch.stream()
                         .map(CheckpointRequest::update)
-                        .filter(u -> u != null)
+                        .filter(Objects::nonNull)
                         .toList();
 
-                var response = client.checkpoint(durableExecutionArn, tokenSupplier.get(), updates);
-                logger.debug("DAR backend called: {}.", response);
+                var response = client.checkpoint(durableExecutionArn, checkpointToken, updates);
+                logger.debug("DAR checkpointDurableExecution called: {}.", response);
 
                 // Notify callback of completion
                 // TODO: sam local backend returns no new execution state when called with zero
                 // updates. WHY?
                 // This means the polling will never receive an operation update and complete
                 // the Phaser.
-                if (response.newExecutionState() != null
-                        && response.newExecutionState().operations() != null
-                        && !response.newExecutionState().operations().isEmpty()) {
-                    callback.onComplete(
-                            response.checkpointToken(),
-                            response.newExecutionState().operations());
-                } else {
-                    callback.onComplete(response.checkpointToken(), List.of());
+                checkpointToken = response.checkpointToken();
+                if (response.newExecutionState() != null) {
+                    var operations = fetchAllPages(
+                            response.newExecutionState().operations(),
+                            response.newExecutionState().nextMarker());
+                    if (!operations.isEmpty()) {
+                        callback.onComplete(operations);
+                    }
                 }
 
                 batch.forEach(req -> req.completion().complete(null));
@@ -122,7 +137,7 @@ private List<CheckpointRequest> collectBatch() {
             var itemSize = estimateSize(req.update());
 
             if (currentSize + itemSize > MAX_BATCH_SIZE_BYTES && !batch.isEmpty()) {
-                queue.offer(req);
+                queue.add(req);
                 break;
             }
 

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

@@ -17,8 +17,7 @@ interface CheckpointCallback {
     /**
      * Called when a checkpoint completes successfully.
      *
-     * @param newToken New checkpoint token from backend
      * @param newOperations Updated operations from backend
      */
-    void onComplete(String newToken, List<Operation> newOperations);
+    void onComplete(List<Operation> newOperations);
 }

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

@@ -11,9 +11,9 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.lambda.model.Operation;
@@ -43,9 +43,8 @@ public class ExecutionManager {
     private static final Logger logger = LoggerFactory.getLogger(ExecutionManager.class);
 
     // ===== Execution State =====
-    private final Map<String, Operation> operations = new ConcurrentHashMap<>();
+    private final Map<String, Operation> operations;
     private final String executionOperationId;
-    private volatile String checkpointToken;
     private final String durableExecutionArn;
     private final AtomicReference<ExecutionMode> executionMode;
 
@@ -57,46 +56,28 @@ public class ExecutionManager {
 
     // ===== Checkpoint Batching =====
     private final CheckpointBatcher checkpointBatcher;
-    private final DurableExecutionClient client;
 
     public ExecutionManager(
             String durableExecutionArn,
             String checkpointToken,
             InitialExecutionState initialExecutionState,
             DurableExecutionClient client) {
         this.durableExecutionArn = durableExecutionArn;
-        this.checkpointToken = checkpointToken;
-        this.client = client;
         this.executionOperationId = initialExecutionState.operations().get(0).id();
-        loadAllOperations(initialExecutionState);
+
+        // Create checkpoint batcher for internal coordination
+        this.checkpointBatcher =
+                new CheckpointBatcher(client, durableExecutionArn, checkpointToken, this::onCheckpointComplete);
+
+        this.operations =
+                checkpointBatcher
+                        .fetchAllPages(initialExecutionState.operations(), initialExecutionState.nextMarker())
+                        .stream()
+                        .collect(Collectors.toConcurrentMap(Operation::id, op -> op));
 
         // Start in REPLAY mode if we have more than just the initial EXECUTION operation
         this.executionMode =
                 new AtomicReference<>(operations.size() > 1 ? ExecutionMode.REPLAY : ExecutionMode.EXECUTION);
-
-        // Create checkpoint manager using common pool for internal coordination
-        this.checkpointBatcher = new CheckpointBatcher(
-                client, durableExecutionArn, this::getCheckpointToken, this::onCheckpointComplete);
-    }
-
-    private void loadAllOperations(InitialExecutionState initialExecutionState) {
-        var initialOperations = initialExecutionState.operations();
-        initialOperations.forEach(op -> operations.put(op.id(), op));
-
-        var nextMarker = initialExecutionState.nextMarker();
-        while (nextMarker != null && !nextMarker.isEmpty()) {
-            var response = client.getExecutionState(durableExecutionArn, checkpointToken, nextMarker);
-            response.operations().forEach(op -> operations.put(op.id(), op));
-            nextMarker = response.nextMarker();
-        }
-    }
-
-    // ===== Checkpoint Completion Handler =====
-
-    /** Called by CheckpointManager when a checkpoint completes. Updates state and advances phasers. */
-    private void onCheckpointComplete(String newToken, List<Operation> newOperations) {
-        this.checkpointToken = newToken;
-        updateOperations(newOperations);
     }
 
     // ===== State Management =====
@@ -109,11 +90,9 @@ public boolean isReplaying() {
         return executionMode.get() == ExecutionMode.REPLAY;
     }
 
-    public String getCheckpointToken() {
-        return checkpointToken;
-    }
-
-    private void updateOperations(List<Operation> newOperations) {
+    // ===== Checkpoint Completion Handler =====
+    /** Called by CheckpointManager when a checkpoint completes. Updates state and advances phasers. */
+    private void onCheckpointComplete(List<Operation> newOperations) {
         // Update operation storage
         newOperations.forEach(op -> operations.put(op.id(), op));