From 1b92765dbcbf373801e69fbe875971587eac61d2 Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Mon, 24 Aug 2026 13:23:29 -0700 Subject: [PATCH] fix(otel): End all recording spans on non-terminal invocations --- .../durable/otel/ExecutionOtelPlugin.java | 282 +++++++++--------- .../durable/otel/InvocationOtelPlugin.java | 108 ++++--- .../durable/otel/ExecutionOtelPluginTest.java | 148 ++++++++- .../otel/InvocationOtelPluginTest.java | 90 ++++++ .../otel/LifecycleTrackingSpanProcessor.java | 58 ++++ 5 files changed, 507 insertions(+), 179 deletions(-) create mode 100644 otel-plugin/src/test/java/software/amazon/lambda/durable/otel/LifecycleTrackingSpanProcessor.java diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java index 9cd8ac33a..09214055c 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionOtelPlugin.java @@ -41,13 +41,15 @@ * @@ -77,6 +79,12 @@ * so it is not exported this invocation (effectively {@link StatusCode#UNSET}). * * + *

Deferred operation spans and context operations. An operation span is created only when the operation + * completes ({@code onOperationEnd}), so each operation is exported once even across suspend/resume. Before completion + * it is represented by a deterministic, non-recording {@link SpanContext}. A CONTEXT operation makes this placeholder + * current, so {@code Span.current()} enrichment is not recorded on the final operation span. The placeholder uses the + * Invocation span's resolved sampling metadata when available. + * *

Thread-safe: uses {@link ConcurrentHashMap} for span/scope storage since the SDK runs user code on multiple * threads. */ @@ -94,7 +102,6 @@ public class ExecutionOtelPlugin implements DurableExecutionPlugin { // Per-invocation state private volatile boolean tracingEnabled; - private volatile Span workflowSpan; private volatile Span invocationSpan; private volatile String durableExecutionArn; @@ -109,16 +116,25 @@ public class ExecutionOtelPlugin implements DurableExecutionPlugin { /** Immutable snapshot of the resolved execution trace, read atomically through a single volatile reference. */ private record ExecutionTrace(String traceId, TraceFlags flags) {} - // Thread-safe storage for operation spans (keyed by operationId) — open spans that need ending - private final ConcurrentHashMap operationSpans = new ConcurrentHashMap<>(); + // Between invocations the Workflow span exists only as a deterministic context that operations parent onto; the + // recording span is started and ended in a single call on the terminal invocation, so it is never left open. The + // execution ancestor and start time are retained so that span can be built at invocation end. + private volatile SpanContext workflowSpanContext; + private volatile SpanContext executionAncestor; + private volatile Instant executionStartTime; // Thread-safe storage for attempt spans/scopes (keyed by operationId + "-" + attempt) private final ConcurrentHashMap attemptSpans = new ConcurrentHashMap<>(); private final ConcurrentHashMap attemptScopes = new ConcurrentHashMap<>(); - // Store operation span contexts for parent resolution (keyed by operationId) + // Deterministic operation contexts (keyed by operationId), held between start and end so children and attempts can + // parent onto an operation whose recording span is not created until onOperationEnd. private final ConcurrentHashMap operationContexts = new ConcurrentHashMap<>(); + // Start timestamps captured at onOperationStart (keyed by operationId), used when onOperationEnd carries none — + // virtual map/parallel child contexts report null timestamps at end. + private final ConcurrentHashMap operationStartTimes = new ConcurrentHashMap<>(); + /** * Creates a Workflow-rooted OTel plugin with default settings: X-Ray context extraction, MDC enabled, root span * named {@code "Workflow"}. @@ -221,18 +237,8 @@ public void onInvocationStart(InvocationInfo info) { var sampled = OtelPluginSupport.isSampled(decision); var execCtx = ExecutionTraceContext.resolve(extracted, canonicalTraceId, arn(), idGenerator, () -> sampled); executionTrace = new ExecutionTrace(canonicalTraceId, execCtx.traceFlags()); - - // Workflow root span — parented onto the execution ancestor so it joins the execution trace, with a - // deterministic span ID from the ARN. Recreated every invocation with the same ID so it is exported once as a - // single logical span (on the terminal invocation only). Its start time is the backend execution start time. - var workflowSpanBuilder = tracer.spanBuilder(workflowSpanName) - .setSpanKind(SpanKind.INTERNAL) - .setParent(withDurableDecision(Context.root().with(Span.wrap(execCtx.executionAncestor())))) - .setAttribute(DURABLE_EXECUTION_ARN, info.durableExecutionArn()) - .setStartTimestamp(info.executionStartTime()); - var workflowSpanId = idGenerator.generateWorkflowSpanId(info.durableExecutionArn()); - // Force the span ID only; the trace ID comes from the parent so the Workflow span joins the execution trace. - workflowSpan = startDurableSpan(workflowSpanBuilder, null, workflowSpanId); + executionAncestor = execCtx.executionAncestor(); + executionStartTime = info.executionStartTime(); // Invocation span — child of the ambient Lambda span when it is on the execution trace, otherwise a child of // the execution ancestor so it stays within the same trace. @@ -249,6 +255,13 @@ public void onInvocationStart(InvocationInfo info) { invocationSpan = startDurableSpan(spanBuilder); + // Defer the recording Workflow span until terminal completion. The placeholder uses the Invocation span's + // resolved sampling metadata so operation parents/links match the span that is eventually exported. + var workflowSpanId = idGenerator.generateWorkflowSpanId(info.durableExecutionArn()); + var invocationContext = invocationSpan.getSpanContext(); + workflowSpanContext = SpanContext.create( + canonicalTraceId, workflowSpanId, invocationContext.getTraceFlags(), invocationContext.getTraceState()); + // Inject MDC on the handler thread so handler-level logs (between steps) have trace context. if (enableMdc) { MDC.put( @@ -270,20 +283,19 @@ public void onInvocationEnd(InvocationEndInfo info) { MdcSpanEnricher.clear(); } - // Reset per-invocation operation state WITHOUT ending open operation spans. Matching the JS/Python - // ExecutionOtelPlugin, an operation span is only ended in onOperationEnd. An operation still open when the - // invocation suspends is left un-exported here and is re-materialized once (with its deterministic span ID, - // plus a link to the invocation that completes it) when onOperationEnd fires in a later invocation. - operationSpans.clear(); + // Drop placeholder state. Open operations have no recording span to abandon. operationContexts.clear(); + operationStartTimes.clear(); - // Defensively close any lingering attempt scopes so OTel context is not leaked on worker threads (normally - // every onUserFunctionStart is paired with onUserFunctionEnd within the invocation). The attempt spans - // themselves are left un-ended rather than force-ended, consistent with not ending open spans here. + // Release OTel context on worker threads, then end any attempt spans still open so no recording span is + // abandoned. Attempt spans normally start and end within one user-function call, so this is a safeguard. for (var scope : attemptScopes.values()) { scope.close(); } attemptScopes.clear(); + for (var span : attemptSpans.values()) { + span.end(); + } attemptSpans.clear(); // End the invocation span every invocation. @@ -294,30 +306,36 @@ public void onInvocationEnd(InvocationEndInfo info) { invocationSpan.end(); invocationSpan = null; } - samplingIntent = null; - // End the Workflow span only on a terminal status, so it is exported exactly once per execution. - if (workflowSpan != null) { - if (isTerminal(info)) { - workflowSpan.setAttribute( - DURABLE_EXECUTION_STATUS, info.invocationStatus().name()); - switch (info.invocationStatus()) { - case FAILED -> { - var message = info.executionError() != null - ? info.executionError().getMessage() - : null; - workflowSpan.setStatus(StatusCode.ERROR, message); - if (info.executionError() != null) { - workflowSpan.recordException(info.executionError()); - } + // Materialize the Workflow span only on terminal status. + if (isTerminal(info) && workflowSpanContext != null && executionAncestor != null) { + var workflowSpanBuilder = tracer.spanBuilder(workflowSpanName) + .setSpanKind(SpanKind.INTERNAL) + .setParent(withDurableDecision(Context.root().with(Span.wrap(executionAncestor)))) + .setAttribute(DURABLE_EXECUTION_ARN, durableExecutionArn) + .setAttribute( + DURABLE_EXECUTION_STATUS, info.invocationStatus().name()) + .setStartTimestamp(executionStartTime != null ? executionStartTime : Instant.now()); + // Force only the deterministic span ID; the parent supplies the execution trace ID. + var workflowSpan = startDurableSpan(workflowSpanBuilder, null, workflowSpanContext.getSpanId()); + switch (info.invocationStatus()) { + case FAILED -> { + var message = info.executionError() != null + ? info.executionError().getMessage() + : null; + workflowSpan.setStatus(StatusCode.ERROR, message); + if (info.executionError() != null) { + workflowSpan.recordException(info.executionError()); } - default -> workflowSpan.setStatus(StatusCode.OK); // SUCCEEDED } - workflowSpan.end(); + default -> workflowSpan.setStatus(StatusCode.OK); // SUCCEEDED } - // Non-terminal (PENDING/RETRYING): leave the Workflow span un-ended (not exported this invocation). - workflowSpan = null; + workflowSpan.end(); } + workflowSpanContext = null; + executionAncestor = null; + executionStartTime = null; + samplingIntent = null; // Flush spans before Lambda freezes if (sdkTracerProvider != null) { @@ -335,6 +353,30 @@ public void onOperationStart(OperationInfo info) { if (!tracingEnabled) return; if (info.id() == null) return; + // Retain only a deterministic placeholder. Its flags/state come from the Invocation span's resolved sampling + // decision because CONTEXT operations make this placeholder current. + var trace = executionTrace; + var spanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, info.id()); + operationContexts.put( + info.id(), SpanContext.create(trace.traceId(), spanId, effectiveTraceFlags(), effectiveTraceState())); + + // Retain the start time for onOperationEnd, which may receive none (virtual FLAT map/parallel operations). + if (info.startTimestamp() != null) { + operationStartTimes.put(info.id(), info.startTimestamp()); + } + } + + @Override + public void onOperationEnd(OperationEndInfo info) { + if (!tracingEnabled) return; + if (info.id() == null) return; + + // Start and end the operation's single span here, using its deterministic span ID and linking to the + // invocation that completed it. This covers operations that ran in this invocation and ones resumed from an + // earlier one, and it is the only place an operation span is created — so none is ever left open. + operationContexts.remove(info.id()); + var capturedStart = operationStartTimes.remove(info.id()); + var parentContext = resolveParentContext(info.parentId()); var spanBuilder = tracer.spanBuilder(spanName(info.type(), info.subType(), info.name())) @@ -344,8 +386,12 @@ public void onOperationStart(OperationInfo info) { .setAttribute(DURABLE_OPERATION_TYPE, info.type()); addInvocationLink(spanBuilder); - if (info.startTimestamp() != null) { - spanBuilder.setStartTimestamp(info.startTimestamp()); + // Use the earliest known start so the operation span never starts after its own attempt/child spans (which were + // created earlier, at onUserFunctionStart/onOperationStart). onOperationEnd's start timestamp can be a later + // re-observed value than the start captured at onOperationStart, so take the minimum of the two. + var startTimestamp = earliest(capturedStart, info.startTimestamp()); + if (startTimestamp != null) { + spanBuilder.setStartTimestamp(startTimestamp); } if (info.name() != null) { spanBuilder.setAttribute(DURABLE_OPERATION_NAME, info.name()); @@ -357,86 +403,25 @@ public void onOperationStart(OperationInfo info) { var operationSpanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, info.id()); var span = startDurableSpan(spanBuilder, null, operationSpanId); - // Store the open span — will be ended in onOperationEnd or onInvocationEnd - operationSpans.put(info.id(), span); - operationContexts.put(info.id(), span.getSpanContext()); - } - - @Override - public void onOperationEnd(OperationEndInfo info) { - if (!tracingEnabled) return; - if (info.id() == null) return; - - var span = operationSpans.remove(info.id()); - - if (span != null) { - // Operation was started in this invocation — end normally - if (info.status() != null) { - span.setAttribute(DURABLE_OPERATION_STATUS, info.status()); - } - // Total attempts for retriable operations (STEP, WAIT_FOR_CONDITION) — emitted only at end. - if (info.attempt() != null) { - span.setAttribute(DURABLE_ATTEMPT_NUMBER, info.attempt().longValue()); - } - if (info.error() != null) { - span.setStatus(StatusCode.ERROR, info.error().getMessage()); - span.recordException(info.error()); - } else if ("SUCCEEDED".equals(info.status()) || info.status() == null) { - // Only stamp OK on genuine success. onOperationEnd fires for every terminal status, and - // extractErrorFromOperation returns null for CANCELLED (always) and for FAILED/TIMED_OUT/STOPPED - // with no attached error object — those carry a non-null, non-SUCCEEDED status and must stay UNSET. - // A null status is a successful statusless virtual (FLAT CONTEXT) operation, which is OK. - span.setStatus(StatusCode.OK); - } - endSpan(span, info.endTimestamp()); - } else { - // Operation completed between invocations: its onOperationStart ran in a prior invocation, whose - // in-memory span was dropped un-exported at that invocation's end. Emit the operation's single span - // now, using its deterministic span ID (stable across the execution), plus a link to the invocation - // that completed it. - operationContexts.remove(info.id()); - - var parentContext = resolveParentContext(info.parentId()); - - var spanBuilder = tracer.spanBuilder(spanName(info.type(), info.subType(), info.name())) - .setParent(parentContext) - .setAttribute(DURABLE_EXECUTION_ARN, durableExecutionArn) - .setAttribute(DURABLE_OPERATION_ID, info.id()) - .setAttribute(DURABLE_OPERATION_TYPE, info.type()); - addInvocationLink(spanBuilder); - - if (info.startTimestamp() != null) { - spanBuilder.setStartTimestamp(info.startTimestamp()); - } - if (info.name() != null) { - spanBuilder.setAttribute(DURABLE_OPERATION_NAME, info.name()); - } - if (info.subType() != null) { - spanBuilder.setAttribute(DURABLE_OPERATION_SUBTYPE, info.subType()); - } - - var operationSpanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, info.id()); - var continuationSpan = startDurableSpan(spanBuilder, null, operationSpanId); - - if (info.status() != null) { - continuationSpan.setAttribute(DURABLE_OPERATION_STATUS, info.status()); - } - // Total attempts for retriable operations (STEP, WAIT_FOR_CONDITION) — emitted only at end. - if (info.attempt() != null) { - continuationSpan.setAttribute( - DURABLE_ATTEMPT_NUMBER, info.attempt().longValue()); - } - if (info.error() != null) { - continuationSpan.setStatus(StatusCode.ERROR, info.error().getMessage()); - continuationSpan.recordException(info.error()); - } else if ("SUCCEEDED".equals(info.status()) || info.status() == null) { - // See onOperationEnd (this-invocation branch): only genuine success (or a successful statusless - // virtual operation) is OK; error-less non-success statuses stay UNSET. - continuationSpan.setStatus(StatusCode.OK); - } - - endSpan(continuationSpan, info.endTimestamp()); + if (info.status() != null) { + span.setAttribute(DURABLE_OPERATION_STATUS, info.status()); + } + // Total attempts for retriable operations (STEP, WAIT_FOR_CONDITION) — emitted only at end. + if (info.attempt() != null) { + span.setAttribute(DURABLE_ATTEMPT_NUMBER, info.attempt().longValue()); + } + if (info.error() != null) { + span.setStatus(StatusCode.ERROR, info.error().getMessage()); + span.recordException(info.error()); + } else if ("SUCCEEDED".equals(info.status()) || info.status() == null) { + // Only stamp OK on genuine success. onOperationEnd fires for every terminal status, and + // extractErrorFromOperation returns null for CANCELLED (always) and for FAILED/TIMED_OUT/STOPPED with no + // attached error object — those carry a non-null, non-SUCCEEDED status and must stay UNSET. A null status + // is a successful statusless virtual (FLAT CONTEXT) operation, which is OK. + span.setStatus(StatusCode.OK); } + + endSpan(span, info.endTimestamp()); } // ─── User function hooks ───────────────────────────────────────────── @@ -446,11 +431,12 @@ public void onUserFunctionStart(UserFunctionStartInfo info) { if (!tracingEnabled) return; // Skip attempt spans for CONTEXT operations — they are a scoping construct, not a retriable unit of work. Still - // make the operation span current so auto-instrumented calls become children. + // make the operation's context current so auto-instrumented calls become children of the (deferred) operation + // span. The context is non-recording until onOperationEnd, which is enough for parent propagation. if ("CONTEXT".equals(info.type())) { - var operationSpan = operationSpans.get(info.id()); - if (operationSpan != null) { - var scope = operationSpan.makeCurrent(); + var operationContext = operationContexts.get(info.id()); + if (operationContext != null) { + var scope = Span.wrap(operationContext).makeCurrent(); var key = attemptKey(info.id(), info.attempt()); attemptScopes.put(key, scope); } @@ -627,13 +613,14 @@ private Context resolveParentContext(String parentId) { // Parent operation from a prior invocation — non-recording placeholder with its deterministic ID. var deterministicParentSpanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, parentId); var placeholderContext = SpanContext.create( - trace.traceId(), deterministicParentSpanId, trace.flags(), TraceState.getDefault()); + trace.traceId(), deterministicParentSpanId, effectiveTraceFlags(), effectiveTraceState()); return withDurableDecision(Context.current().with(Span.wrap(placeholderContext))); } } - // No usable parent operation — hang off the Workflow root span. - if (workflowSpan != null) { - return withDurableDecision(Context.current().with(workflowSpan)); + // No usable parent operation — hang off the deferred Workflow span via its deterministic context. + var workflowContext = workflowSpanContext; + if (workflowContext != null) { + return withDurableDecision(Context.current().with(Span.wrap(workflowContext))); } return withDurableDecision(Context.current()); } @@ -648,6 +635,20 @@ private Context withDurableDecision(Context context) { return intent != null ? DurableSamplingDecision.store(context, intent) : context; } + private TraceFlags effectiveTraceFlags() { + var invocation = invocationSpan; + if (invocation != null) { + return invocation.getSpanContext().getTraceFlags(); + } + var trace = executionTrace; + return trace != null ? trace.flags() : TraceFlags.getDefault(); + } + + private TraceState effectiveTraceState() { + var invocation = invocationSpan; + return invocation != null ? invocation.getSpanContext().getTraceState() : TraceState.getDefault(); + } + /** * Starts a durable span with the execution's sampling intent published on the current thread for the duration of * the sampler call, so {@link DurableSampler} applies it even when the plugin and the agent-installed sampler run @@ -683,6 +684,17 @@ private static void endSpan(Span span, Instant endTimestamp) { } } + /** Returns the earlier of two timestamps, ignoring nulls; null only when both are null. */ + private static Instant earliest(Instant a, Instant b) { + if (a == null) { + return b; + } + if (b == null) { + return a; + } + return a.isBefore(b) ? a : b; + } + private static String spanName(String type, String subType, String name) { if (name != null) { return name; 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 307ec84ac..e89804bdc 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 @@ -39,9 +39,10 @@ *

Creates spans at these levels: * *