diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java index 90551a7a9..f0391567a 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -139,13 +139,12 @@ public void execute() { validateReplay(existing); if (ExecutionManager.isTerminalStatus(existing.status())) { replayCompletedOperation.set(true); - } else if (getType() == OperationType.STEP || getType() == OperationType.CONTEXT) { - // Non-terminal STEP/CONTEXT operations are being re-executed (user code runs again). - // Fire onOperationStart so the OTel plugin can create a parent span for attempt spans. + } else { + // Non-terminal operations are being replayed. Fire onOperationStart so plugins + // can observe all in-progress operations during replay, including WAIT/INVOKE/CALLBACK + // that are still pending. fireOnOperationStart(existing); } - // WAIT/INVOKE/CALLBACK in non-terminal status just poll — no onOperationStart needed. - // They'll get a continuation span via onOperationEnd when they complete. // Fire onOperationEnd for operations that completed during suspension (between invocations). // The OTel plugin handles the missing onOperationStart by creating a continuation span linked diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java new file mode 100644 index 000000000..463a824ea --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/BaseDurableOperationPluginTest.java @@ -0,0 +1,166 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.operation; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.EnumSource; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.execution.ThreadContext; +import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.OperationInfo; + +/** + * Unit tests verifying that BaseDurableOperation.execute() fires onOperationStart with isReplay=true for all + * non-terminal operations during replay, regardless of operation type. + * + *

This mirrors the Python SDK's TestPluginExecutorOnOperationReplay tests. + */ +class BaseDurableOperationPluginTest { + + private static final String EXECUTION_OP_ID = "exec-123"; + private static final String EXECUTION_ARN = + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/exec-name/" + EXECUTION_OP_ID; + private static final String OPERATION_ID = TestUtils.hashOperationId("1"); + private static final String OPERATION_NAME = "test-op"; + + @Test + void execute_firesOnOperationStart_withIsReplayTrue_forNonTerminalWait() { + var plugin = new RecordingPlugin(); + var waitOp = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .type(OperationType.WAIT) + .subType("Wait") + .status(OperationStatus.STARTED) + .build(); + + var executionManager = createExecutionManager(List.of(waitOp), plugin); + var durableContext = mockDurableContext(executionManager, plugin); + + var operation = new WaitOperation( + OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT), + Duration.ofMinutes(5), + durableContext); + + operation.execute(); + + // onOperationStart should fire with isReplay=true for the non-terminal WAIT + assertEquals(1, plugin.operationStarts.size(), "Should fire exactly one onOperationStart"); + var info = plugin.operationStarts.get(0); + assertEquals(OPERATION_NAME, info.name()); + assertEquals("WAIT", info.type()); + assertTrue(info.isReplay(), "isReplay should be true for a replayed non-terminal wait"); + } + + @ParameterizedTest + @EnumSource( + value = OperationStatus.class, + names = {"SUCCEEDED", "FAILED", "TIMED_OUT", "STOPPED"}) + void execute_doesNotFireOnOperationStart_forTerminalOperation(OperationStatus terminalStatus) { + var plugin = new RecordingPlugin(); + var waitOp = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .type(OperationType.WAIT) + .subType("Wait") + .status(terminalStatus) + .build(); + + var executionManager = createExecutionManager(List.of(waitOp), plugin); + var durableContext = mockDurableContext(executionManager, plugin); + + var operation = new WaitOperation( + OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT), + Duration.ofMinutes(5), + durableContext); + + operation.execute(); + + // onOperationStart should NOT fire for terminal operations during replay + assertTrue( + plugin.operationStarts.isEmpty(), + "Should not fire onOperationStart for terminal " + terminalStatus + " operation"); + } + + @Test + void execute_firesOnOperationStart_withIsReplayFalse_forFirstExecution() { + var plugin = new RecordingPlugin(); + // No existing operations — first execution + var executionManager = createExecutionManager(List.of(), plugin); + var durableContext = mockDurableContext(executionManager, plugin); + + var operation = new WaitOperation( + OperationIdentifier.of(OPERATION_ID, OPERATION_NAME, OperationSubType.WAIT), + Duration.ofMinutes(5), + durableContext); + + operation.execute(); + + // onOperationStart should fire with isReplay=false on first execution + assertEquals(1, plugin.operationStarts.size(), "Should fire exactly one onOperationStart"); + var info = plugin.operationStarts.get(0); + assertEquals(OPERATION_NAME, info.name()); + assertFalse(info.isReplay(), "isReplay should be false on first execution"); + } + + // ─── Helpers ───────────────────────────────────────────────────────── + + private ExecutionManager createExecutionManager(List additionalOps, RecordingPlugin plugin) { + var client = TestUtils.createMockClient(); + var operations = new ArrayList(); + operations.add(Operation.builder() + .id(EXECUTION_OP_ID) + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .build()); + operations.addAll(additionalOps); + var initialState = + CheckpointUpdatedExecutionState.builder().operations(operations).build(); + var config = DurableConfig.builder() + .withDurableExecutionClient(client) + .withPlugins(plugin) + .build(); + var executionManager = new ExecutionManager( + new DurableExecutionInput(EXECUTION_ARN, "test-token", initialState), config, null); + executionManager.setCurrentThreadContext(new ThreadContext("Root", ThreadType.CONTEXT)); + return executionManager; + } + + private DurableContextImpl mockDurableContext(ExecutionManager executionManager, RecordingPlugin plugin) { + var durableContext = mock(DurableContextImpl.class); + when(durableContext.getExecutionManager()).thenReturn(executionManager); + when(durableContext.getDurableConfig()) + .thenReturn(DurableConfig.builder().withPlugins(plugin).build()); + return durableContext; + } + + /** Plugin that records onOperationStart calls. */ + private static class RecordingPlugin implements DurableExecutionPlugin { + final List operationStarts = Collections.synchronizedList(new ArrayList<>()); + + @Override + public void onOperationStart(OperationInfo info) { + operationStarts.add(info); + } + } +}