Skip to content

Commit 51ce9f7

Browse files
ayushiahjoliaayushiahjolia
authored
feat(sdk): Add UpdatedOperationIds support for replay status tracking (#439)
1 parent 537a061 commit 51ce9f7

15 files changed

Lines changed: 449 additions & 42 deletions

File tree

‎examples/src/main/java/software/amazon/lambda/durable/examples/general/PluginExample.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
* <p>Expected output for a successful run:
1818
*
1919
* <pre>
20-
* [PLUGIN] onInvocationStart: requestId=..., executionArn=..., firstInvocation=true
20+
* [PLUGIN] onInvocationStart: requestId=..., durableExecutionArn=..., firstInvocation=true
2121
* [PLUGIN] onOperationStart: name=create-greeting, type=STEP
2222
* [PLUGIN] onUserFunctionStart: name=create-greeting, attempt=1
2323
* [PLUGIN] onUserFunctionEnd: name=create-greeting, succeeded=true
@@ -54,8 +54,8 @@ static class LoggingPlugin implements DurableExecutionPlugin {
5454
@Override
5555
public void onInvocationStart(InvocationInfo info) {
5656
System.out.printf(
57-
"[PLUGIN] onInvocationStart: requestId=%s, executionArn=%s, firstInvocation=%s%n",
58-
info.requestId(), info.executionArn(), info.isFirstInvocation());
57+
"[PLUGIN] onInvocationStart: requestId=%s, durableExecutionArn=%s, firstInvocation=%s%n",
58+
info.requestId(), info.durableExecutionArn(), info.isFirstInvocation());
5959
}
6060

6161
@Override

‎otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DeterministicIdGenerator.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,15 @@ public class DeterministicIdGenerator implements IdGenerator {
2626

2727
private final AtomicReference<String> executionTraceId = new AtomicReference<>(null);
2828
private final ThreadLocal<String> pendingSpanOperationId = new ThreadLocal<>();
29-
private final AtomicReference<String> executionArn = new AtomicReference<>(null);
29+
private final AtomicReference<String> durableExecutionArn = new AtomicReference<>(null);
3030

3131
/**
3232
* Sets the execution ARN used for generating deterministic IDs.
3333
*
3434
* @param arn the durable execution ARN
3535
*/
36-
public void setExecutionArn(String arn) {
37-
this.executionArn.set(arn);
36+
public void setDurableExecutionArn(String arn) {
37+
this.durableExecutionArn.set(arn);
3838
this.executionTraceId.set(generateTraceIdFromArn(arn));
3939
}
4040

@@ -96,7 +96,7 @@ private String generateTraceIdFromArn(String arn) {
9696
* <p>Uses SHA-256 hash truncated to 16 hex chars (64 bits) for the span ID.
9797
*/
9898
private String generateSpanIdFromOperation(String operationId) {
99-
var arn = executionArn.get();
99+
var arn = durableExecutionArn.get();
100100
var input = arn != null ? arn + ":" + operationId : operationId;
101101
var hash = sha256(input);
102102
// Span ID is 16 hex chars (8 bytes)

‎otel-plugin/src/main/java/software/amazon/lambda/durable/otel/MdcSpanEnricher.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,16 +37,16 @@ private MdcSpanEnricher() {}
3737
/**
3838
* Injects the current span's trace ID and span ID into MDC.
3939
*
40-
* @param executionArn the durable execution ARN (may be null)
40+
* @param durableExecutionArn the durable execution ARN (may be null)
4141
*/
42-
public static void inject(String executionArn) {
42+
public static void inject(String durableExecutionArn) {
4343
var span = Span.current();
4444
if (span.getSpanContext().isValid()) {
4545
MDC.put(MDC_TRACE_ID, span.getSpanContext().getTraceId());
4646
MDC.put(MDC_SPAN_ID, span.getSpanContext().getSpanId());
4747
}
48-
if (executionArn != null) {
49-
MDC.put(MDC_EXECUTION_ARN, executionArn);
48+
if (durableExecutionArn != null) {
49+
MDC.put(MDC_EXECUTION_ARN, durableExecutionArn);
5050
}
5151
}
5252

‎otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OpenTelemetryDurablePlugin.java‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ public class OpenTelemetryDurablePlugin implements DurableExecutionPlugin {
6060

6161
// Per-invocation state
6262
private volatile Span invocationSpan;
63-
private volatile String executionArn;
63+
private volatile String durableExecutionArn;
6464

6565
// Thread-safe storage for operation spans (keyed by operationId) — open spans that need ending
6666
private final ConcurrentHashMap<String, Span> operationSpans = new ConcurrentHashMap<>();
@@ -116,16 +116,16 @@ public OpenTelemetryDurablePlugin(
116116

117117
@Override
118118
public void onInvocationStart(InvocationInfo info) {
119-
this.executionArn = info.executionArn();
120-
idGenerator.setExecutionArn(info.executionArn());
119+
this.durableExecutionArn = info.durableExecutionArn();
120+
idGenerator.setDurableExecutionArn(info.durableExecutionArn());
121121

122122
// Extract parent context from Lambda environment (X-Ray, W3C, etc.)
123123
var extractedParentContext = contextExtractor.extract();
124124

125125
// Create invocation span as child of extracted context
126126
var spanBuilder = tracer.spanBuilder("durable.invocation")
127127
.setParent(extractedParentContext)
128-
.setAttribute(DURABLE_EXECUTION_ARN, info.executionArn())
128+
.setAttribute(DURABLE_EXECUTION_ARN, info.durableExecutionArn())
129129
.setAttribute(DURABLE_FIRST_INVOCATION, info.isFirstInvocation());
130130

131131
if (info.requestId() != null) {
@@ -189,7 +189,7 @@ public void onOperationStart(OperationInfo info) {
189189

190190
var spanBuilder = tracer.spanBuilder(spanName(info.type(), info.subType(), info.name()))
191191
.setParent(parentContext)
192-
.setAttribute(DURABLE_EXECUTION_ARN, executionArn)
192+
.setAttribute(DURABLE_EXECUTION_ARN, durableExecutionArn)
193193
.setAttribute(DURABLE_OPERATION_ID, info.id())
194194
.setAttribute(DURABLE_OPERATION_TYPE, info.type());
195195

@@ -239,7 +239,7 @@ public void onUserFunctionStart(UserFunctionStartInfo info) {
239239
.setParent(parentContext)
240240
.setStartTimestamp(info.startTimestamp() != null ? info.startTimestamp() : Instant.now());
241241

242-
spanBuilder.setAttribute(DURABLE_EXECUTION_ARN, executionArn);
242+
spanBuilder.setAttribute(DURABLE_EXECUTION_ARN, durableExecutionArn);
243243
spanBuilder.setAttribute(DURABLE_OPERATION_ID, info.id());
244244

245245
if (info.type() != null) {
@@ -261,7 +261,7 @@ public void onUserFunctionStart(UserFunctionStartInfo info) {
261261

262262
// Inject trace context into MDC for log-trace correlation
263263
if (enableMdc) {
264-
MdcSpanEnricher.inject(executionArn);
264+
MdcSpanEnricher.inject(durableExecutionArn);
265265
}
266266
}
267267

‎otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DeterministicIdGeneratorTest.java‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ void generateTraceId_withoutArn_returnsRandom() {
2929

3030
@Test
3131
void generateTraceId_withArn_returnsDeterministic() {
32-
generator.setExecutionArn("arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1");
32+
generator.setDurableExecutionArn("arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1");
3333

3434
var id1 = generator.generateTraceId();
3535
var id2 = generator.generateTraceId();
@@ -40,10 +40,10 @@ void generateTraceId_withArn_returnsDeterministic() {
4040

4141
@Test
4242
void generateTraceId_differentArns_produceDifferentIds() {
43-
generator.setExecutionArn("arn:exec1");
43+
generator.setDurableExecutionArn("arn:exec1");
4444
var id1 = generator.generateTraceId();
4545

46-
generator.setExecutionArn("arn:exec2");
46+
generator.setDurableExecutionArn("arn:exec2");
4747
var id2 = generator.generateTraceId();
4848

4949
assertNotEquals(id1, id2);
@@ -60,7 +60,7 @@ void generateSpanId_withoutOperationId_returnsRandom() {
6060

6161
@Test
6262
void generateSpanId_withOperationId_returnsDeterministic() {
63-
generator.setExecutionArn("arn:exec1");
63+
generator.setDurableExecutionArn("arn:exec1");
6464
generator.setNextSpanOperationId("op-hash-1");
6565
var id1 = generator.generateSpanId();
6666

@@ -73,7 +73,7 @@ void generateSpanId_withOperationId_returnsDeterministic() {
7373

7474
@Test
7575
void generateSpanId_differentOperationIds_produceDifferentIds() {
76-
generator.setExecutionArn("arn:exec1");
76+
generator.setDurableExecutionArn("arn:exec1");
7777

7878
generator.setNextSpanOperationId("op-1");
7979
var id1 = generator.generateSpanId();
@@ -96,7 +96,7 @@ void generateSpanId_consumesPendingId() {
9696

9797
@Test
9898
void generateSpanIdForOperation_doesNotConsumePending() {
99-
generator.setExecutionArn("arn:exec1");
99+
generator.setDurableExecutionArn("arn:exec1");
100100
generator.setNextSpanOperationId("op-1");
101101

102102
// This should NOT consume the pending
@@ -112,7 +112,7 @@ void generateSpanIdForOperation_doesNotConsumePending() {
112112

113113
@Test
114114
void generateSpanIdForOperation_isDeterministic() {
115-
generator.setExecutionArn("arn:exec1");
115+
generator.setDurableExecutionArn("arn:exec1");
116116

117117
var id1 = generator.generateSpanIdForOperation("op-1");
118118
var id2 = generator.generateSpanIdForOperation("op-1");
@@ -122,15 +122,15 @@ void generateSpanIdForOperation_isDeterministic() {
122122

123123
@Test
124124
void traceId_isValidHex() {
125-
generator.setExecutionArn("arn:exec1");
125+
generator.setDurableExecutionArn("arn:exec1");
126126
var traceId = generator.generateTraceId();
127127

128128
assertTrue(traceId.matches("[0-9a-f]{32}"), "Trace ID should be 32 hex chars: " + traceId);
129129
}
130130

131131
@Test
132132
void spanId_isValidHex() {
133-
generator.setExecutionArn("arn:exec1");
133+
generator.setDurableExecutionArn("arn:exec1");
134134
generator.setNextSpanOperationId("op-1");
135135
var spanId = generator.generateSpanId();
136136

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

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,96 @@ void waitAndResume_producesSpansAcrossInvocations() {
158158
assertEquals(2, invocationSpans.size(), "Should have 2 invocation spans (one per run)");
159159
}
160160

161+
@Test
162+
void waitCompletedDuringSuspension_producesOperationSpan() {
163+
var runner = LocalDurableTestRunner.create(
164+
String.class,
165+
(input, ctx) -> {
166+
ctx.step("before-wait", String.class, stepCtx -> "pre");
167+
ctx.wait("pause", Duration.ofMinutes(1));
168+
ctx.step("after-wait", String.class, stepCtx -> "post");
169+
return "done";
170+
},
171+
otelConfig);
172+
173+
// First invocation: step completes, wait starts, suspend
174+
var result1 = runner.run("input");
175+
assertEquals(ExecutionStatus.PENDING, result1.getStatus());
176+
177+
// The wait operation span should be open (ended with PENDING status at invocation end)
178+
var spansAfterFirst = spanExporter.getFinishedSpanItems();
179+
var waitSpansAfterFirst = spansAfterFirst.stream()
180+
.filter(s -> s.getName().equals("durable.wait:pause"))
181+
.toList();
182+
assertEquals(1, waitSpansAfterFirst.size(), "Wait span should exist after first invocation (ended as PENDING)");
183+
184+
// Advance time (wait completes externally) and resume
185+
runner.advanceTime();
186+
var result2 = runner.run("input");
187+
assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
188+
189+
// After second invocation, the wait should have a second operation span
190+
// (fired by onOperationEnd via updatedOperationIds) showing it completed
191+
var allSpans = spanExporter.getFinishedSpanItems();
192+
var waitSpansTotal = allSpans.stream()
193+
.filter(s -> s.getName().equals("durable.wait:pause"))
194+
.toList();
195+
assertEquals(
196+
2,
197+
waitSpansTotal.size(),
198+
"Wait should have 2 spans: one PENDING from first invocation, one completed from second. Got: "
199+
+ allSpans.stream()
200+
.map(SpanData::getName)
201+
.filter(n -> n.contains("pause"))
202+
.toList());
203+
}
204+
205+
@Test
206+
void invokeFailedDuringSuspension_producesErrorSpan() {
207+
var runner = LocalDurableTestRunner.create(
208+
String.class,
209+
(input, ctx) -> {
210+
try {
211+
ctx.invoke("call-target", "target-fn", "{}", String.class);
212+
} catch (Exception e) {
213+
// expected failure
214+
}
215+
return "handled";
216+
},
217+
otelConfig);
218+
219+
// First invocation: invoke starts, suspends waiting for target
220+
var result1 = runner.run("input");
221+
assertEquals(ExecutionStatus.PENDING, result1.getStatus());
222+
223+
// Target fails while Lambda is frozen
224+
runner.failChainedInvoke(
225+
"call-target",
226+
software.amazon.awssdk.services.lambda.model.ErrorObject.builder()
227+
.errorType("TargetError")
228+
.errorMessage("target function failed")
229+
.build());
230+
231+
// Resume
232+
var result2 = runner.run("input");
233+
assertEquals(ExecutionStatus.SUCCEEDED, result2.getStatus());
234+
235+
// The invoke operation span from the second invocation should have ERROR status
236+
var allSpans = spanExporter.getFinishedSpanItems();
237+
var invokeSpans = allSpans.stream()
238+
.filter(s -> s.getName().contains("call-target"))
239+
.toList();
240+
assertTrue(invokeSpans.size() >= 2, "Should have at least 2 invoke spans (one PENDING, one ended with error)");
241+
242+
// The span from the second invocation should be marked as error
243+
var errorSpan = invokeSpans.stream()
244+
.filter(s -> s.getStatus().getStatusCode() == io.opentelemetry.api.trace.StatusCode.ERROR)
245+
.findFirst();
246+
assertTrue(
247+
errorSpan.isPresent(),
248+
"Should have an invoke span with ERROR status when invoke fails during suspension");
249+
}
250+
161251
@Test
162252
void childContext_producesNestedSpans() {
163253
var runner = LocalDurableTestRunner.create(

0 commit comments

Comments
 (0)