From 6962f5ab76b6bccf4b7908d5974bfd0a2e4b1ea4 Mon Sep 17 00:00:00 2001 From: Zhongke Chen Date: Fri, 7 Aug 2026 02:53:41 +0000 Subject: [PATCH] fix(otel): close spans child first --- .../durable/otel/InvocationOtelPlugin.java | 46 +++++++++++-------- .../otel/InvocationOtelPluginTest.java | 32 +++++++++++++ 2 files changed, 60 insertions(+), 18 deletions(-) diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java index 9111b80d2..0806aa4d7 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/InvocationOtelPlugin.java @@ -22,6 +22,7 @@ import io.opentelemetry.semconv.ServiceAttributes; import java.time.Instant; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedDeque; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; @@ -99,6 +100,9 @@ public class InvocationOtelPlugin implements DurableExecutionPlugin { // Store operation span contexts for parent resolution (keyed by operationId) private final ConcurrentHashMap operationContexts = new ConcurrentHashMap<>(); + // Operation start order, drained in reverse so children end before parents + private final ConcurrentLinkedDeque operationStartOrder = new ConcurrentLinkedDeque<>(); + /** * Creates an OTel plugin with default settings: X-Ray context extraction, MDC enabled. * @@ -230,24 +234,7 @@ public void onInvocationStart(InvocationInfo info) { public void onInvocationEnd(InvocationEndInfo info) { if (invocationSpan == null) return; - // End still-open operation spans without stamping a status — no terminal - // durable.operation.status means still running (STARTED). A later invocation's - // onOperationEnd emits a continuation span with the real terminal status. - for (var entry : operationSpans.entrySet()) { - entry.getValue().end(); - } - operationSpans.clear(); - operationContexts.clear(); - - // End any attempt spans that are still open (e.g., crash before onUserFunctionEnd) - for (var entry : attemptScopes.entrySet()) { - entry.getValue().close(); - } - attemptScopes.clear(); - for (var entry : attemptSpans.entrySet()) { - entry.getValue().end(); - } - attemptSpans.clear(); + endOpenSpansChildFirst(); // End invocation span invocationSpan.setAttribute( @@ -350,6 +337,7 @@ public void onOperationStart(OperationInfo info) { // Store the open span — will be ended in onOperationEnd or onInvocationEnd operationSpans.put(info.id(), span); operationContexts.put(info.id(), span.getSpanContext()); + operationStartOrder.addLast(info.id()); } @Override @@ -519,6 +507,28 @@ public void onUserFunctionEnd(UserFunctionEndInfo info) { // ─── Helpers ───────────────────────────────────────────────────────── + private void endOpenSpansChildFirst() { + // Attempt spans are children of operation spans. + for (var scope : attemptScopes.values()) { + scope.close(); + } + attemptScopes.clear(); + for (var span : attemptSpans.values()) { + span.end(); + } + attemptSpans.clear(); + + String operationId; + while ((operationId = operationStartOrder.pollLast()) != null) { + var span = operationSpans.remove(operationId); + if (span != null) { + span.end(); + } + } + operationSpans.clear(); + operationContexts.clear(); + } + private Context resolveParentContext(String parentId) { if (parentId != null) { var parentSpanContext = operationContexts.get(parentId); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java index 2536d0ad3..0e6ec348a 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginTest.java @@ -345,6 +345,38 @@ void operationNotCompleted_spanEndedAtInvocationEnd() { assertEquals("my-wait", operationSpan.getName()); } + @Test + void invocationEnd_closesNestedSpansChildFirst() { + var parentId = "op-parent"; + var childId = "op-child"; + plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec1", true, Instant.now())); + plugin.onOperationStart(new OperationInfo( + parentId, "parent-context", "CONTEXT", "RunInChildContext", null, Instant.now(), null, false)); + plugin.onOperationStart( + new OperationInfo(childId, "child-step", "STEP", "Step", parentId, Instant.now(), null, false)); + plugin.onUserFunctionStart( + new UserFunctionStartInfo(childId, "child-step", "STEP", "Step", parentId, Instant.now(), false, 1)); + + plugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec1", true, InvocationStatus.PENDING, null)); + + var parentSpan = spanByName("parent-context"); + var childSpan = spanByName("child-step"); + var attemptSpan = spanByName("child-step attempt 1"); + var spans = spanExporter.getFinishedSpanItems(); + assertTrue( + spans.indexOf(attemptSpan) < spans.indexOf(childSpan), + "Attempt span must be exported before its operation span"); + assertTrue( + spans.indexOf(childSpan) < spans.indexOf(parentSpan), + "Child operation span must be exported before its parent operation span"); + assertTrue( + attemptSpan.getEndEpochNanos() <= childSpan.getEndEpochNanos(), + "Attempt span must end before its operation span"); + assertTrue( + childSpan.getEndEpochNanos() <= parentSpan.getEndEpochNanos(), + "Child operation span must end before its parent operation span"); + } + @Test void sampling_disabled_producesNoSpans() { spanExporter = InMemorySpanExporter.create();