fix(sdk): report invalid suspension without pending operations

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java

@@ -19,6 +19,7 @@
 import software.amazon.lambda.durable.config.NestingType;
 import software.amazon.lambda.durable.config.ParallelConfig;
 import software.amazon.lambda.durable.config.WaitForConditionConfig;
+import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
 import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.model.ParallelResult;
@@ -1369,4 +1370,29 @@ void testParallelWithFirstSuccessful_earlyTermination(NestingType nestingType, i
             assertTrue(result.getHistoryEvents().size() <= events);
         }
     }
+
+    @Test
+    void testInvalidSuspensionWhenWaitingOnSkippedBranch() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
+            var config = ParallelConfig.builder()
+                    .completionConfig(CompletionConfig.firstSuccessful())
+                    .maxConcurrency(1)
+                    .build();
+            var futures = new ArrayList<DurableFuture<String>>();
+            var parallel = context.parallel("first-success-parallel", config);
+
+            try (parallel) {
+                futures.add(parallel.branch("branch-1", String.class, ctx -> "success-1"));
+                futures.add(parallel.branch("branch-2", String.class, ctx -> "skipped-2"));
+            }
+
+            parallel.get();
+            return futures.get(1).get();
+        });
+
+        var ex = assertThrows(IllegalDurableOperationException.class, () -> runner.run("test"));
+        assertTrue(ex.isRetryable());
+        assertTrue(ex.getMessage().contains("Invalid suspension. No operation is pending."));
+        assertTrue(ex.getMessage().contains("branch-2"));
+    }
 }

sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java

@@ -981,9 +981,12 @@ void plugin_userFunctionEnd_unwrapsCompletionExceptionForSuspension() {
 
         var runner = LocalDurableTestRunner.create(
                 String.class,
-                (input, context) -> context.step("wrapped-suspension", String.class, step -> {
-                    throw new CompletionException(suspension);
-                }),
+                (input, context) -> {
+                    context.waitAsync("hold", Duration.ofMinutes(5));
+                    return context.step("wrapped-suspension", String.class, step -> {
+                        throw new CompletionException(suspension);
+                    });
+                },
                 config);
 
         var result = runner.run("input");

sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java

@@ -11,9 +11,15 @@
 import java.util.Collections;
 import java.util.List;
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.TypeToken;
+import software.amazon.lambda.durable.context.DurableContextImpl;
+import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
 import software.amazon.lambda.durable.model.ExecutionStatus;
+import software.amazon.lambda.durable.model.OperationIdentifier;
+import software.amazon.lambda.durable.model.OperationSubType;
+import software.amazon.lambda.durable.operation.BaseDurableOperation;
 import software.amazon.lambda.durable.plugin.DurableExecutionPlugin;
 import software.amazon.lambda.durable.plugin.InvocationInfo;
 
@@ -114,4 +120,36 @@ public void onInvocationStart(InvocationInfo info) {
         assertNotNull(executionStartTimes.get(0));
         assertEquals(executionStartTimes.get(0), executionStartTimes.get(1));
     }
+
+    @Test
+    void testInvalidSuspensionThrowsIllegalDurableOperationException() {
+        var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> {
+            var durableContext = (DurableContextImpl) ctx;
+            class WaitingOperation extends BaseDurableOperation {
+                WaitingOperation() {
+                    super(
+                            OperationIdentifier.of("unregistered-op", "unregistered-op", OperationSubType.STEP),
+                            durableContext,
+                            null);
+                }
+
+                @Override
+                protected void start() {}
+
+                @Override
+                protected void replay(Operation existing) {}
+
+                void await() {
+                    waitForOperationCompletion();
+                }
+            }
+            new WaitingOperation().await();
+            return "done";
+        });
+
+        var ex = assertThrows(IllegalDurableOperationException.class, () -> runner.run("input"));
+        assertTrue(ex.isRetryable());
+        assertTrue(ex.getMessage().contains("Invalid suspension. No operation is pending."));
+        assertTrue(ex.getMessage().contains("unregistered-op"));
+    }
 }

sdk/src/main/java/software/amazon/lambda/durable/exception/IllegalDurableOperationException.java

@@ -7,9 +7,15 @@
 /** An illegal operation is detected. The execution will be immediately terminated. */
 public class IllegalDurableOperationException extends UnrecoverableDurableExecutionException {
     public IllegalDurableOperationException(String message) {
-        super(ErrorObject.builder()
-                .errorType(IllegalDurableOperationException.class.getName())
-                .errorMessage(message)
-                .build());
+        this(message, false);
+    }
+
+    public IllegalDurableOperationException(String message, boolean retryable) {
+        super(
+                ErrorObject.builder()
+                        .errorType(IllegalDurableOperationException.class.getName())
+                        .errorMessage(message)
+                        .build(),
+                retryable);
     }
 }

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

@@ -44,6 +44,11 @@ class CheckpointManager {
     private final BooleanSupplier tryStartCheckpointProcessing;
     private final Runnable finishCheckpointProcessing;
     private String checkpointToken;
+    private volatile boolean checkpointTokenRevoked;
+
+    boolean isCheckpointTokenRevoked() {
+        return checkpointTokenRevoked;
+    }
 
     CheckpointManager(
             DurableConfig config,
@@ -234,6 +239,9 @@ private void checkpointBatch(List<OperationUpdate> updates) {
 
                 // Notify callback of completion
                 checkpointToken = response.checkpointToken();
+                if (checkpointToken == null) {
+                    checkpointTokenRevoked = true;
+                }
                 if (response.newExecutionState() != null) {
                     // fetch all pages of operations
                     var operations = fetchAllPages(response.newExecutionState());

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

@@ -181,6 +181,10 @@ private DurableExecutionOutput finishInvocation(O value, Throwable failure) {
         }
 
         private DurableExecutionOutput finishFailure(Throwable cause) {
+            if (cause instanceof SuspendExecutionException) {
+                var invalidSuspension = manager.getInvalidSuspensionException();
+                if (invalidSuspension != null) cause = invalidSuspension;
+            }
             var status = failureStatus(cause);
             if (status == InvocationStatus.FAILED) return finishTerminalFailure(cause);
             try {

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

@@ -80,6 +80,7 @@ public class ExecutionManager implements SafeCloseable {
     private final CompletableFuture<Void> executionExceptionFuture = new CompletableFuture<>();
     // Requests and polling continuations; guarded by activeThreads so admission is atomic with last-thread suspension.
     private int checkpointRequestsInFlight;
+    private volatile IllegalDurableOperationException invalidSuspensionException;
     // Guarded by activeThreads; an observation future can be cancelled independently of its actual task.
     private final Map<Object, BaseDurableOperation> checkpointContinuations = new HashMap<>();
     private boolean closing;
@@ -691,8 +692,18 @@ private void preSuspendCheck() {
             default -> false;
         });
 
-        if (!hasAnyPendingOperation) {
-            logger.warn("Invalid suspension. No operation is pending");
+        if (!hasAnyPendingOperation && !isCheckpointTokenRevoked()) {
+            var waitingOperations = registeredOperations.values().stream()
+                    .filter(op -> !op.getCompletionFuture().isDone())
+                    .map(op -> String.format(
+                            "%s (id=%s)", Objects.toString(op.getName(), "<unnamed>"), op.getOperationId()))
+                    .distinct()
+                    .toList();
+
+            var errorMessage = "Invalid suspension. No operation is pending. Suspended threads were waiting for: "
+                    + waitingOperations;
+            logger.error(errorMessage);
+            this.invalidSuspensionException = new IllegalDurableOperationException(errorMessage, true);
         }
     }
 
@@ -829,6 +840,14 @@ public boolean isExecutionCompletedExceptionally() {
         return executionExceptionFuture.isCompletedExceptionally();
     }
 
+    public IllegalDurableOperationException getInvalidSuspensionException() {
+        return invalidSuspensionException;
+    }
+
+    public boolean isCheckpointTokenRevoked() {
+        return checkpointManager.isCheckpointTokenRevoked();
+    }
+
     private void stopAllOperations(Throwable cause) {
         registeredOperations.values().forEach(op -> op.getCompletionFuture().completeExceptionally(cause));
     }

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

@@ -6,8 +6,12 @@
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 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.mock;
+import static org.mockito.Mockito.when;
 import static software.amazon.lambda.durable.TypeToken.get;
 
 import java.time.Instant;
@@ -18,22 +22,31 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse;
 import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
 import software.amazon.awssdk.services.lambda.model.ErrorObject;
 import software.amazon.awssdk.services.lambda.model.ExecutionDetails;
 import software.amazon.awssdk.services.lambda.model.Operation;
+import software.amazon.awssdk.services.lambda.model.OperationAction;
 import software.amazon.awssdk.services.lambda.model.OperationStatus;
 import software.amazon.awssdk.services.lambda.model.OperationType;
+import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.awssdk.services.lambda.model.StepDetails;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.TestUtils;
+import software.amazon.lambda.durable.client.DurableExecutionClient;
 import software.amazon.lambda.durable.context.DurableContextImpl;
+import software.amazon.lambda.durable.exception.IllegalDurableOperationException;
 import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.model.DurableExecutionInput;
 import software.amazon.lambda.durable.model.ExecutionStatus;
 import software.amazon.lambda.durable.model.OperationIdentifier;
 import software.amazon.lambda.durable.model.OperationSubType;
 import software.amazon.lambda.durable.operation.BaseDurableOperation;
+import software.amazon.lambda.durable.plugin.DurableExecutionPlugin;
+import software.amazon.lambda.durable.plugin.InvocationEndInfo;
+import software.amazon.lambda.durable.plugin.InvocationInfo;
+import software.amazon.lambda.durable.plugin.InvocationStatus;
 
 class DurableExecutionTest {
 
@@ -458,6 +471,133 @@ void testExecutorNotShutdownAfterMultipleHandlerInvocations() {
         assertTrue(output2.result().contains("Result 2: test-input-2"));
     }
 
+    @Test
+    void invalidSuspensionThrowsIllegalDurableOperationExceptionAndNotifiesPluginsAsRetrying() {
+        var input = new DurableExecutionInput(
+                EXECUTION_ARN,
+                "token1",
+                CheckpointUpdatedExecutionState.builder()
+                        .operations(List.of(executionOp()))
+                        .build());
+
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(TestUtils.createMockClient())
+                .withPlugins(plugin)
+                .build();
+
+        var exception = assertThrows(
+                IllegalDurableOperationException.class,
+                () -> DurableExecutor.execute(
+                        input,
+                        null,
+                        get(String.class),
+                        (userInput, ctx) -> {
+                            var durableContext = (DurableContextImpl) ctx;
+                            class WaitingOperation extends BaseDurableOperation {
+                                WaitingOperation() {
+                                    super(
+                                            OperationIdentifier.of(
+                                                    "my-waiting-op", "my-waiting-op", OperationSubType.STEP),
+                                            durableContext,
+                                            null);
+                                }
+
+                                @Override
+                                protected void start() {}
+
+                                @Override
+                                protected void replay(Operation existing) {}
+
+                                void await() {
+                                    waitForOperationCompletion();
+                                }
+                            }
+                            var op = new WaitingOperation();
+                            op.await();
+                            return "should-not-reach";
+                        },
+                        config));
+
+        assertTrue(exception.isRetryable());
+        assertTrue(exception.getMessage().contains("Invalid suspension. No operation is pending."));
+        assertTrue(exception.getMessage().contains("my-waiting-op"));
+        assertEquals(InvocationStatus.RETRYING, plugin.invocationStatus);
+        assertEquals(exception, plugin.invocationError);
+        assertNotNull(plugin.startThread);
+        assertSame(plugin.startThread, plugin.endThread);
+        assertEquals(1, plugin.endCalls);
+    }
+
+    @Test
+    void revokedCheckpointTokenDoesNotRecordInvalidSuspension() {
+        var mockClient = mock(DurableExecutionClient.class);
+        when(mockClient.checkpoint(any(), any(), any()))
+                .thenReturn(CheckpointDurableExecutionResponse.builder()
+                        .checkpointToken(null)
+                        .build());
+
+        var input = new DurableExecutionInput(
+                EXECUTION_ARN,
+                "token1",
+                CheckpointUpdatedExecutionState.builder()
+                        .operations(List.of(executionOp()))
+                        .build());
+
+        var plugin = new RecordingPlugin();
+        var config = DurableConfig.builder()
+                .withDurableExecutionClient(mockClient)
+                .withPlugins(plugin)
+                .build();
+
+        var output = DurableExecutor.execute(
+                input,
+                null,
+                get(String.class),
+                (userInput, ctx) -> {
+                    var durableContext = (DurableContextImpl) ctx;
+                    durableContext
+                            .getExecutionManager()
+                            .sendOperationUpdate(OperationUpdate.builder()
+                                    .id("dummy")
+                                    .type(OperationType.STEP)
+                                    .action(OperationAction.START)
+                                    .build())
+                            .join();
+                    durableContext.getExecutionManager().suspendExecution();
+                    return "should-not-reach";
+                },
+                config);
+
+        assertEquals(ExecutionStatus.PENDING, output.status());
+        assertEquals(InvocationStatus.PENDING, plugin.invocationStatus);
+        assertNull(plugin.invocationError);
+        assertNotNull(plugin.startThread);
+        assertSame(plugin.startThread, plugin.endThread);
+        assertEquals(1, plugin.endCalls);
+    }
+
+    private static class RecordingPlugin implements DurableExecutionPlugin {
+        private InvocationStatus invocationStatus;
+        private Throwable invocationError;
+        private Thread startThread;
+        private Thread endThread;
+        private int endCalls;
+
+        @Override
+        public void onInvocationStart(InvocationInfo info) {
+            startThread = Thread.currentThread();
+        }
+
+        @Override
+        public void onInvocationEnd(InvocationEndInfo info) {
+            endThread = Thread.currentThread();
+            endCalls++;
+            this.invocationStatus = info.invocationStatus();
+            this.invocationError = info.executionError();
+        }
+    }
+
     private Operation executionOp() {
         return Operation.builder()
                 .id(EXECUTION_OP_ID)