fix: prevent invalid PENDING suspension after checkpoint completion

.github/workflows/e2e-tests.yml

@@ -89,6 +89,8 @@ jobs:
             --resolve-image-repos --resolve-s3 --parameter-overrides \
             'ParameterKey=Architecture,ParameterValue=x86_64 ParameterKey=JavaVersion,ParameterValue=java${{ matrix.java }} ParameterKey=FunctionNamePrefix,ParameterValue=Java${{ matrix.java }}- ParameterKey=RoleArn,ParameterValue=${{ secrets.TEST_LAMBDA_EXECUTION_ROLE_ARN }}'
         working-directory: ./examples
+      - name: Record E2E log start time
+        run: echo "E2E_LOG_START_TIME_MS=$(date +%s%3N)" >> "$GITHUB_ENV"
       - name: Cloud Based Integration Tests
         run: |
           mvn clean test -B \
@@ -102,6 +104,45 @@ jobs:
             -Djunit.jupiter.execution.parallel.config.strategy=fixed \
             -Djunit.jupiter.execution.parallel.config.fixed.parallelism=${{ env.E2E_TEST_PARALLELISM }}
         working-directory: ./examples
+      - name: Check checkpoint invariant logs
+        env:
+          E2E_STACK_NAME: Java${{ matrix.java }}-JavaSDKCloudBasedIntegrationTestStack
+        run: |
+          set -euo pipefail
+
+          # Allow the final Lambda invocation logs to reach CloudWatch before checking every
+          # managed function log group created by the E2E stack.
+          sleep 10
+
+          log_groups=$(aws cloudformation list-stack-resources \
+            --stack-name "$E2E_STACK_NAME" \
+            --query "StackResourceSummaries[?ResourceType=='AWS::Logs::LogGroup'].PhysicalResourceId" \
+            --output text)
+
+          if [[ -z "$log_groups" ]]; then
+            echo "::error::No managed Lambda log groups found for $E2E_STACK_NAME"
+            exit 1
+          fi
+
+          invariant_violated=false
+          for log_group in $log_groups; do
+            messages=$(aws logs filter-log-events \
+              --log-group-name "$log_group" \
+              --start-time "$E2E_LOG_START_TIME_MS" \
+              --filter-pattern '"Checkpoint invariant violation"' \
+              --query "events[].message" \
+              --output text)
+
+            if [[ -n "$messages" ]]; then
+              invariant_violated=true
+              echo "::error title=Checkpoint invariant violation::Found in $log_group"
+              echo "$messages"
+            fi
+          done
+
+          if [[ "$invariant_violated" == true ]]; then
+            exit 1
+          fi
       - name: Publish test case summary
         if: always()
         env:

sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java

@@ -11,6 +11,7 @@
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.BooleanSupplier;
 import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -40,17 +41,31 @@ class CheckpointManager {
     private final Map<String, List<CompletableFuture<Operation>>> pollingFutures = new ConcurrentHashMap<>();
     private final ApiRequestDelayedBatcher<OperationUpdate> checkpointApiRequestDelayedBatcher;
     private final DurableConfig config;
+    private final BooleanSupplier tryStartCheckpointProcessing;
+    private final Runnable finishCheckpointProcessing;
     private String checkpointToken;
 
     CheckpointManager(
             DurableConfig config,
             String durableExecutionArn,
             String checkpointToken,
             Consumer<List<Operation>> callback) {
+        this(config, durableExecutionArn, checkpointToken, callback, () -> true, () -> {});
+    }
+
+    CheckpointManager(
+            DurableConfig config,
+            String durableExecutionArn,
+            String checkpointToken,
+            Consumer<List<Operation>> callback,
+            BooleanSupplier tryStartCheckpointProcessing,
+            Runnable finishCheckpointProcessing) {
         this.config = config;
         this.durableExecutionArn = durableExecutionArn;
         this.callback = callback;
         this.checkpointToken = checkpointToken;
+        this.tryStartCheckpointProcessing = tryStartCheckpointProcessing;
+        this.finishCheckpointProcessing = finishCheckpointProcessing;
         this.checkpointApiRequestDelayedBatcher = new ApiRequestDelayedBatcher<>(
                 MAX_ITEM_COUNT, MAX_BATCH_SIZE_BYTES, CheckpointManager::estimateSize, this::checkpointBatch);
     }
@@ -191,6 +206,19 @@ private void checkpointBatch(List<OperationUpdate> updates) {
                 return;
             }
 
+            // Starting the backend request is coordinated with the last-thread suspension decision. Once suspension
+            // wins that race, no later poll/checkpoint may advance backend state behind the PENDING response.
+            if (!tryStartCheckpointProcessing.getAsBoolean()) {
+                if (!request.isEmpty()) {
+                    logger.error(
+                            "Checkpoint invariant violation: skipping {} operation updates because execution has already"
+                                    + " completed",
+                            request.size());
+                }
+                logger.debug("Skipping checkpoint API call because execution has already completed");
+                return;
+            }
+
             var startTime = System.nanoTime();
             logger.debug("Calling durable checkpoint API with {} updates: {}", updates.size(), request);
             try {
@@ -228,6 +256,8 @@ private void checkpointBatch(List<OperationUpdate> updates) {
                 }
             } catch (AwsServiceException e) {
                 throw DurableApiErrorClassifier.classifyException(e);
+            } finally {
+                finishCheckpointProcessing.run();
             }
         }
     }

sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java

@@ -71,6 +71,8 @@ public class ExecutionManager implements SafeCloseable {
     private final Set<String> activeThreads = Collections.synchronizedSet(new HashSet<>());
     private static final ThreadLocal<ThreadContext> currentThreadContext = new ThreadLocal<>();
     private final CompletableFuture<Void> executionExceptionFuture = new CompletableFuture<>();
+    // Guarded by activeThreads so starting a checkpoint request is atomic with the last-thread suspension decision.
+    private int checkpointRequestsInFlight;
 
     // ===== Checkpoint Batching =====
     private final CheckpointManager checkpointManager;
@@ -85,8 +87,13 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte
                 input.updatedOperationIds() != null ? Set.copyOf(input.updatedOperationIds()) : Collections.emptySet();
 
         // Create checkpoint batcher for internal coordination
-        this.checkpointManager =
-                new CheckpointManager(config, durableExecutionArn, input.checkpointToken(), this::onCheckpointComplete);
+        this.checkpointManager = new CheckpointManager(
+                config,
+                durableExecutionArn,
+                input.checkpointToken(),
+                this::onCheckpointComplete,
+                this::tryStartCheckpointProcessing,
+                this::finishCheckpointProcessing);
 
         this.operationStorage = checkpointManager.fetchAllPages(input.initialExecutionState()).stream()
                 .collect(Collectors.toConcurrentMap(Operation::id, op -> op));
@@ -187,20 +194,23 @@ public void registerOperation(BaseDurableOperation operation) {
 
     // ===== Checkpoint Completion Handler =====
     /** Called by CheckpointManager when a checkpoint completes. Updates operationStorage and notify operations . */
-    private void onCheckpointComplete(List<Operation> newOperations) {
+    void onCheckpointComplete(List<Operation> newOperations) {
         var updatedOperations = new ArrayList<Operation>();
         newOperations.forEach(op -> {
             // Detect a status change against the previously stored operation
             var previous = operationStorage.get(op.id());
             if (previous == null || previous.status() != op.status()) {
                 updatedOperations.add(op);
             }
-            // Update operation storage
-            operationStorage.put(op.id(), op);
-            // call registered operation's onCheckpointComplete method for completed operations
-            registeredOperations.computeIfPresent(op.id(), (id, operation) -> {
-                operation.onCheckpointComplete(op);
-                return operation;
+            // Publish the updated state and notify its waiter atomically. Otherwise, a waiter can observe the terminal
+            // state before its completion future is completed and attempt to suspend with no pending operations.
+            registeredOperations.compute(op.id(), (id, registeredOperation) -> {
+                if (registeredOperation == null) {
+                    operationStorage.put(op.id(), op);
+                } else {
+                    registeredOperation.processCheckpointUpdate(op, () -> operationStorage.put(op.id(), op));
+                }
+                return registeredOperation;
             });
         });
 
@@ -301,29 +311,57 @@ public void registerActiveThread(String threadId) {
      * @param threadId the thread ID to deregister
      */
     public void deregisterActiveThread(String threadId) {
-        // Skip if already suspended
-        if (executionExceptionFuture.isDone()) {
-            return;
-        }
-
         // Add synchronized block to avoid remove then check race condition and make sure that
         // the suspendExecution is called only once
         synchronized (activeThreads) {
+            // Skip if already suspended
+            if (executionExceptionFuture.isDone()) {
+                return;
+            }
+
             boolean removed = activeThreads.remove(threadId);
             if (removed) {
                 logger.trace("Deregistered thread '{}' Active threads: {}", threadId, activeThreads.size());
             } else {
                 logger.warn("Thread '{}' not active, cannot deregister", threadId);
             }
 
-            if (activeThreads.isEmpty()) {
+            if (shouldSuspendExecution()) {
                 logger.info("No active threads remaining - suspending execution");
                 preSuspendCheck();
                 suspendExecution();
             }
         }
     }
 
+    boolean tryStartCheckpointProcessing() {
+        synchronized (activeThreads) {
+            if (executionExceptionFuture.isDone()) {
+                return false;
+            }
+            checkpointRequestsInFlight++;
+            return true;
+        }
+    }
+
+    void finishCheckpointProcessing() {
+        synchronized (activeThreads) {
+            if (checkpointRequestsInFlight == 0) {
+                throw new IllegalStateException("No checkpoint request is in flight");
+            }
+            checkpointRequestsInFlight--;
+            if (shouldSuspendExecution()) {
+                logger.info("Checkpoint processing completed with no active threads - suspending execution");
+                preSuspendCheck();
+                signalSuspension();
+            }
+        }
+    }
+
+    private boolean shouldSuspendExecution() {
+        return activeThreads.isEmpty() && checkpointRequestsInFlight == 0 && !executionExceptionFuture.isDone();
+    }
+
     private void preSuspendCheck() {
         var hasAnyPendingOperation = operationStorage.values().stream().anyMatch(o -> switch (o.type()) {
             case STEP -> o.status() == OperationStatus.PENDING;
@@ -428,10 +466,14 @@ public void terminateExecution(UnrecoverableDurableExecutionException exception)
 
     /** Suspends the execution by completing the execution exception future with a {@link SuspendExecutionException}. */
     public void suspendExecution() {
+        throw signalSuspension();
+    }
+
+    private SuspendExecutionException signalSuspension() {
         var ex = new SuspendExecutionException();
         stopAllOperations(ex);
         executionExceptionFuture.completeExceptionally(ex);
-        throw ex;
+        return ex;
     }
 
     /**

sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java

@@ -234,7 +234,7 @@ protected Operation waitForOperationCompletion() {
         // is between `isOperationCompleted` and `thenRun`.
         // If this operation is a branch/iteration of a ConcurrencyOperation (map or parallel), the branches/iterations
         // must be completed sequentially to avoid race conditions.
-        synchronized (parentOperation == null ? completionFuture : parentOperation.completionFuture) {
+        synchronized (completionLock()) {
             if (!isOperationCompleted()) {
                 // Add a completion stage to completionFuture so that when the completionFuture is completed,
                 // it will register the current Context thread synchronously to make sure it is always registered
@@ -396,6 +396,19 @@ public void onCheckpointComplete(Operation operation) {
         }
     }
 
+    /**
+     * Publishes a checkpoint update and notifies this operation under the same lock used by operation waiters.
+     *
+     * @param operation the updated operation state
+     * @param publishUpdate publishes the operation to execution state
+     */
+    public final void processCheckpointUpdate(Operation operation, Runnable publishUpdate) {
+        synchronized (completionLock()) {
+            publishUpdate.run();
+            onCheckpointComplete(operation);
+        }
+    }
+
     /** Marks the operation as already completed (in replay). */
     protected void markAlreadyCompleted() {
         // When the operation is already completed in a replay, we complete completionFuture immediately
@@ -407,7 +420,7 @@ protected void markAlreadyCompleted() {
     private void markCompletionFutureCompleted() {
         // It's important that we synchronize access to the future, otherwise the processing could happen
         // on someone else's thread and cause a race condition.
-        synchronized (parentOperation == null ? completionFuture : parentOperation.completionFuture) {
+        synchronized (completionLock()) {
             // Completing the future here will also run any other completion stages that have been attached
             // to the future. In our case, other contexts may have attached a function to reactivate themselves,
             // so they will definitely have a chance to reactivate before we finish completing and deactivating
@@ -416,6 +429,10 @@ private void markCompletionFutureCompleted() {
         }
     }
 
+    private CompletableFuture<BaseDurableOperation> completionLock() {
+        return parentOperation == null ? completionFuture : parentOperation.completionFuture;
+    }
+
     /**
      * Terminates the execution with the given exception.
      *

sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java

@@ -75,6 +75,23 @@ void checkpoint_sendsUpdateAndReturnsCompletedFuture() throws Exception {
         assertTrue(future.isDone());
     }
 
+    @Test
+    void checkpoint_skipsApiCallWhenExecutionAlreadyCompleted() throws Exception {
+        var finishCheckpointProcessing = mock(Runnable.class);
+        var guardedBatcher = new CheckpointManager(
+                config, "arn:test", "token-1", callbackOperations::addAll, () -> false, finishCheckpointProcessing);
+        var update = OperationUpdate.builder()
+                .id("op-1")
+                .type(OperationType.STEP)
+                .action(OperationAction.START)
+                .build();
+
+        guardedBatcher.checkpoint(update).get(200, TimeUnit.MILLISECONDS);
+
+        verifyNoInteractions(client);
+        verify(finishCheckpointProcessing, never()).run();
+    }
+
     @Test
     void pollForUpdate_completesWhenOperationReturned() throws Exception {
         var operation = Operation.builder()