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 5bd794b9c..648488034 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 @@ -4,11 +4,15 @@ import static org.junit.jupiter.api.Assertions.*; +import java.io.IOException; import java.time.Duration; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.OperationStatus; @@ -19,7 +23,10 @@ import software.amazon.lambda.durable.model.WaitForConditionResult; import software.amazon.lambda.durable.plugin.*; import software.amazon.lambda.durable.retry.RetryStrategies; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.testing.LocalDurableTestRunner; +import software.amazon.lambda.durable.util.ExceptionHelper; /** Integration tests verifying plugin hooks fire correctly during durable execution lifecycle. */ class PluginIntegrationTest { @@ -122,6 +129,167 @@ void plugin_receivesInvocationEnd_withFailedStatus_onError() { assertNotNull(plugin.invocationEnds.get(0).executionError()); } + @Test + void plugin_invocationHooks_carryExecutionInputAndResult_onSuccess() { + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder().withPlugins(plugin).build(); + + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> context.step("greet", String.class, stepCtx -> "Hello " + input), + config); + + var result = runner.runUntilComplete("World"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + // The deserialized handler input reaches the start hook ... + assertEquals("World", plugin.invocationStarts.get(0).executionInput()); + // ... and the end hook, alongside the value the handler returned. + var end = plugin.invocationEnds.get(0); + assertEquals("World", end.executionInput()); + assertEquals("Hello World", end.executionResult()); + } + + @Test + void plugin_invocationEnd_omitsExecutionResult_onFailure() { + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder().withPlugins(plugin).build(); + + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> context.step( + "failing", + String.class, + stepCtx -> { + throw new RuntimeException("boom"); + }, + StepConfig.builder() + .retryStrategy(RetryStrategies.Presets.NO_RETRY) + .build()), + config); + + var result = runner.run("input"); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + + var end = plugin.invocationEnds.get(0); + assertEquals("input", end.executionInput()); + assertNull(end.executionResult(), "a failed execution produced no result"); + } + + @Test + void plugin_invocationEnd_omitsExecutionResult_onSuspension() { + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder().withPlugins(plugin).build(); + + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + context.wait("pause", Duration.ofMinutes(5)); + return "complete"; + }, + config); + + var result = runner.run("input"); + + assertEquals(ExecutionStatus.PENDING, result.getStatus()); + + var end = plugin.invocationEnds.get(0); + assertEquals(InvocationStatus.PENDING, end.invocationStatus()); + assertEquals("input", end.executionInput()); + assertNull(end.executionResult(), "a suspended invocation has not produced a result yet"); + } + + @Test + void plugin_executionInput_isDeserializedOnce_andSharedWithHandler() { + var serDes = new CountingSerDes(); + var plugin = new RecordingPlugin(); + var config = + DurableConfig.builder().withPlugins(plugin).withSerDes(serDes).build(); + var handlerInput = new AtomicReference(); + + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + handlerInput.set(input); + return context.step("greet", String.class, stepCtx -> "Hello " + input); + }, + config); + + var result = runner.runUntilComplete("World"); + + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + // The handler input is deserialized once and that same instance reaches the hooks — a second + // deserialization would double the cost and re-run side effects in a stateful SerDes. + assertEquals(1, serDes.inputDeserializations("World")); + assertSame(handlerInput.get(), plugin.invocationStarts.get(0).executionInput()); + assertSame(handlerInput.get(), plugin.invocationEnds.get(0).executionInput()); + } + + /** SerDes that counts how many times each payload is deserialized. */ + static class CountingSerDes implements SerDes { + private final JacksonSerDes delegate = new JacksonSerDes(); + private final Map deserializations = new ConcurrentHashMap<>(); + + @Override + public String serialize(Object value) { + return delegate.serialize(value); + } + + @Override + public T deserialize(String data, TypeToken typeToken) { + deserializations + .computeIfAbsent(data, ignored -> new AtomicInteger()) + .incrementAndGet(); + return delegate.deserialize(data, typeToken); + } + + int inputDeserializations(String value) { + var counter = deserializations.get(delegate.serialize(value)); + return counter == null ? 0 : counter.get(); + } + } + + @Test + void plugin_hooksStayPaired_whenSerDesSneakyThrowsCheckedException() { + var plugin = new RecordingPlugin(); + var config = DurableConfig.builder() + .withPlugins(plugin) + .withSerDes(new SneakyThrowingSerDes()) + .build(); + + var runner = LocalDurableTestRunner.create(String.class, (input, context) -> "unreachable", config); + + var result = runner.run("input"); + + assertEquals(ExecutionStatus.FAILED, result.getStatus()); + // SerDes.deserialize declares no checked exceptions, so an implementation can sneaky-throw one. The start + // hook must still fire, otherwise a plugin keying state on it sees an end with no start. + assertEquals(1, plugin.invocationStarts.size(), "onInvocationStart must fire before the failure surfaces"); + assertNull(plugin.invocationStarts.get(0).executionInput(), "input could not be deserialized"); + assertEquals(1, plugin.invocationEnds.size()); + assertEquals(InvocationStatus.FAILED, plugin.invocationEnds.get(0).invocationStatus()); + } + + /** + * SerDes that sneaky-throws a checked exception from deserialize, as DurableInputOutputSerDes does on IO errors. + */ + static class SneakyThrowingSerDes implements SerDes { + private final JacksonSerDes delegate = new JacksonSerDes(); + + @Override + public String serialize(Object value) { + return delegate.serialize(value); + } + + @Override + public T deserialize(String data, TypeToken typeToken) { + ExceptionHelper.sneakyThrow(new IOException("sneaky checked failure")); + return null; + } + } + // ─── Operation-level hooks ─────────────────────────────────────────── @Test 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 6ab6bdbc3..649c7a600 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 @@ -7,6 +7,7 @@ import java.nio.charset.StandardCharsets; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CompletionException; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.BiFunction; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -60,10 +61,28 @@ public static DurableExecutionOutput execute( var executionArn = input.durableExecutionArn(); executionManager.registerActiveThread(null); + // Captured for onInvocationEnd, which runs outside the handler thread below. + var pluginExecutionInput = new AtomicReference<>(); var handlerFuture = CompletableFuture.supplyAsync( () -> { executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT)); + // Deserialize once and share the value with the plugin hooks and the handler below. A second + // deserialization would double the cost, hand plugins a different object than the handler, and + // re-run any side effects in a stateful custom SerDes. A failure is captured rather than thrown + // so onInvocationStart still fires before it surfaces, keeping the start/end hooks paired. + // SerDes is a public extension point whose deserialize declares no checked exceptions, so an + // implementation may sneaky-throw one; capture every Throwable and rethrow it unchanged. + I userInput = null; + Throwable inputFailure = null; + try { + userInput = extractUserInput( + executionManager.getExecutionOperation(), config.getSerDes(), inputType); + } catch (Throwable t) { + inputFailure = t; + } + pluginExecutionInput.set(userInput); + // onInvocationStart runs on the user thread so plugins can // inject ThreadLocal objects, update MDC, etc. // executionStartTime comes from the initial EXECUTION operation in the first backend event. @@ -71,10 +90,12 @@ public static DurableExecutionOutput execute( requestId, executionArn, isFirstInvocation, - executionManager.getExecutionOperation().startTimestamp())); + executionManager.getExecutionOperation().startTimestamp(), + userInput)); + if (inputFailure != null) { + ExceptionHelper.sneakyThrow(inputFailure); + } - var userInput = extractUserInput( - executionManager.getExecutionOperation(), config.getSerDes(), inputType); var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext); DurableContextImpl.setCurrentContext(context); // use a try-with-resources to clear logger properties @@ -103,6 +124,8 @@ public static DurableExecutionOutput execute( executionArn, isFirstInvocation, InvocationStatus.PENDING, + null, + pluginExecutionInput.get(), null); return DurableExecutionOutput.pending(); } @@ -119,7 +142,9 @@ public static DurableExecutionOutput execute( executionArn, isFirstInvocation, InvocationStatus.RETRYING, - cause); + cause, + pluginExecutionInput.get(), + null); throw unrecoverableDurableExecutionException; } @@ -131,7 +156,9 @@ public static DurableExecutionOutput execute( executionArn, isFirstInvocation, InvocationStatus.FAILED, - cause); + cause, + pluginExecutionInput.get(), + null); return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes())); } // user handler complete successfully @@ -145,7 +172,9 @@ public static DurableExecutionOutput execute( executionArn, isFirstInvocation, InvocationStatus.SUCCEEDED, - null); + null, + pluginExecutionInput.get(), + result); return output; }) .join(); @@ -163,8 +192,11 @@ private static void fireOnInvocationEnd( String executionArn, boolean isFirstInvocation, InvocationStatus status, - Throwable error) { - pluginRunner.onInvocationEnd(new InvocationEndInfo(requestId, executionArn, isFirstInvocation, status, error)); + Throwable error, + Object executionInput, + Object executionResult) { + pluginRunner.onInvocationEnd(new InvocationEndInfo( + requestId, executionArn, isFirstInvocation, status, error, executionInput, executionResult)); } private static String handleLargePayload(ExecutionManager executionManager, String outputPayload) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationEndInfo.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationEndInfo.java index dd49f0571..c187700f6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationEndInfo.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationEndInfo.java @@ -12,10 +12,52 @@ * @param isFirstInvocation true if this is the first invocation of the execution * @param invocationStatus the invocation outcome (SUCCEEDED, FAILED, or PENDING) * @param executionError non-null if the execution failed; this component is experimental + * @param executionInput the deserialized execution input passed to the user handler, or null when no plugins are + * registered or the input could not be deserialized; this component is experimental + * @param executionResult the value the user handler returned, or null unless the invocation completed the execution + * successfully; this component is experimental */ public record InvocationEndInfo( String requestId, String durableExecutionArn, boolean isFirstInvocation, InvocationStatus invocationStatus, - @Experimental Throwable executionError) {} + @Experimental Throwable executionError, + @Experimental Object executionInput, + @Experimental Object executionResult) { + + /** + * Creates invocation-end information without the execution input or result. + * + *

Retained so callers written before {@code executionInput} and {@code executionResult} were added keep + * compiling; both resolve to null. + * + * @param requestId the Lambda request ID for this invocation + * @param durableExecutionArn the durable execution ARN + * @param isFirstInvocation true if this is the first invocation of the execution + * @param invocationStatus the invocation outcome + * @param executionError non-null if the execution failed + */ + public InvocationEndInfo( + String requestId, + String durableExecutionArn, + boolean isFirstInvocation, + InvocationStatus invocationStatus, + Throwable executionError) { + this(requestId, durableExecutionArn, isFirstInvocation, invocationStatus, executionError, null, null); + } + + /** + * Returns a representation that omits {@code executionInput} and {@code executionResult}. + * + *

The generated representation would render both payloads, so plugins that log this object whole would start + * emitting customer inputs and results, potentially including secrets or personal data. Read the components + * explicitly to record them. + */ + @Override + public String toString() { + return "InvocationEndInfo[requestId=" + requestId + ", durableExecutionArn=" + durableExecutionArn + + ", isFirstInvocation=" + isFirstInvocation + ", invocationStatus=" + invocationStatus + + ", executionError=" + executionError + "]"; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationInfo.java b/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationInfo.java index 89af7e189..c384cf887 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationInfo.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationInfo.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.plugin; import java.time.Instant; +import software.amazon.lambda.durable.annotations.Experimental; /** * Invocation-level information available to plugin hooks. @@ -12,6 +13,42 @@ * @param isFirstInvocation true if this is the first invocation of the execution (not a replay invocation) * @param executionStartTime the start timestamp of the durable execution, taken from the initial EXECUTION operation in * the first event delivered by the backend. Stable across all invocations of the same execution. + * @param executionInput the deserialized execution input passed to the user handler, or null when no plugins are + * registered or the input could not be deserialized; this component is experimental */ public record InvocationInfo( - String requestId, String durableExecutionArn, boolean isFirstInvocation, Instant executionStartTime) {} + String requestId, + String durableExecutionArn, + boolean isFirstInvocation, + Instant executionStartTime, + @Experimental Object executionInput) { + + /** + * Creates invocation information without an execution input. + * + *

Retained so callers written before {@code executionInput} was added keep compiling; {@code executionInput} + * resolves to null. + * + * @param requestId the Lambda request ID for this invocation + * @param durableExecutionArn the durable execution ARN + * @param isFirstInvocation true if this is the first invocation of the execution + * @param executionStartTime the start timestamp of the durable execution + */ + public InvocationInfo( + String requestId, String durableExecutionArn, boolean isFirstInvocation, Instant executionStartTime) { + this(requestId, durableExecutionArn, isFirstInvocation, executionStartTime, null); + } + + /** + * Returns a representation that omits {@code executionInput}. + * + *

The generated representation would render the execution input, so plugins that log this object whole would + * start emitting customer payloads, potentially including secrets or personal data. Read the component explicitly + * to record it. + */ + @Override + public String toString() { + return "InvocationInfo[requestId=" + requestId + ", durableExecutionArn=" + durableExecutionArn + + ", isFirstInvocation=" + isFirstInvocation + ", executionStartTime=" + executionStartTime + "]"; + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java index d34b5a8ca..ed15499c9 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/plugin/PluginRunnerTest.java @@ -133,6 +133,65 @@ void pluginRunner_isImmutable() { assertEquals(List.of("p1:onInvocationStart"), calls); } + // ─── Execution input / result components ───────────────────────────── + + @Test + void invocationInfo_compatibilityConstructor_leavesExecutionInputNull() { + var info = new InvocationInfo("req-123", "arn:test", false, Instant.now()); + + assertNull(info.executionInput()); + } + + @Test + void invocationInfo_carriesExecutionInput() { + var input = Map.of("name", "World"); + + var info = new InvocationInfo("req-123", "arn:test", true, Instant.now(), input); + + assertEquals(input, info.executionInput()); + } + + @Test + void invocationEndInfo_compatibilityConstructor_leavesExecutionInputAndResultNull() { + var info = new InvocationEndInfo("req-123", "arn:test", false, InvocationStatus.SUCCEEDED, null); + + assertNull(info.executionInput()); + assertNull(info.executionResult()); + } + + @Test + void invocationEndInfo_carriesExecutionInputAndResult() { + var info = new InvocationEndInfo( + "req-123", "arn:test", false, InvocationStatus.SUCCEEDED, null, "World", "Hello World"); + + assertEquals("World", info.executionInput()); + assertEquals("Hello World", info.executionResult()); + } + + @Test + void invocationInfo_toString_omitsExecutionInput() { + var info = new InvocationInfo("req-123", "arn:test", true, Instant.now(), "s3cret-input"); + + var rendered = info.toString(); + + assertFalse(rendered.contains("s3cret-input"), "execution input must not leak into logs"); + assertTrue(rendered.contains("req-123")); + assertTrue(rendered.contains("arn:test")); + } + + @Test + void invocationEndInfo_toString_omitsExecutionInputAndResult() { + var info = new InvocationEndInfo( + "req-123", "arn:test", false, InvocationStatus.SUCCEEDED, null, "s3cret-input", "s3cret-result"); + + var rendered = info.toString(); + + assertFalse(rendered.contains("s3cret-input"), "execution input must not leak into logs"); + assertFalse(rendered.contains("s3cret-result"), "execution result must not leak into logs"); + assertTrue(rendered.contains("req-123")); + assertTrue(rendered.contains("SUCCEEDED")); + } + // ─── Helper methods ────────────────────────────────────────────────── private static InvocationInfo invocationInfo() {