From b23ad7132c6afaac000ccd30628b3e6cb102b4f2 Mon Sep 17 00:00:00 2001 From: Harini L Date: Tue, 13 Oct 2026 14:18:27 +0530 Subject: [PATCH] fix(sdk): report invalid suspension without pending operations --- .../durable/ParallelIntegrationTest.java | 26 ++++ .../lambda/durable/PluginIntegrationTest.java | 9 +- .../testing/LocalDurableTestRunnerTest.java | 38 ++++++ .../IllegalDurableOperationException.java | 14 +- .../durable/execution/CheckpointManager.java | 8 ++ .../durable/execution/DurableExecutor.java | 15 +++ .../durable/execution/ExecutionManager.java | 23 +++- .../execution/DurableExecutionTest.java | 122 ++++++++++++++++++ 8 files changed, 246 insertions(+), 9 deletions(-) diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java index 7a1c45f03..8ec7e2057 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/ParallelIntegrationTest.java +++ b/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>(); + 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")); + } } diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java index f8e3bceb6..d8b42c050 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/PluginIntegrationTest.java +++ b/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"); diff --git a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java b/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java index 36f1bbced..3060f4651 100644 --- a/sdk-testing/src/test/java/software/amazon/lambda/durable/testing/LocalDurableTestRunnerTest.java +++ b/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")); + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/exception/IllegalDurableOperationException.java b/sdk/src/main/java/software/amazon/lambda/durable/exception/IllegalDurableOperationException.java index 875a1c063..1ab3686af 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/exception/IllegalDurableOperationException.java +++ b/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); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java index 34f44134c..23c557497 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java +++ b/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, @@ -229,6 +234,9 @@ private void checkpointBatch(List 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()); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java index d8db91326..0c0d97500 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java @@ -127,6 +127,21 @@ public static DurableExecutionOutput execute( // return PENDING if it's SuspendExecutionException if (cause instanceof SuspendExecutionException) { + var invalidSuspensionException = executionManager.getInvalidSuspensionException(); + if (invalidSuspensionException != null) { + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.RETRYING, + invalidSuspensionException, + pluginExecutionInput.get(), + null); + throw invalidSuspensionException; + } + fireOnInvocationEnd( pluginRunner, executionManager, diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java index 4c7feeda4..0fa642129 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java @@ -77,6 +77,7 @@ public class ExecutionManager implements SafeCloseable { private final CompletableFuture executionExceptionFuture = new CompletableFuture<>(); // Guarded by activeThreads so starting a checkpoint request is atomic with the last-thread suspension decision. private int checkpointRequestsInFlight; + private volatile IllegalDurableOperationException invalidSuspensionException; /** * Per-wait state used to coordinate the caller thread with the future completion callback. @@ -556,8 +557,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(), ""), 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); } } @@ -670,6 +681,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)); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java index e0773aaeb..e5c442a59 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java @@ -8,6 +8,9 @@ import static org.junit.jupiter.api.Assertions.assertNull; 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 +21,30 @@ 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.InvocationStatus; class DurableExecutionTest { @@ -448,6 +459,117 @@ 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); + } + + @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); + } + + private static class RecordingPlugin implements DurableExecutionPlugin { + private InvocationStatus invocationStatus; + private Throwable invocationError; + + @Override + public void onInvocationEnd(InvocationEndInfo info) { + this.invocationStatus = info.invocationStatus(); + this.invocationError = info.executionError(); + } + } + private Operation executionOp() { return Operation.builder() .id(EXECUTION_OP_ID)