Skip to content

Commit 4c290bd

Browse files
wangyb-AAlex Wang
andauthored
feat(plugin): surface execution input and result on invocation hooks (#622)
* feat(plugin): surface execution input and result on invocation hooks Adds InvocationInfo.executionInput (the deserialized handler input) plus InvocationEndInfo.executionInput/executionResult (the value the handler returned, populated only when the invocation completed the execution successfully), so instrumentation plugins can record execution I/O. Both records keep a constructor at their previous arity, so existing callers compile unchanged. The plugin-facing deserialization is skipped when no plugins are registered and yields null on failure, leaving the authoritative input extraction to surface errors at its original point. Mirrors the JS SDK's InvocationBaseInfo.executionInput / InvocationEndInfo.executionResult and the Python SDK's execution_input/execution_result (aws-durable-execution-sdk-python#616). * style(plugin): mark new hook components @experimental Matches the convention established by #620, which replaced the preview-API @deprecated markers with @experimental on the individual record components (as on InvocationEndInfo.executionError). * fix(plugin): share one deserialized input and keep payloads out of toString Review feedback: - Deserialize the handler input once and hand the same instance to the plugin hooks and the handler. The previous plugin-only extraction deserialized a second time, doubling the cost, giving plugins a different object than the handler, and re-running side effects in a stateful custom SerDes. A deserialization failure is now captured and rethrown after onInvocationStart, so the start and end hooks stay paired as before. - Override toString() on both invocation records to omit executionInput and executionResult. Records render every component, so plugins that log the info object whole would have begun emitting customer payloads (possibly secrets or personal data). Output is otherwise unchanged from before this branch; mirrors repr=False on the Python fields. * fix(plugin): keep hooks paired for sneaky checked exceptions SerDes.deserialize declares no checked exceptions, so an implementation may sneaky-throw one (as DurableInputOutputSerDes does for IOException). A RuntimeException-only capture let such a failure bypass onInvocationStart while the future's failure path still fired onInvocationEnd(FAILED), leaving a plugin that keys state on the start hook with an unpaired end. Capture Throwable and rethrow it unchanged via ExceptionHelper.sneakyThrow after the start hook. Adds a regression test with a SerDes that sneaky-throws IOException: verified failing (start hook count 0) with the narrower capture. --------- Co-authored-by: Alex Wang <wangyb@amazon.com>
1 parent 8ce3f23 commit 4c290bd

5 files changed

Lines changed: 348 additions & 10 deletions

File tree

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

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,15 @@
44

55
import static org.junit.jupiter.api.Assertions.*;
66

7+
import java.io.IOException;
78
import java.time.Duration;
89
import java.util.ArrayList;
910
import java.util.Collections;
1011
import java.util.List;
12+
import java.util.Map;
13+
import java.util.concurrent.ConcurrentHashMap;
1114
import java.util.concurrent.atomic.AtomicInteger;
15+
import java.util.concurrent.atomic.AtomicReference;
1216
import org.junit.jupiter.api.Test;
1317
import software.amazon.awssdk.services.lambda.model.ErrorObject;
1418
import software.amazon.awssdk.services.lambda.model.OperationStatus;
@@ -19,7 +23,10 @@
1923
import software.amazon.lambda.durable.model.WaitForConditionResult;
2024
import software.amazon.lambda.durable.plugin.*;
2125
import software.amazon.lambda.durable.retry.RetryStrategies;
26+
import software.amazon.lambda.durable.serde.JacksonSerDes;
27+
import software.amazon.lambda.durable.serde.SerDes;
2228
import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
29+
import software.amazon.lambda.durable.util.ExceptionHelper;
2330

2431
/** Integration tests verifying plugin hooks fire correctly during durable execution lifecycle. */
2532
class PluginIntegrationTest {
@@ -122,6 +129,167 @@ void plugin_receivesInvocationEnd_withFailedStatus_onError() {
122129
assertNotNull(plugin.invocationEnds.get(0).executionError());
123130
}
124131

132+
@Test
133+
void plugin_invocationHooks_carryExecutionInputAndResult_onSuccess() {
134+
var plugin = new RecordingPlugin();
135+
var config = DurableConfig.builder().withPlugins(plugin).build();
136+
137+
var runner = LocalDurableTestRunner.create(
138+
String.class,
139+
(input, context) -> context.step("greet", String.class, stepCtx -> "Hello " + input),
140+
config);
141+
142+
var result = runner.runUntilComplete("World");
143+
144+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
145+
146+
// The deserialized handler input reaches the start hook ...
147+
assertEquals("World", plugin.invocationStarts.get(0).executionInput());
148+
// ... and the end hook, alongside the value the handler returned.
149+
var end = plugin.invocationEnds.get(0);
150+
assertEquals("World", end.executionInput());
151+
assertEquals("Hello World", end.executionResult());
152+
}
153+
154+
@Test
155+
void plugin_invocationEnd_omitsExecutionResult_onFailure() {
156+
var plugin = new RecordingPlugin();
157+
var config = DurableConfig.builder().withPlugins(plugin).build();
158+
159+
var runner = LocalDurableTestRunner.create(
160+
String.class,
161+
(input, context) -> context.step(
162+
"failing",
163+
String.class,
164+
stepCtx -> {
165+
throw new RuntimeException("boom");
166+
},
167+
StepConfig.builder()
168+
.retryStrategy(RetryStrategies.Presets.NO_RETRY)
169+
.build()),
170+
config);
171+
172+
var result = runner.run("input");
173+
174+
assertEquals(ExecutionStatus.FAILED, result.getStatus());
175+
176+
var end = plugin.invocationEnds.get(0);
177+
assertEquals("input", end.executionInput());
178+
assertNull(end.executionResult(), "a failed execution produced no result");
179+
}
180+
181+
@Test
182+
void plugin_invocationEnd_omitsExecutionResult_onSuspension() {
183+
var plugin = new RecordingPlugin();
184+
var config = DurableConfig.builder().withPlugins(plugin).build();
185+
186+
var runner = LocalDurableTestRunner.create(
187+
String.class,
188+
(input, context) -> {
189+
context.wait("pause", Duration.ofMinutes(5));
190+
return "complete";
191+
},
192+
config);
193+
194+
var result = runner.run("input");
195+
196+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
197+
198+
var end = plugin.invocationEnds.get(0);
199+
assertEquals(InvocationStatus.PENDING, end.invocationStatus());
200+
assertEquals("input", end.executionInput());
201+
assertNull(end.executionResult(), "a suspended invocation has not produced a result yet");
202+
}
203+
204+
@Test
205+
void plugin_executionInput_isDeserializedOnce_andSharedWithHandler() {
206+
var serDes = new CountingSerDes();
207+
var plugin = new RecordingPlugin();
208+
var config =
209+
DurableConfig.builder().withPlugins(plugin).withSerDes(serDes).build();
210+
var handlerInput = new AtomicReference<Object>();
211+
212+
var runner = LocalDurableTestRunner.create(
213+
String.class,
214+
(input, context) -> {
215+
handlerInput.set(input);
216+
return context.step("greet", String.class, stepCtx -> "Hello " + input);
217+
},
218+
config);
219+
220+
var result = runner.runUntilComplete("World");
221+
222+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
223+
// The handler input is deserialized once and that same instance reaches the hooks — a second
224+
// deserialization would double the cost and re-run side effects in a stateful SerDes.
225+
assertEquals(1, serDes.inputDeserializations("World"));
226+
assertSame(handlerInput.get(), plugin.invocationStarts.get(0).executionInput());
227+
assertSame(handlerInput.get(), plugin.invocationEnds.get(0).executionInput());
228+
}
229+
230+
/** SerDes that counts how many times each payload is deserialized. */
231+
static class CountingSerDes implements SerDes {
232+
private final JacksonSerDes delegate = new JacksonSerDes();
233+
private final Map<String, AtomicInteger> deserializations = new ConcurrentHashMap<>();
234+
235+
@Override
236+
public String serialize(Object value) {
237+
return delegate.serialize(value);
238+
}
239+
240+
@Override
241+
public <T> T deserialize(String data, TypeToken<T> typeToken) {
242+
deserializations
243+
.computeIfAbsent(data, ignored -> new AtomicInteger())
244+
.incrementAndGet();
245+
return delegate.deserialize(data, typeToken);
246+
}
247+
248+
int inputDeserializations(String value) {
249+
var counter = deserializations.get(delegate.serialize(value));
250+
return counter == null ? 0 : counter.get();
251+
}
252+
}
253+
254+
@Test
255+
void plugin_hooksStayPaired_whenSerDesSneakyThrowsCheckedException() {
256+
var plugin = new RecordingPlugin();
257+
var config = DurableConfig.builder()
258+
.withPlugins(plugin)
259+
.withSerDes(new SneakyThrowingSerDes())
260+
.build();
261+
262+
var runner = LocalDurableTestRunner.create(String.class, (input, context) -> "unreachable", config);
263+
264+
var result = runner.run("input");
265+
266+
assertEquals(ExecutionStatus.FAILED, result.getStatus());
267+
// SerDes.deserialize declares no checked exceptions, so an implementation can sneaky-throw one. The start
268+
// hook must still fire, otherwise a plugin keying state on it sees an end with no start.
269+
assertEquals(1, plugin.invocationStarts.size(), "onInvocationStart must fire before the failure surfaces");
270+
assertNull(plugin.invocationStarts.get(0).executionInput(), "input could not be deserialized");
271+
assertEquals(1, plugin.invocationEnds.size());
272+
assertEquals(InvocationStatus.FAILED, plugin.invocationEnds.get(0).invocationStatus());
273+
}
274+
275+
/**
276+
* SerDes that sneaky-throws a checked exception from deserialize, as DurableInputOutputSerDes does on IO errors.
277+
*/
278+
static class SneakyThrowingSerDes implements SerDes {
279+
private final JacksonSerDes delegate = new JacksonSerDes();
280+
281+
@Override
282+
public String serialize(Object value) {
283+
return delegate.serialize(value);
284+
}
285+
286+
@Override
287+
public <T> T deserialize(String data, TypeToken<T> typeToken) {
288+
ExceptionHelper.sneakyThrow(new IOException("sneaky checked failure"));
289+
return null;
290+
}
291+
}
292+
125293
// ─── Operation-level hooks ───────────────────────────────────────────
126294

127295
@Test

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

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import java.nio.charset.StandardCharsets;
88
import java.util.concurrent.CompletableFuture;
99
import java.util.concurrent.CompletionException;
10+
import java.util.concurrent.atomic.AtomicReference;
1011
import java.util.function.BiFunction;
1112
import org.slf4j.Logger;
1213
import org.slf4j.LoggerFactory;
@@ -60,21 +61,41 @@ public static <I, O> DurableExecutionOutput execute(
6061
var executionArn = input.durableExecutionArn();
6162

6263
executionManager.registerActiveThread(null);
64+
// Captured for onInvocationEnd, which runs outside the handler thread below.
65+
var pluginExecutionInput = new AtomicReference<>();
6366
var handlerFuture = CompletableFuture.supplyAsync(
6467
() -> {
6568
executionManager.setCurrentThreadContext(new ThreadContext(null, ThreadType.CONTEXT));
6669

70+
// Deserialize once and share the value with the plugin hooks and the handler below. A second
71+
// deserialization would double the cost, hand plugins a different object than the handler, and
72+
// re-run any side effects in a stateful custom SerDes. A failure is captured rather than thrown
73+
// so onInvocationStart still fires before it surfaces, keeping the start/end hooks paired.
74+
// SerDes is a public extension point whose deserialize declares no checked exceptions, so an
75+
// implementation may sneaky-throw one; capture every Throwable and rethrow it unchanged.
76+
I userInput = null;
77+
Throwable inputFailure = null;
78+
try {
79+
userInput = extractUserInput(
80+
executionManager.getExecutionOperation(), config.getSerDes(), inputType);
81+
} catch (Throwable t) {
82+
inputFailure = t;
83+
}
84+
pluginExecutionInput.set(userInput);
85+
6786
// onInvocationStart runs on the user thread so plugins can
6887
// inject ThreadLocal objects, update MDC, etc.
6988
// executionStartTime comes from the initial EXECUTION operation in the first backend event.
7089
pluginRunner.onInvocationStart(new InvocationInfo(
7190
requestId,
7291
executionArn,
7392
isFirstInvocation,
74-
executionManager.getExecutionOperation().startTimestamp()));
93+
executionManager.getExecutionOperation().startTimestamp(),
94+
userInput));
95+
if (inputFailure != null) {
96+
ExceptionHelper.sneakyThrow(inputFailure);
97+
}
7598

76-
var userInput = extractUserInput(
77-
executionManager.getExecutionOperation(), config.getSerDes(), inputType);
7899
var context = DurableContextImpl.createRootContext(executionManager, config, lambdaContext);
79100
DurableContextImpl.setCurrentContext(context);
80101
// use a try-with-resources to clear logger properties
@@ -103,6 +124,8 @@ public static <I, O> DurableExecutionOutput execute(
103124
executionArn,
104125
isFirstInvocation,
105126
InvocationStatus.PENDING,
127+
null,
128+
pluginExecutionInput.get(),
106129
null);
107130
return DurableExecutionOutput.pending();
108131
}
@@ -119,7 +142,9 @@ public static <I, O> DurableExecutionOutput execute(
119142
executionArn,
120143
isFirstInvocation,
121144
InvocationStatus.RETRYING,
122-
cause);
145+
cause,
146+
pluginExecutionInput.get(),
147+
null);
123148
throw unrecoverableDurableExecutionException;
124149
}
125150

@@ -131,7 +156,9 @@ public static <I, O> DurableExecutionOutput execute(
131156
executionArn,
132157
isFirstInvocation,
133158
InvocationStatus.FAILED,
134-
cause);
159+
cause,
160+
pluginExecutionInput.get(),
161+
null);
135162
return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes()));
136163
}
137164
// user handler complete successfully
@@ -145,7 +172,9 @@ public static <I, O> DurableExecutionOutput execute(
145172
executionArn,
146173
isFirstInvocation,
147174
InvocationStatus.SUCCEEDED,
148-
null);
175+
null,
176+
pluginExecutionInput.get(),
177+
result);
149178
return output;
150179
})
151180
.join();
@@ -163,8 +192,11 @@ private static void fireOnInvocationEnd(
163192
String executionArn,
164193
boolean isFirstInvocation,
165194
InvocationStatus status,
166-
Throwable error) {
167-
pluginRunner.onInvocationEnd(new InvocationEndInfo(requestId, executionArn, isFirstInvocation, status, error));
195+
Throwable error,
196+
Object executionInput,
197+
Object executionResult) {
198+
pluginRunner.onInvocationEnd(new InvocationEndInfo(
199+
requestId, executionArn, isFirstInvocation, status, error, executionInput, executionResult));
168200
}
169201

170202
private static String handleLargePayload(ExecutionManager executionManager, String outputPayload) {

‎sdk/src/main/java/software/amazon/lambda/durable/plugin/InvocationEndInfo.java‎

Lines changed: 43 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,52 @@
1212
* @param isFirstInvocation true if this is the first invocation of the execution
1313
* @param invocationStatus the invocation outcome (SUCCEEDED, FAILED, or PENDING)
1414
* @param executionError non-null if the execution failed; this component is experimental
15+
* @param executionInput the deserialized execution input passed to the user handler, or null when no plugins are
16+
* registered or the input could not be deserialized; this component is experimental
17+
* @param executionResult the value the user handler returned, or null unless the invocation completed the execution
18+
* successfully; this component is experimental
1519
*/
1620
public record InvocationEndInfo(
1721
String requestId,
1822
String durableExecutionArn,
1923
boolean isFirstInvocation,
2024
InvocationStatus invocationStatus,
21-
@Experimental Throwable executionError) {}
25+
@Experimental Throwable executionError,
26+
@Experimental Object executionInput,
27+
@Experimental Object executionResult) {
28+
29+
/**
30+
* Creates invocation-end information without the execution input or result.
31+
*
32+
* <p>Retained so callers written before {@code executionInput} and {@code executionResult} were added keep
33+
* compiling; both resolve to null.
34+
*
35+
* @param requestId the Lambda request ID for this invocation
36+
* @param durableExecutionArn the durable execution ARN
37+
* @param isFirstInvocation true if this is the first invocation of the execution
38+
* @param invocationStatus the invocation outcome
39+
* @param executionError non-null if the execution failed
40+
*/
41+
public InvocationEndInfo(
42+
String requestId,
43+
String durableExecutionArn,
44+
boolean isFirstInvocation,
45+
InvocationStatus invocationStatus,
46+
Throwable executionError) {
47+
this(requestId, durableExecutionArn, isFirstInvocation, invocationStatus, executionError, null, null);
48+
}
49+
50+
/**
51+
* Returns a representation that omits {@code executionInput} and {@code executionResult}.
52+
*
53+
* <p>The generated representation would render both payloads, so plugins that log this object whole would start
54+
* emitting customer inputs and results, potentially including secrets or personal data. Read the components
55+
* explicitly to record them.
56+
*/
57+
@Override
58+
public String toString() {
59+
return "InvocationEndInfo[requestId=" + requestId + ", durableExecutionArn=" + durableExecutionArn
60+
+ ", isFirstInvocation=" + isFirstInvocation + ", invocationStatus=" + invocationStatus
61+
+ ", executionError=" + executionError + "]";
62+
}
63+
}

0 commit comments

Comments
 (0)