diff --git a/.github/workflows/otel-conformance-tests.yml b/.github/workflows/otel-conformance-tests.yml index 1e69aacd7..a6e3379c9 100644 --- a/.github/workflows/otel-conformance-tests.yml +++ b/.github/workflows/otel-conformance-tests.yml @@ -61,13 +61,13 @@ jobs: actions: write contents: read id-token: write - uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@91740c98b496409fa9f1bb8e8e6c329ca8b0185f + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@f9998f305f1e26423baf0c58148a3bd69120d5ef with: language: java resource_prefix: j sdk_repository: aws/aws-durable-execution-sdk-java sdk_ref: ${{ github.event.pull_request.head.sha || github.sha }} - conformance_test_ref: ${{ inputs.conformance_test_ref || '91740c98b496409fa9f1bb8e8e6c329ca8b0185f' }} + conformance_test_ref: ${{ inputs.conformance_test_ref || 'f9998f305f1e26423baf0c58148a3bd69120d5ef' }} checkout_sdk: true # Build the handlers from this repo's checked-out module instead of the conformance repo's # bundled examples/java. Path is relative to the conformance workspace where the SDK is diff --git a/otel-plugin/README.md b/otel-plugin/README.md index 107354f57..7104ce5a4 100644 --- a/otel-plugin/README.md +++ b/otel-plugin/README.md @@ -1,11 +1,11 @@ # AWS Durable Execution SDK - OpenTelemetry Plugin -OpenTelemetry instrumentation plugin for the AWS Lambda Durable Execution SDK for Java. Emits a deterministic Workflow trace for durable-execution correlation while keeping each Invocation span in the ambient Lambda trace. +OpenTelemetry instrumentation plugin for the AWS Lambda Durable Execution SDK for Java. Anchors every durable execution on one trace so the Workflow span and its per-invocation spans stay correlated, joining the propagated backend trace when one is present. ## Features -- **Deterministic Workflow Traces**: Workflow trace IDs are derived from the execution start time and ARN; stable span IDs are derived from the ARN -- **Ambient Invocation Traces**: Invocation spans inherit the active Lambda/X-Ray context, or receive a fresh provider-generated root trace ID +- **Backend-parented execution trace**: The Workflow span parents onto the execution ancestor resolved at invocation start — a propagated remote context, or a synthetic execution root — for one trace ID that is stable across all invocations, plus a stable span ID derived from the ARN +- **Ambient Invocation Traces**: Invocation spans inherit the active Lambda/X-Ray context, or join the execution ancestor so they stay on the execution trace - **Scoped ID Generation**: Unrelated instrumentation scopes retain their provider's normal root trace ID generation - **Span-per-Operation**: Each durable operation (step, wait, map, etc.) gets its own span with accurate timing - **Attempt Spans**: Each user function execution (step attempt, child context run) gets a span, including retries @@ -92,7 +92,7 @@ Build the plugin layer ZIP with the OTel plugin JAR at `java/lib/aws-durable-exe ### 2. AWS X-Ray Active Tracing -Enable active tracing on your Lambda function so the `_X_AMZN_TRACE_ID` environment variable is populated at invocation time. The plugin uses this header to parent Invocation spans to the ambient Lambda/X-Ray trace. The Workflow trace remains independent and deterministic. +Enable active tracing on your Lambda function so the `_X_AMZN_TRACE_ID` environment variable is populated at invocation time. The plugin uses this header both to parent Invocation spans to the ambient Lambda/X-Ray trace and to anchor the execution trace on the propagated context when it carries a complete parent and an explicit sampling decision. **AWS Console:** Lambda > Configuration > Monitoring and operations tools > Active tracing > Enable @@ -157,29 +157,40 @@ The function's execution role needs the `AWSXRayDaemonWriteAccess` managed polic ## Trace Structure -With `InvocationOtelPlugin`, the plugin creates two correlated traces: +The whole execution shares one trace, anchored at the execution ancestor resolved at invocation start. When the backend propagates a valid remote server span (`Root` and `Parent`), that span is the ancestor and the Workflow and Invocation spans nest under it, alongside the ambient Lambda spans on the same trace: ``` -Workflow trace: -Workflow (deterministic trace/span IDs, exported once) - -Ambient invocation trace: -Lambda/X-Ray parent -└── Invocation - ├── fetch-data - │ └── fetch-data attempt 1 - ├── cool-down - └── process - └── process attempt 1 +Remote backend server span (Root / Parent) +├── Workflow (stable span ID, exported once) +├── Ambient Lambda span 1 +│ └── Invocation 1 +├── Ambient Lambda span 2 +│ └── Invocation 2 +└── Invocation N (direct child when no same-trace ambient span exists) ``` -- **Workflow span** — one logical root per durable execution with a deterministic, X-Ray-compatible trace ID derived from the execution start time and ARN, plus a stable span ID derived from the ARN. Exported only on the terminal invocation (SUCCEEDED/FAILED). -- **Invocation span** — one per Lambda invocation, parented to ambient context when available +When no valid remote parent can be constructed, a synthetic execution root anchors the trace instead and both spans parent onto it: + +``` +Synthetic execution root +├── Workflow +├── Invocation 1 +├── Invocation 2 +└── Invocation N +``` + +- **Execution ancestor** — the common parent both the Workflow and Invocation spans resolve onto. A valid remote server span (`Root` and `Parent`) is used directly, whether or not `Sampled` is present; only when a valid remote parent cannot be constructed does a synthetic execution root take its place. It is a non-recording context, not an exported span. +- **Workflow span** — one logical span per durable execution, joining the execution trace with a stable span ID derived from the ARN. Exported only on the terminal invocation (SUCCEEDED/FAILED). +- **Invocation span** — one per Lambda invocation, parented to the ambient span only when it is on the execution trace, otherwise to the execution ancestor - **Operation span** — one per durable operation, named after your step/wait names - **Attempt span** — one per user function execution (retries produce additional attempt spans) Operation and attempt spans link to the Workflow span. `ExecutionOtelPlugin` reverses that relationship: operations are children of Workflow and link to the current Invocation span. +### Sampling + +`Sampled=1` and `Sampled=0` in the propagated header are authoritative and preserved. When `Sampled` is absent or unusable, the sampled flag is left unset rather than treated as `Sampled=0`: a parent-based sampler will then treat the trace as not sampled, while a directly configured trace-ID-ratio sampler can still decide from the canonical trace ID (a decision that stays stable across reinvocations because the trace ID is stable). When there is no valid remote parent, the configured root sampler decides. + ## Span Attributes ### Invocation Span @@ -305,7 +316,7 @@ The plugin's spans do not appear as nested subsegments of the Lambda platform se ### Workflow Span -The Workflow span appears in a separate deterministic trace because it uses `setNoParent()`. Invocation spans remain in the ambient Lambda/X-Ray trace. Links correlate durable operations with the other trace. +The Workflow span joins the execution trace by parenting onto the execution ancestor: the propagated remote server span when one is valid, otherwise a synthetic execution root. Either way it shares the execution trace ID and keeps its stable, ARN-derived span ID. ## Verification @@ -314,8 +325,8 @@ After deploying your function with the plugin configured: 1. **Invoke your durable function** — trigger at least one execution that includes multiple steps or a wait/resume cycle. 2. **Check CloudWatch console** — Navigate to CloudWatch > Traces. Enable "Group by nodes" to see: - - A deterministic Workflow trace covering the entire execution - - Ambient Lambda traces containing one Invocation span per Lambda invocation + - One execution trace covering the whole execution, with the Workflow span and each Invocation span sharing its trace ID + - One Invocation span per Lambda invocation - Child spans for each durable operation (named after your step names) - Links between durable Workflow/operation spans and Invocation spans diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DeterministicIdGenerator.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DeterministicIdGenerator.java index c81b6c706..88e2cb821 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DeterministicIdGenerator.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DeterministicIdGenerator.java @@ -129,7 +129,22 @@ String generateTraceIdForExecution(String arn, Instant executionStartTime) { } String generateWorkflowSpanId(String arn) { - var seed = "workflow:" + (arn != null ? arn : ""); + return deterministicSpanId("workflow:" + (arn != null ? arn : "")); + } + + /** + * Generates the deterministic span ID for the synthetic execution root from the execution ARN, using a seed + * namespace distinct from the Workflow and operation span IDs. Stable across reinvocations so the synthetic root is + * the same common ancestor every invocation. + * + * @param arn the durable execution ARN + * @return a deterministic 16-char hex span ID + */ + String generateExecutionRootSpanId(String arn) { + return deterministicSpanId("execution-root:" + (arn != null ? arn : "")); + } + + private static String deterministicSpanId(String seed) { var spanId = sha256(seed).substring(0, 16); if (spanId.equals("0000000000000000")) { spanId = "0000000000000001"; 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 8170ce008..c2fb8b003 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 @@ -5,6 +5,7 @@ import static software.amazon.lambda.durable.otel.SpanAttributes.*; import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.SpanContext; @@ -38,14 +39,17 @@ * durable-execution hierarchy: * * @@ -56,11 +60,12 @@ * it. Both plugins share {@link DeterministicIdGenerator}, {@link ContextExtractor}, {@link SpanAttributes}, and * {@link MdcSpanEnricher}. * - *

The Workflow trace ID is derived from the execution start time and ARN, and is independent of the ambient - * Lambda/X-Ray trace. Invocation spans inherit the active ambient context, or extracted upstream context as a fallback. - * When using {@link #ExecutionOtelPlugin()}, the plugin resolves the global provider at invocation start. If the - * OpenTelemetry Java agent is not initialized yet, telemetry is disabled for that entire invocation and provider - * resolution is retried on the next invocation. + *

The Workflow and Invocation spans share one execution trace, anchored at the execution ancestor resolved at + * invocation start: a valid propagated remote server span becomes that ancestor directly, otherwise a synthetic + * execution root anchors the trace. The trace ID is stable across invocations of the same execution. When using + * {@link #ExecutionOtelPlugin()}, the plugin resolves the global provider at invocation start. If the OpenTelemetry + * Java agent is not initialized yet, telemetry is disabled for that entire invocation and provider resolution is + * retried on the next invocation. * *

Status mapping (parity with the Python/JS references): * @@ -91,21 +96,30 @@ public class ExecutionOtelPlugin implements DurableExecutionPlugin { // Per-invocation state private volatile boolean tracingEnabled; - private volatile Span workflowSpan; private volatile Span invocationSpan; private volatile String durableExecutionArn; private volatile String workflowTraceId; + private volatile TraceFlags executionTraceFlags; - // 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"}. @@ -183,41 +197,38 @@ public void onInvocationStart(InvocationInfo info) { this.durableExecutionArn = info.durableExecutionArn(); - // Prefer the active Java-agent span, then fall back to explicitly extracted upstream context. - var invocationParent = extractCurrentSpanContext(); - if (invocationParent == null) { - invocationParent = contextExtractor.extract(); - } - - // Workflow root span — deterministic span ID from the ARN, no parent. 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 execution start time from the backend. - var workflowSpanBuilder = tracer.spanBuilder(workflowSpanName) - .setSpanKind(SpanKind.INTERNAL) - .setNoParent() - .setAttribute(DURABLE_EXECUTION_ARN, info.durableExecutionArn()) - .setStartTimestamp(info.executionStartTime()); - workflowTraceId = - idGenerator.generateTraceIdForExecution(info.durableExecutionArn(), info.executionStartTime()); + // Resolve the one execution ancestor both spans parent onto, so they share a trace and a sampling decision. + var extracted = contextExtractor.extract(); + var canonicalTraceId = + ExecutionTraceContext.canonicalTraceId(extracted, arn(), info.executionStartTime(), idGenerator); + var execCtx = ExecutionTraceContext.resolve( + extracted, + canonicalTraceId, + arn(), + idGenerator, + () -> OtelPluginSupport.rootSampled( + sdkTracerProvider, + canonicalTraceId, + workflowSpanName, + Attributes.of(DURABLE_EXECUTION_ARN, arn()))); + workflowTraceId = canonicalTraceId; + executionTraceFlags = execCtx.traceFlags(); + executionAncestor = execCtx.executionAncestor(); + executionStartTime = info.executionStartTime(); + + // Workflow span — deferred. Between invocations it exists only as this deterministic context (same trace and + // span ID every invocation) that operations parent onto. The recording span is started and ended together on + // the terminal invocation, so it is never left open on a non-terminal status. var workflowSpanId = idGenerator.generateWorkflowSpanId(info.durableExecutionArn()); - workflowSpan = idGenerator.startSpan(workflowSpanBuilder, workflowTraceId, workflowSpanId); - - Context parentContext; - if (invocationParent != null && invocationParent.parentSpanId() != null) { - var parentSpanContext = SpanContext.createFromRemoteParent( - invocationParent.traceId(), - invocationParent.parentSpanId(), - TraceFlags.getSampled(), - TraceState.getDefault()); - parentContext = Context.root().with(Span.wrap(parentSpanContext)); - } else { - parentContext = Context.root(); - } + workflowSpanContext = + SpanContext.create(canonicalTraceId, workflowSpanId, execCtx.traceFlags(), TraceState.getDefault()); - // Invocation span — child of the ambient Lambda span when available, otherwise a root. + // 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. + var invocationParent = invocationParentContext(execCtx, canonicalTraceId); var spanBuilder = tracer.spanBuilder("Invocation") .setSpanKind(SpanKind.INTERNAL) - .setParent(parentContext) + .setParent(invocationParent) .setAttribute(DURABLE_EXECUTION_ARN, info.durableExecutionArn()) .setAttribute(DURABLE_FIRST_INVOCATION, info.isFirstInvocation()); @@ -248,20 +259,21 @@ 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 the deterministic operation state. An operation still open here has no recording span to abandon — its + // single span is created and ended by the invocation that completes it (onOperationEnd), keyed by the stable + // deterministic span ID so it remains one logical span across invocations. 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. @@ -273,28 +285,37 @@ public void onInvocationEnd(InvocationEndInfo info) { invocationSpan = 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 on a terminal status only: started and ended in a single call, parented onto + // the execution ancestor, so it is exported exactly once per execution and never left open. Non-terminal + // statuses (PENDING/RETRYING) export no Workflow span this invocation. + if (isTerminal(info) && workflowSpanContext != null && executionAncestor != null) { + var workflowSpanBuilder = tracer.spanBuilder(workflowSpanName) + .setSpanKind(SpanKind.INTERNAL) + .setParent(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 the span ID only; the trace ID comes from the parent so the Workflow span joins the execution + // trace. The deterministic span ID matches the context operations parented onto. + var workflowSpan = idGenerator.startSpan(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; // Flush spans before Lambda freezes if (sdkTracerProvider != null) { @@ -312,6 +333,29 @@ public void onOperationStart(OperationInfo info) { if (!tracingEnabled) return; if (info.id() == null) return; + // Retain only the deterministic context so children and attempts can parent onto the operation. The recording + // span is created and ended in onOperationEnd, so a suspended operation never leaves an open span and a + // suspended-then-resumed operation stitches into a single logical span. + var spanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, info.id()); + operationContexts.put(info.id(), SpanContext.create(workflowTraceId, spanId, flags(), TraceState.getDefault())); + + // 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())) @@ -321,8 +365,10 @@ public void onOperationStart(OperationInfo info) { .setAttribute(DURABLE_OPERATION_TYPE, info.type()); addInvocationLink(spanBuilder); - if (info.startTimestamp() != null) { - spanBuilder.setStartTimestamp(info.startTimestamp()); + // Prefer the end info's start timestamp, falling back to the one captured at operation start. + var startTimestamp = info.startTimestamp() != null ? info.startTimestamp() : capturedStart; + if (startTimestamp != null) { + spanBuilder.setStartTimestamp(startTimestamp); } if (info.name() != null) { spanBuilder.setAttribute(DURABLE_OPERATION_NAME, info.name()); @@ -334,86 +380,25 @@ public void onOperationStart(OperationInfo info) { var operationSpanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, info.id()); var span = idGenerator.startSpan(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 = idGenerator.startSpan(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 ───────────────────────────────────────────── @@ -423,11 +408,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); } @@ -564,6 +550,27 @@ private static boolean isTerminal(InvocationEndInfo info) { }; } + private String arn() { + return durableExecutionArn; + } + + private TraceFlags flags() { + var resolved = executionTraceFlags; + return resolved != null ? resolved : TraceFlags.getDefault(); + } + + /** + * The parent context for the Invocation span: the active ambient span when it is already on the execution trace, + * otherwise the execution ancestor so the Invocation span stays within the same trace. + */ + private Context invocationParentContext(ExecutionTraceContext execCtx, String canonicalTraceId) { + var ambient = Span.current().getSpanContext(); + if (ambient.isValid() && ambient.getTraceId().equals(canonicalTraceId)) { + return Context.root().with(Span.current()); + } + return Context.root().with(Span.wrap(execCtx.executionAncestor())); + } + /** Adds a link to the current invocation span, if one exists, for correlation. */ private void addInvocationLink(SpanBuilder spanBuilder) { var currentInvocationSpan = invocationSpan; @@ -580,13 +587,14 @@ private Context resolveParentContext(String parentId) { } // Parent operation from a prior invocation — create a non-recording placeholder with its deterministic ID. var deterministicParentSpanId = idGenerator.generateSpanIdForOperation(durableExecutionArn, parentId); - var placeholderContext = SpanContext.create( - workflowTraceId, deterministicParentSpanId, TraceFlags.getSampled(), TraceState.getDefault()); + var placeholderContext = + SpanContext.create(workflowTraceId, deterministicParentSpanId, flags(), TraceState.getDefault()); return Context.current().with(Span.wrap(placeholderContext)); } - // No parent operation — hang off the Workflow root span. - if (workflowSpan != null) { - return Context.current().with(workflowSpan); + // No parent operation — hang off the deferred Workflow span via its deterministic context. + var workflowContext = workflowSpanContext; + if (workflowContext != null) { + return Context.current().with(Span.wrap(workflowContext)); } return Context.current(); } @@ -617,8 +625,4 @@ private static String attemptSpanName(String type, String subType, String name, private static String attemptKey(String operationId, Integer attempt) { return operationId + "-" + (attempt != null ? attempt : "ctx"); } - - private static ExtractedContext extractCurrentSpanContext() { - return OtelPluginSupport.extractCurrentSpanContext(); - } } diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionTraceContext.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionTraceContext.java new file mode 100644 index 000000000..615021ea4 --- /dev/null +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExecutionTraceContext.java @@ -0,0 +1,94 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import java.time.Instant; +import java.util.function.BooleanSupplier; + +/** + * The per-execution trace context resolved once at invocation start, shared by both plugins. + * + *

It selects the common execution ancestor that the Workflow and Invocation spans parent onto, so they share one + * trace: + * + *

+ * + *

Sampling flags: an explicit upstream decision is preserved; when it is absent the flag is left unset for a remote + * parent (the configured sampler decides at span start), while for a synthetic root the supplied decision is applied. + * + *

The ancestor is a non-recording context: it is either the external backend server span or a synthetic root the SDK + * does not export. + * + * @param executionAncestor the common parent context for the Workflow and Invocation spans + */ +record ExecutionTraceContext(SpanContext executionAncestor) { + + String traceId() { + return executionAncestor.getTraceId(); + } + + TraceFlags traceFlags() { + return executionAncestor.getTraceFlags(); + } + + /** + * Resolves the execution trace context from the propagated header and the canonical trace ID. + * + * @param extracted the context parsed from the backend header, or null when none is present + * @param canonicalTraceId the trace ID to anchor a synthetic root on (the remote trace ID, or an ARN-derived one) + * @param arn the durable execution ARN + * @param idGenerator the deterministic ID generator + * @param rootSampled decides a synthetic root's sampled flag when the header carries no explicit decision + */ + static ExecutionTraceContext resolve( + ExtractedContext extracted, + String canonicalTraceId, + String arn, + DeterministicIdGenerator idGenerator, + BooleanSupplier rootSampled) { + + // A valid remote parent is the authoritative ancestor, regardless of whether Sampled is present. An absent + // upstream decision leaves the flag unset so the configured sampler applies at span start. + if (extracted != null && extracted.hasCompleteRemoteParent()) { + var flags = explicitFlags(extracted.sampling(), TraceFlags.getDefault()); + var remoteParent = SpanContext.createFromRemoteParent( + extracted.traceId(), extracted.parentSpanId(), flags, TraceState.getDefault()); + return new ExecutionTraceContext(remoteParent); + } + + // No valid remote parent: synthesize an execution root on the canonical trace. An explicit upstream decision is + // still preserved; otherwise the supplied root-sampling decision applies. + var sampling = extracted != null ? extracted.sampling() : ExtractedContext.Sampling.UNDECIDED; + var flags = + explicitFlags(sampling, rootSampled.getAsBoolean() ? TraceFlags.getSampled() : TraceFlags.getDefault()); + var syntheticRoot = SpanContext.create( + canonicalTraceId, idGenerator.generateExecutionRootSpanId(arn), flags, TraceState.getDefault()); + return new ExecutionTraceContext(syntheticRoot); + } + + /** The canonical trace ID: the remote trace ID when present, else one derived from the ARN and start time. */ + static String canonicalTraceId( + ExtractedContext extracted, String arn, Instant executionStartTime, DeterministicIdGenerator idGenerator) { + return extracted != null && extracted.traceId() != null + ? extracted.traceId() + : idGenerator.generateTraceIdForExecution(arn, executionStartTime); + } + + /** Flags for an explicit upstream decision, or {@code whenUndecided} when the upstream did not decide. */ + private static TraceFlags explicitFlags(ExtractedContext.Sampling sampling, TraceFlags whenUndecided) { + return switch (sampling) { + case SAMPLED -> TraceFlags.getSampled(); + case NOT_SAMPLED -> TraceFlags.getDefault(); + case UNDECIDED -> whenUndecided; + }; + } +} diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExtractedContext.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExtractedContext.java index d381732cb..544849424 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExtractedContext.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ExtractedContext.java @@ -5,10 +5,33 @@ /** * Trace context extracted from the Lambda runtime environment. * - *

Contains the trace ID (always present) and an optional parent span ID used to parent an Invocation span to ambient - * Lambda/X-Ray context. + *

Carries the trace ID, an optional parent span ID (the remote durable backend server span), and the upstream + * sampling decision. The sampling decision is tri-state: an absent or unusable value is {@link Sampling#UNDECIDED}, + * which is distinct from an explicit {@link Sampling#NOT_SAMPLED}. * - * @param traceId 32-character lowercase hex trace ID (OTel format, no dashes) - * @param parentSpanId 16-character lowercase hex parent span ID (may be null if no parent available) + * @param traceId 32-character lowercase hex trace ID (OTel format, no dashes), or null when no valid Root was present + * @param parentSpanId 16-character lowercase hex parent span ID, or null when no valid Parent was present + * @param sampling the upstream sampling decision */ -public record ExtractedContext(String traceId, String parentSpanId) {} +public record ExtractedContext(String traceId, String parentSpanId, Sampling sampling) { + + /** Upstream sampling decision carried by the propagated context. */ + public enum Sampling { + /** The upstream explicitly decided to sample (X-Ray {@code Sampled=1}). */ + SAMPLED, + /** The upstream explicitly decided not to sample (X-Ray {@code Sampled=0}). */ + NOT_SAMPLED, + /** No usable upstream decision; the configured sampler decides. */ + UNDECIDED + } + + /** Creates a context with an undecided sampling decision. */ + public ExtractedContext(String traceId, String parentSpanId) { + this(traceId, parentSpanId, Sampling.UNDECIDED); + } + + /** A complete remote context has both a trace ID and a parent span ID, so it can serve as a remote parent. */ + public boolean hasCompleteRemoteParent() { + return traceId != null && parentSpanId != null; + } +} 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 75df1f60b..1777c5b36 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 @@ -5,6 +5,7 @@ import static software.amazon.lambda.durable.otel.SpanAttributes.*; import io.opentelemetry.api.common.AttributeKey; +import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; import io.opentelemetry.api.trace.SpanBuilder; import io.opentelemetry.api.trace.SpanContext; @@ -38,28 +39,20 @@ *

Creates spans at these levels: * *

* - *

Workflow span behavior by provider: - * - *

- * - *

The Workflow trace ID is derived from the execution start time and ARN, and is independent of the ambient - * Lambda/X-Ray trace. Invocation spans inherit the active ambient context, or extracted upstream context as a fallback. + *

The Workflow span is parented onto the execution ancestor resolved at invocation start (the propagated remote + * server span when one is valid, otherwise a synthetic execution root), so it joins the execution trace with a trace ID + * stable across invocations. It serves as a correlation anchor: operation and attempt spans link to it while remaining + * parented to the per-invocation span. The Invocation span parents onto the same-trace ambient span when available, + * otherwise onto the execution ancestor so it stays on the execution trace. * *

Requires the ADOT Lambda Layer for trace export. Configure with: * @@ -76,8 +69,8 @@ * operation, attempt) do not appear as nested subsegments of the Lambda platform segment. This is a known limitation of * the OTLP-to-X-Ray conversion: the ADOT collector cannot attach OTLP-exported spans as subsegments of the Lambda * service's native X-Ray segment because that segment is created outside the OTLP pipeline. Use the "Group by nodes" - * view to inspect parent-child relationships within the ambient Invocation trace and the links to the independent - * Workflow trace. + * view to inspect parent-child relationships within the shared execution trace and the links between operation spans + * and the Workflow span. * *

Thread-safe: uses {@link ConcurrentHashMap} for span/scope storage since the SDK runs user code on multiple * threads. @@ -96,9 +89,21 @@ public class InvocationOtelPlugin implements DurableExecutionPlugin { // Per-invocation state private volatile boolean tracingEnabled; - private volatile Span workflowSpan; private volatile Span invocationSpan; private volatile String durableExecutionArn; + // Trace ID and flags of the execution trace, published together as one snapshot so readers never pair a trace ID + // with mismatched flags. + private volatile ExecutionTrace executionTrace; + + // Workflow span — deferred. Between invocations it exists only as this deterministic context, which operations + // link to; the recording span is started and ended in a single call on the terminal invocation, parented onto the + // execution ancestor, so it is never left open on a non-terminal status. + private volatile SpanContext workflowSpanContext; + private volatile SpanContext executionAncestor; + private volatile Instant executionStartTime; + + /** 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<>(); @@ -198,40 +203,36 @@ public void onInvocationStart(InvocationInfo info) { this.durableExecutionArn = info.durableExecutionArn(); - // Prefer the active Java-agent span, then fall back to explicitly extracted upstream context. - var invocationParent = extractCurrentSpanContext(); - if (invocationParent == null) { - invocationParent = contextExtractor.extract(); - } - - // Workflow root span — one logical span per durable execution, created unconditionally (independent of the - // X-Ray parent below). Deterministic span ID from the ARN so it is the same across invocations; exported once, - // on the terminal invocation. Operation and attempt spans link to it for execution-level correlation while - // remaining parented to the per-invocation span (this plugin stays invocation-rooted). - var workflowSpanBuilder = tracer.spanBuilder(workflowSpanName) - .setSpanKind(SpanKind.INTERNAL) - .setNoParent() - .setAttribute(DURABLE_EXECUTION_ARN, info.durableExecutionArn()) - .setStartTimestamp(info.executionStartTime()); - var workflowTraceId = - idGenerator.generateTraceIdForExecution(info.durableExecutionArn(), info.executionStartTime()); + var extracted = contextExtractor.extract(); + + // Resolve the execution ancestor the Workflow span parents onto so it joins the execution trace. + var canonicalTraceId = ExecutionTraceContext.canonicalTraceId( + extracted, info.durableExecutionArn(), info.executionStartTime(), idGenerator); + var execCtx = ExecutionTraceContext.resolve( + extracted, + canonicalTraceId, + info.durableExecutionArn(), + idGenerator, + () -> OtelPluginSupport.rootSampled( + sdkTracerProvider, + canonicalTraceId, + workflowSpanName, + Attributes.of(DURABLE_EXECUTION_ARN, info.durableExecutionArn()))); + executionTrace = new ExecutionTrace(canonicalTraceId, execCtx.traceFlags()); + executionAncestor = execCtx.executionAncestor(); + executionStartTime = info.executionStartTime(); + + // Workflow span — deferred. Between invocations it exists only as this deterministic context (same trace and + // span ID every invocation) that operation and attempt spans link to. The recording span is started and ended + // together on the terminal invocation, so it is never left open on a non-terminal status. Operations remain + // parented to the per-invocation span (this plugin stays invocation-rooted). var workflowSpanId = idGenerator.generateWorkflowSpanId(info.durableExecutionArn()); - workflowSpan = idGenerator.startSpan(workflowSpanBuilder, workflowTraceId, workflowSpanId); - - // Determine parent context for the invocation span. - Context parentContext; - if (invocationParent != null && invocationParent.parentSpanId() != null) { - // Reconstruct a remote parent from the extracted trace context (X-Ray header or current span). - // This connects plugin spans to the Lambda service's X-Ray segments. - var parentSpanContext = SpanContext.createFromRemoteParent( - invocationParent.traceId(), - invocationParent.parentSpanId(), - TraceFlags.getSampled(), - TraceState.getDefault()); - parentContext = Context.root().with(Span.wrap(parentSpanContext)); - } else { - parentContext = Context.root(); - } + workflowSpanContext = + SpanContext.create(canonicalTraceId, workflowSpanId, execCtx.traceFlags(), TraceState.getDefault()); + + // Invocation span parent — the same-trace ambient span when available, then the execution ancestor, so the + // Invocation span stays on the execution trace. + var parentContext = invocationParentContext(execCtx, canonicalTraceId); // Create an INTERNAL span for the invocation. var spanBuilder = tracer.spanBuilder("Invocation") @@ -301,28 +302,37 @@ public void onInvocationEnd(InvocationEndInfo info) { invocationSpan.end(); invocationSpan = null; - // End the Workflow span only on a terminal status, so it is exported exactly once per execution - // (SUCCEEDED -> OK, FAILED -> ERROR; non-terminal statuses leave it un-ended / not exported this invocation). - 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 on a terminal status only: started and ended in a single call, parented onto + // the execution ancestor, so it is exported exactly once per execution and never left open. Non-terminal + // statuses (PENDING/RETRYING) export no Workflow span this invocation. + if (isTerminal(info) && workflowSpanContext != null && executionAncestor != null) { + var workflowSpanBuilder = tracer.spanBuilder(workflowSpanName) + .setSpanKind(SpanKind.INTERNAL) + .setParent(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 the span ID only; the trace ID comes from the parent so the Workflow span joins the execution + // trace. The deterministic span ID matches the context operation and attempt spans linked to. + var workflowSpan = idGenerator.startSpan(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 } - workflowSpan = null; + workflowSpan.end(); } + workflowSpanContext = null; + executionAncestor = null; + executionStartTime = null; if (sdkTracerProvider != null) { // Flush spans before Lambda freezes @@ -359,6 +369,12 @@ public void onOperationStart(OperationInfo info) { spanBuilder.setAttribute(DURABLE_OPERATION_SUBTYPE, info.subType()); } + // On replay, this span is a distinct segment of an operation whose initial span ran in an earlier invocation; + // link back to that initial logical operation span. + if (info.isReplay()) { + addInitialOperationLink(spanBuilder, info.id()); + } + var span = info.isReplay() ? spanBuilder.startSpan() : idGenerator.startSpan( @@ -406,6 +422,8 @@ public void onOperationEnd(OperationEndInfo info) { .setAttribute(DURABLE_OPERATION_ID, info.id()) .setAttribute(DURABLE_OPERATION_TYPE, info.type()); addWorkflowLink(spanBuilder); + // This continuation segment completes an operation whose initial span ran earlier; link back to it. + addInitialOperationLink(spanBuilder, info.id()); if (info.name() != null) { spanBuilder.setAttribute(DURABLE_OPERATION_NAME, info.name()); @@ -589,6 +607,18 @@ private void endOpenSpansChildFirst() { operationContexts.clear(); } + /** + * The parent context for the Invocation span: the active ambient span when it is on the execution trace, otherwise + * the execution ancestor so the Invocation span stays within the same trace. + */ + private Context invocationParentContext(ExecutionTraceContext execCtx, String canonicalTraceId) { + var ambient = Span.current().getSpanContext(); + if (ambient.isValid() && ambient.getTraceId().equals(canonicalTraceId)) { + return Context.root().with(Span.current()); + } + return Context.root().with(Span.wrap(execCtx.executionAncestor())); + } + private Context resolveParentContext(String parentId) { if (parentId != null) { var parentSpanContext = operationContexts.get(parentId); @@ -603,14 +633,35 @@ private Context resolveParentContext(String parentId) { return Context.current(); } - /** Adds a link to the Workflow span, if one exists, for execution-level correlation. */ + /** + * Adds a link to the Workflow span, if one is set, for execution-level correlation. Uses the deterministic Workflow + * context (the recording span is deferred to the terminal invocation, but shares this span ID). + */ private void addWorkflowLink(SpanBuilder spanBuilder) { - var currentWorkflowSpan = workflowSpan; - if (currentWorkflowSpan != null) { - spanBuilder.addLink(currentWorkflowSpan.getSpanContext()); + var workflowContext = workflowSpanContext; + if (workflowContext != null) { + spanBuilder.addLink(workflowContext); } } + /** + * Links a continuation or replay operation span back to the initial logical operation span, whose ID is + * deterministic on the execution trace, so the segments of one logical operation stay correlated across + * invocations. + */ + private void addInitialOperationLink(SpanBuilder spanBuilder, String operationId) { + var trace = executionTrace; + if (trace == null || operationId == null) { + return; + } + var initial = SpanContext.create( + trace.traceId(), + idGenerator.generateSpanIdForOperation(durableExecutionArn, operationId), + trace.flags(), + TraceState.getDefault()); + spanBuilder.addLink(initial); + } + private static boolean isTerminal(InvocationEndInfo info) { return switch (info.invocationStatus()) { case SUCCEEDED, FAILED -> true; @@ -636,8 +687,4 @@ private static String attemptSpanName(String type, String subType, String name, private static String attemptKey(String operationId, Integer attempt) { return operationId + "-" + (attempt != null ? attempt : "ctx"); } - - private static ExtractedContext extractCurrentSpanContext() { - return OtelPluginSupport.extractCurrentSpanContext(); - } } diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginAutoConfigurationCustomizerProvider.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginAutoConfigurationCustomizerProvider.java index 74c7fc1d3..cbce7a40e 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginAutoConfigurationCustomizerProvider.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginAutoConfigurationCustomizerProvider.java @@ -5,7 +5,10 @@ import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer; import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider; -/** Wraps the Java agent's configured ID generator with scoped durable-execution overrides. */ +/** + * Wraps the Java agent's configured ID generator with scoped durable-execution overrides, and publishes the sampler's + * decision for root spans so the application can apply it when the propagated context carries no explicit decision. + */ public final class OtelPluginAutoConfigurationCustomizerProvider implements AutoConfigurationCustomizerProvider { @Override @@ -15,6 +18,11 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) { DeterministicIdGenerator.installOn(builder); return builder; }); + // The sampler is only reachable here, inside the agent. Publish its root-span decision and return it unchanged. + autoConfiguration.addSamplerCustomizer((sampler, config) -> { + WorkflowRootSamplingBridge.publish(sampler); + return sampler; + }); } @Override diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java index cac13d10f..8246c7917 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/OtelPluginSupport.java @@ -3,12 +3,16 @@ package software.amazon.lambda.durable.otel; import io.opentelemetry.api.GlobalOpenTelemetry; -import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; import io.opentelemetry.api.trace.Tracer; import io.opentelemetry.api.trace.TracerProvider; +import io.opentelemetry.context.Context; import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Collections; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -24,6 +28,36 @@ static DeterministicIdGenerator createDefaultIdGenerator() { return new DeterministicIdGenerator(); } + /** + * Decides whether a root span on the given trace would be sampled, used only when the propagated context carries no + * explicit decision. Queries the provider's sampler directly when reachable; otherwise uses the decision the agent + * published across the classloader boundary; otherwise treats the trace as sampled, matching the default provider + * configuration. + * + * @param sdkTracerProvider the resolved provider, or null when it is not visible to the application + * @param traceId the trace ID to evaluate + * @param spanName the span name passed to the sampler + * @param attributes the attributes the span is started with + */ + static boolean rootSampled( + SdkTracerProvider sdkTracerProvider, String traceId, String spanName, Attributes attributes) { + if (sdkTracerProvider == null) { + var published = WorkflowRootSamplingBridge.publishedRootSampled(); + return published == null || published; + } + return sdkTracerProvider + .getSampler() + .shouldSample( + Context.root(), + traceId, + spanName, + SpanKind.INTERNAL, + attributes, + Collections.emptyList()) + .getDecision() + == SamplingDecision.RECORD_AND_SAMPLE; + } + /** The tracer provider and tracer resolved from the global OpenTelemetry instance. */ record ProviderSetup(SdkTracerProvider sdkTracerProvider, Tracer tracer) {} @@ -71,15 +105,6 @@ static ProviderSetup tryResolveGlobalProvider(String instrumentationName, String getSdkTracerProviderForFlush(tracerProvider, pluginName), tracerProvider.get(instrumentationName)); } - /** Extracts trace context from the current OTel span (fallback when X-Ray header is unavailable). */ - static ExtractedContext extractCurrentSpanContext() { - var spanContext = Span.current().getSpanContext(); - if (!spanContext.isValid()) { - return null; - } - return new ExtractedContext(spanContext.getTraceId(), spanContext.getSpanId()); - } - /** Returns the SdkTracerProvider for flushing, or null if the provider is wrapped by the agent classloader. */ static SdkTracerProvider getSdkTracerProviderForFlush(TracerProvider tracerProvider, String pluginName) { if (tracerProvider instanceof SdkTracerProvider sdkTracerProvider) { diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/WorkflowRootSamplingBridge.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/WorkflowRootSamplingBridge.java new file mode 100644 index 000000000..097c02d60 --- /dev/null +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/WorkflowRootSamplingBridge.java @@ -0,0 +1,75 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import java.util.Collections; + +/** + * Carries the configured sampler's decision for root spans from the Java agent to the application. + * + *

The agent's tracer provider is often hidden behind a classloader wrapper, so the application cannot reach the + * sampler to ask how a root span would be sampled. The plugin's {@code AutoConfigurationCustomizerProvider} runs inside + * the agent, where the sampler is available, and publishes the answer through a system property that both sides read. + * + *

The execution being traced is unknown when the agent is configured, so the decision is published only when it is + * the same across trace IDs and attributes. Samplers that decide per trace ID or per attribute publish nothing, and + * callers treat the trace as sampled. This probe is best-effort: it cannot detect a sampler keyed on some other input + * that the fixed probes happen to share, so "publish nothing, treat as sampled" is the safe default. + * + *

Assumes a single agent-configured sampler per JVM, which holds for the Lambda execution environment: the property + * reflects the last {@link #publish} call, so a second, differently-configured provider in the same JVM would overwrite + * it. + */ +final class WorkflowRootSamplingBridge { + + private static final String ROOT_SAMPLED_PROPERTY = "software.amazon.lambda.durable.otel.rootSpansSampled"; + + private static final String PROBE_TRACE_ID_A = "0af7651916cd43dd8448eb211c80319c"; + private static final String PROBE_TRACE_ID_B = "f1e2d3c4b5a6978800112233445566aa"; + private static final Attributes PROBE_ATTRIBUTES = + Attributes.of(SpanAttributes.DURABLE_EXECUTION_ARN, "arn:aws:lambda:probe"); + + private WorkflowRootSamplingBridge() {} + + /** Probes the sampler for root spans and publishes the result when every probe agrees. */ + static void publish(Sampler sampler) { + var baseline = isRootSampled(sampler, PROBE_TRACE_ID_A, Attributes.empty()); + var otherTraceId = isRootSampled(sampler, PROBE_TRACE_ID_B, Attributes.empty()); + var withAttributes = isRootSampled(sampler, PROBE_TRACE_ID_A, PROBE_ATTRIBUTES); + if (baseline == otherTraceId && baseline == withAttributes) { + System.setProperty(ROOT_SAMPLED_PROPERTY, Boolean.toString(baseline)); + } else { + System.clearProperty(ROOT_SAMPLED_PROPERTY); + } + } + + /** + * Returns the published decision for root spans, or null when nothing was published because the agent did not + * install the plugin's customizer or the root decision is not the same for every execution. + */ + static Boolean publishedRootSampled() { + var value = System.getProperty(ROOT_SAMPLED_PROPERTY); + return value == null ? null : Boolean.valueOf(value); + } + + private static boolean isRootSampled(Sampler sampler, String traceId, Attributes attributes) { + return sampler.shouldSample( + Context.root(), + traceId, + "Workflow", + SpanKind.INTERNAL, + attributes, + Collections.emptyList()) + .getDecision() + == SamplingDecision.RECORD_AND_SAMPLE; + } + + static void clearForTest() { + System.clearProperty(ROOT_SAMPLED_PROPERTY); + } +} diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/XRayContextExtractor.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/XRayContextExtractor.java index 109f633de..1df80ba31 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/XRayContextExtractor.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/XRayContextExtractor.java @@ -42,6 +42,7 @@ public ExtractedContext extract() { String root = null; String parent = null; + String sampled = null; for (var part : traceHeader.split(";")) { var eqIdx = part.indexOf('='); @@ -52,6 +53,7 @@ public ExtractedContext extract() { switch (key) { case "Root" -> root = value; case "Parent" -> parent = value; + case "Sampled" -> sampled = value; } } @@ -68,7 +70,7 @@ public ExtractedContext extract() { return null; } - // Parent is a 16-char hex span ID + // Parent is a 16-char hex span ID; may be absent (Root-only header is still usable). String parentSpanId = null; if (parent != null) { var normalized = parent.toLowerCase(); @@ -77,7 +79,15 @@ public ExtractedContext extract() { } } - return new ExtractedContext(traceId, parentSpanId); + // Only Sampled=1 and Sampled=0 are authoritative; anything else (missing or unusable) is undecided. + var sampling = + switch (sampled == null ? "" : sampled) { + case "1" -> ExtractedContext.Sampling.SAMPLED; + case "0" -> ExtractedContext.Sampling.NOT_SAMPLED; + default -> ExtractedContext.Sampling.UNDECIDED; + }; + + return new ExtractedContext(traceId, parentSpanId, sampling); } /** diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginIntegrationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginIntegrationTest.java new file mode 100644 index 000000000..068e0f7a1 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginIntegrationTest.java @@ -0,0 +1,214 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.OpenTelemetry; +import io.opentelemetry.context.propagation.ContextPropagators; +import io.opentelemetry.javaagent.testing.FakeJavaAgentTracerProvider; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import java.time.Duration; +import java.util.List; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +/** + * Integration tests for the workflow-rooted {@link ExecutionOtelPlugin} running through the real SDK execution engine + * (LocalDurableTestRunner). Complements {@link InvocationOtelPluginIntegrationTest}, which covers the invocation-rooted + * plugin. + * + *

Execution-view topology: operations are children of the Workflow span and link to the current Invocation span; the + * whole execution shares one trace. + */ +class ExecutionOtelPluginIntegrationTest { + + private InMemorySpanExporter spanExporter; + private DurableConfig otelConfig; + + @BeforeEach + void setUp() { + DeterministicIdGenerator.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + WorkflowRootSamplingBridge.clearForTest(); + spanExporter = InMemorySpanExporter.create(); + + var plugin = new ExecutionOtelPlugin( + SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(spanExporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build()); + + otelConfig = DurableConfig.builder().withPlugins(plugin).build(); + } + + @AfterEach + void tearDown() { + GlobalOpenTelemetry.resetForTest(); + DeterministicIdGenerator.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + WorkflowRootSamplingBridge.clearForTest(); + } + + @Test + void simpleStep_operationsAreChildrenOfWorkflow_andLinkToInvocation() { + var runner = LocalDurableTestRunner.create( + String.class, (input, ctx) -> ctx.step("greet", String.class, stepCtx -> "Hello " + input), otelConfig); + + var result = runner.runUntilComplete("World"); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var spans = spanExporter.getFinishedSpanItems(); + assertTrue(spans.size() >= 4, "Expected Workflow + Invocation + operation + attempt, got " + spans.size()); + + var workflow = spanByName(spans, "Workflow"); + var invocation = spanByName(spans, "Invocation"); + var operation = spanByName(spans, "greet"); + var attempt = spanByName(spans, "greet attempt 1"); + + // The whole execution shares one trace. + assertEquals(workflow.getTraceId(), invocation.getTraceId()); + assertTrue( + spans.stream().allMatch(s -> s.getTraceId().equals(workflow.getTraceId())), + "Every span shares the execution trace"); + + // Execution-view parenting: the operation is a child of the Workflow span, and its attempt a child of it. + assertEquals(workflow.getSpanId(), operation.getParentSpanId(), "Operation is a child of the Workflow span"); + assertEquals(operation.getSpanId(), attempt.getParentSpanId(), "Attempt is a child of its operation span"); + + // Execution-view correlation: the operation links to the current Invocation span. + assertTrue( + operation.getLinks().stream() + .anyMatch(l -> l.getSpanContext().getSpanId().equals(invocation.getSpanId())), + "Operation span links to the current Invocation span"); + } + + @Test + void waitAcrossInvocations_sharesOneTrace_andExportsWorkflowOnceOnTerminal() { + // No propagated context, so the execution trace is anchored on a synthetic root derived from the ARN. The + // runner keeps the execution ARN and start time fixed across reinvocations, matching the backend, so the + // derived trace ID is the same for every invocation. + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> { + ctx.step("before-wait", String.class, stepCtx -> "pre"); + ctx.wait("pause", Duration.ofMinutes(1)); + ctx.step("after-wait", String.class, stepCtx -> "post"); + return "done"; + }, + otelConfig); + + // First invocation: step + wait, then suspend (PENDING). The Workflow span must not be exported yet. + var first = runner.run("input"); + assertEquals(ExecutionStatus.PENDING, first.getStatus()); + assertTrue( + spanExporter.getFinishedSpanItems().stream() + .noneMatch(s -> s.getName().equals("Workflow")), + "Workflow span must not be exported on a non-terminal invocation"); + var firstInvocationTraceId = + spanByName(spanExporter.getFinishedSpanItems(), "Invocation").getTraceId(); + + // Resume and complete. + runner.advanceTime(); + var second = runner.runUntilComplete("input"); + assertEquals(ExecutionStatus.SUCCEEDED, second.getStatus()); + + var spans = spanExporter.getFinishedSpanItems(); + + // Exactly one Workflow span across the whole execution, exported on the terminal invocation. + var workflowSpans = + spans.stream().filter(s -> s.getName().equals("Workflow")).toList(); + assertEquals(1, workflowSpans.size(), "Workflow span is exported exactly once, on the terminal invocation"); + + // Single trace per execution: every span from both invocations shares one trace ID, stable across + // reinvocations. + var executionTraceId = workflowSpans.get(0).getTraceId(); + assertEquals(firstInvocationTraceId, executionTraceId, "The first invocation already used the execution trace"); + assertTrue( + spans.stream().allMatch(s -> s.getTraceId().equals(executionTraceId)), + "Both invocations and the Workflow span share one execution trace"); + + // Two Invocation spans (one per run), both on the execution trace. + assertEquals( + 2, spans.stream().filter(s -> s.getName().equals("Invocation")).count(), "One Invocation span per run"); + } + + @Test + void wrappedAgentProvider_withRootDroppingSampler_keepsExecutionTreeConsistentlySampled() { + // The agent provider is hidden behind a classloader wrapper, so the plugin cannot reach the sampler and the + // bridge published nothing. With no reachable decision, the synthetic execution root is treated as sampled so + // the execution root and everything parented onto it are sampled together — no orphan operation spans exported + // under a dropped root. + OtelPluginAutoConfigurationState.markInstalled(); + GlobalOpenTelemetry.resetForTest(); + WorkflowRootSamplingBridge.clearForTest(); + var globalExporter = InMemorySpanExporter.create(); + var sdkTracerProvider = SdkTracerProvider.builder() + .setIdGenerator(new DeterministicIdGenerator()) + .setSampler(Sampler.parentBased(Sampler.alwaysOff())) + .addSpanProcessor(SimpleSpanProcessor.create(globalExporter)) + .build(); + var javaAgentTracerProvider = new FakeJavaAgentTracerProvider(sdkTracerProvider); + GlobalOpenTelemetry.set(new OpenTelemetry() { + @Override + public io.opentelemetry.api.trace.TracerProvider getTracerProvider() { + return javaAgentTracerProvider; + } + + @Override + public ContextPropagators getPropagators() { + return ContextPropagators.noop(); + } + }); + + var defaultConfig = + DurableConfig.builder().withPlugins(new ExecutionOtelPlugin()).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, ctx) -> ctx.step("wrapped-step", String.class, stepCtx -> "Hello " + input), + defaultConfig); + + var result = runner.runUntilComplete("World"); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + + var spans = globalExporter.getFinishedSpanItems(); + assertSpanExists(spans, "Workflow"); + assertSpanExists(spans, "wrapped-step"); + var workflowTraceId = spans.stream() + .filter(s -> s.getName().equals("Workflow")) + .findFirst() + .orElseThrow() + .getTraceId(); + // Every exported span is on the one execution trace: the root and its children are sampled consistently. + assertTrue( + spans.stream().allMatch(s -> s.getTraceId().equals(workflowTraceId)), + "All spans share the sampled execution trace"); + } + + // Helpers + + private static SpanData spanByName(List spans, String name) { + return spans.stream() + .filter(s -> s.getName().equals(name)) + .findFirst() + .orElseThrow(() -> new AssertionError("No span named '" + name + "' in " + + spans.stream().map(SpanData::getName).toList())); + } + + private static void assertSpanExists(List spans, String expectedName) { + assertTrue( + spans.stream().anyMatch(s -> s.getName().equals(expectedName)), + "Expected span '" + expectedName + "' not found. Got: " + + spans.stream().map(SpanData::getName).toList()); + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java index 887e86666..223ec6d72 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginTest.java @@ -17,7 +17,9 @@ import io.opentelemetry.sdk.resources.Resource; import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.SpanData; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; import java.time.Instant; import java.util.ServiceLoader; import org.junit.jupiter.api.AfterEach; @@ -247,7 +249,7 @@ void workflowSpan_hasInternalKind() { } @Test - void workflowAndInvocationSpans_areIndependentRoots_withoutAmbientContext() { + void workflowAndInvocationSpans_shareExecutionTrace_withoutAmbientContext() { plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); @@ -255,28 +257,65 @@ void workflowAndInvocationSpans_areIndependentRoots_withoutAmbientContext() { var workflowSpan = spanByName(spans, "Workflow"); var invocationSpan = spanByName(spans, "Invocation"); - assertFalse(workflowSpan.getParentSpanContext().isValid(), "Workflow span must be a root"); - assertFalse(invocationSpan.getParentSpanContext().isValid(), "Invocation span must be a root"); - assertNotEquals( - workflowSpan.getTraceId(), invocationSpan.getTraceId(), "Independent roots must not share a trace ID"); + // With no propagated context, a synthetic execution root anchors the trace and both spans parent onto it. + assertEquals( + workflowSpan.getTraceId(), + invocationSpan.getTraceId(), + "Workflow and Invocation spans share the execution trace"); + assertTrue(workflowSpan.getParentSpanContext().isValid(), "Workflow span parents onto the execution ancestor"); + assertTrue( + invocationSpan.getParentSpanContext().isValid(), "Invocation span parents onto the execution ancestor"); + assertEquals( + workflowSpan.getParentSpanId(), + invocationSpan.getParentSpanId(), + "Both spans share the same synthetic execution root as parent"); assertEquals(SpanKind.INTERNAL, invocationSpan.getKind()); } @Test - void invocationStart_usesCurrentSpanContext_whenExtractorReturnsNull() { - var traceId = "5759e988bd862e3fe1be46a994272793"; - var parentSpanId = "53995c3f42cd8ad8"; - var parentSpanContext = - SpanContext.create(traceId, parentSpanId, TraceFlags.getSampled(), TraceState.getDefault()); + void invocationStart_joinsAmbientTrace_whenAmbientIsOnExecutionTrace() { + // Drive an invocation to learn the canonical execution trace ID, then start a fresh invocation with an ambient + // span on that same trace: the Invocation span joins the ambient span directly. + plugin.onInvocationStart(new InvocationInfo("req-0", ARN, true, Instant.now())); + plugin.onInvocationEnd(new InvocationEndInfo("req-0", ARN, true, InvocationStatus.SUCCEEDED, null)); + var canonicalTraceId = + spanByName(spanExporter.getFinishedSpanItems(), "Workflow").getTraceId(); + spanExporter.reset(); + + var ambientSpanId = "53995c3f42cd8ad8"; + var ambient = + SpanContext.create(canonicalTraceId, ambientSpanId, TraceFlags.getSampled(), TraceState.getDefault()); + try (var ignored = Span.wrap(ambient).makeCurrent()) { + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, false, Instant.now())); + } + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, false, InvocationStatus.SUCCEEDED, null)); - try (var ignored = Span.wrap(parentSpanContext).makeCurrent()) { + var invocationSpan = spanByName(spanExporter.getFinishedSpanItems(), "Invocation"); + assertEquals(canonicalTraceId, invocationSpan.getTraceId()); + assertEquals(ambientSpanId, invocationSpan.getParentSpanId(), "Invocation joins the ambient span on its trace"); + } + + @Test + void invocationStart_ignoresUnrelatedAmbientTrace() { + // An ambient span on a different trace must not pull the Invocation span off the execution trace. + var unrelated = SpanContext.create( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "1111111111111111", + TraceFlags.getSampled(), + TraceState.getDefault()); + try (var ignored = Span.wrap(unrelated).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); } plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); - var invocationSpan = spanByName(spanExporter.getFinishedSpanItems(), "Invocation"); - assertEquals(traceId, invocationSpan.getTraceId()); - assertEquals(parentSpanId, invocationSpan.getParentSpanId()); + var spans = spanExporter.getFinishedSpanItems(); + var workflowSpan = spanByName(spans, "Workflow"); + var invocationSpan = spanByName(spans, "Invocation"); + assertEquals( + workflowSpan.getTraceId(), + invocationSpan.getTraceId(), + "Invocation stays on the execution trace, not the unrelated ambient trace"); + assertNotEquals("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", invocationSpan.getTraceId()); } @Test @@ -766,8 +805,9 @@ void operationNotCompleted_notEndedAtInvocationEnd() { plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); var spans = spanExporter.getFinishedSpanItems(); - // Only the invocation span is exported. The still-open operation span is NOT force-ended (no PENDING - // span here), and the Workflow span is not exported on a non-terminal invocation. + // Only the invocation span is exported. A still-open operation has no recording span (creation is deferred to + // onOperationEnd), so there is nothing to abandon, and the Workflow span is not exported on a non-terminal + // invocation. assertEquals(1, spans.size()); assertEquals("Invocation", spans.get(0).getName()); assertTrue( @@ -775,6 +815,43 @@ void operationNotCompleted_notEndedAtInvocationEnd() { "An operation still open at invocation end must not be ended/exported in onInvocationEnd"); } + @Test + void openAttemptSpan_isEndedAtInvocationEnd_notAbandoned() { + // A user function that starts but never ends (e.g. the execution suspends mid-attempt) must not leave a + // recording span abandoned: onInvocationEnd force-ends it so it is exported. + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + plugin.onUserFunctionStart( + new UserFunctionStartInfo("op-1", "stuck", "STEP", "Step", null, Instant.now(), false, 1)); + // No onUserFunctionEnd — the invocation suspends. + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); + + var attemptSpan = spanExporter.getFinishedSpanItems().stream() + .filter(s -> s.getName().contains("stuck")) + .findFirst() + .orElseThrow(() -> new AssertionError("Open attempt span must be ended and exported, not abandoned")); + assertTrue(attemptSpan.hasEnded(), "Attempt span must be ended"); + } + + @Test + void everyRecordingSpanIsEnded_onNonTerminalInvocation() { + // No recording span may be left un-ended when the execution returns a non-terminal status. + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + plugin.onOperationStart( + new OperationInfo("op-1", "step-a", "STEP", "Step", null, Instant.now(), null, null, false)); + plugin.onUserFunctionStart( + new UserFunctionStartInfo("op-1", "step-a", "STEP", "Step", null, Instant.now(), false, 1)); + // Suspend mid-attempt: neither onUserFunctionEnd nor onOperationEnd fires. + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); + + // Whatever spans were exported this invocation must all be ended (InMemorySpanExporter only records ended + // spans, so an abandoned recording span would simply be missing — assert the attempt span made it out ended). + var spans = spanExporter.getFinishedSpanItems(); + assertTrue(spans.stream().allMatch(SpanData::hasEnded), "All exported spans are ended"); + assertTrue( + spans.stream().anyMatch(s -> s.getName().contains("step-a")), + "The open attempt span is force-ended and exported rather than abandoned"); + } + @Test void operationOpenedThenCompletedNextInvocation_exportedOnceOnOperationEnd() { // Invocation 1: operation opens but does not complete. @@ -819,7 +896,7 @@ void operationOpenedThenCompletedNextInvocation_exportedOnceOnOperationEnd() { // ─── Cross-invocation stitching ────────────────────────────────────── @Test - void workflowTraceIsStableAndInvocationRootsAreFresh_acrossInvocations() { + void executionTraceIsStableAcrossInvocations_andSharedByInvocationSpans() { var executionStartTime = Instant.parse("2026-08-15T00:00:00Z"); plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, executionStartTime)); plugin.onOperationStart( @@ -839,7 +916,7 @@ void workflowTraceIsStableAndInvocationRootsAreFresh_acrossInvocations() { null)); plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); var firstSpans = spanExporter.getFinishedSpanItems(); - var workflowTraceId = spanByName(firstSpans, "step-1").getTraceId(); + var executionTraceId = spanByName(firstSpans, "step-1").getTraceId(); var firstInvocationTraceId = spanByName(firstSpans, "Invocation").getTraceId(); spanExporter.reset(); @@ -849,10 +926,10 @@ void workflowTraceIsStableAndInvocationRootsAreFresh_acrossInvocations() { var workflowSpan = spanByName(secondSpans, "Workflow"); var secondInvocationSpan = spanByName(secondSpans, "Invocation"); - assertEquals(workflowTraceId, workflowSpan.getTraceId()); - assertNotEquals(workflowTraceId, firstInvocationTraceId); - assertNotEquals(workflowTraceId, secondInvocationSpan.getTraceId()); - assertNotEquals(firstInvocationTraceId, secondInvocationSpan.getTraceId()); + // The whole execution shares one trace ID, stable across invocations. + assertEquals(executionTraceId, workflowSpan.getTraceId()); + assertEquals(executionTraceId, firstInvocationTraceId); + assertEquals(executionTraceId, secondInvocationSpan.getTraceId()); } @Test @@ -934,12 +1011,17 @@ void sampling_disabled_producesNoSpans() { // ─── X-Ray trace ID ────────────────────────────────────────────────── @Test - void xrayExtraction_keepsWorkflowTraceIndependent() { + void xrayExtraction_undecidedSampling_remoteParentIsAncestor_flagUnset() { var xrayTraceId = "aabbccddee112233445566778899aabb"; var parentSpanId = "53995c3f42cd8ad8"; var exporter = InMemorySpanExporter.create(); + // Two-arg context → UNDECIDED sampling: the valid remote parent is still the authoritative ancestor. A + // non-parent-based alwaysOn sampler exports the spans so the topology is observable (a plain parent-based + // sampler would drop them, since the remote parent's sampled flag is left unset). var xrayPlugin = new ExecutionOtelPlugin( - SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)), + SdkTracerProvider.builder() + .setSampler(Sampler.alwaysOn()) + .addSpanProcessor(SimpleSpanProcessor.create(exporter)), OtelPluginConfig.builder() .contextExtractor(() -> new ExtractedContext(xrayTraceId, parentSpanId)) .enableMdc(false) @@ -968,21 +1050,50 @@ void xrayExtraction_keepsWorkflowTraceIndependent() { var workflowSpan = spanByName(spans, "Workflow"); var invocationSpan = spanByName(spans, "Invocation"); var operationSpan = spanByName(spans, "step-a"); + // The remote parent is the ancestor: Workflow and Invocation both parent onto it, on the remote trace. + assertEquals(xrayTraceId, workflowSpan.getTraceId()); assertEquals(xrayTraceId, invocationSpan.getTraceId()); - assertEquals(parentSpanId, invocationSpan.getParentSpanId()); - assertNotEquals(xrayTraceId, workflowSpan.getTraceId()); - assertEquals(workflowSpan.getTraceId(), operationSpan.getTraceId()); + assertEquals(xrayTraceId, operationSpan.getTraceId()); + assertEquals(parentSpanId, workflowSpan.getParentSpanId(), "Workflow parents onto the remote span"); + assertEquals(parentSpanId, invocationSpan.getParentSpanId(), "Invocation parents onto the remote span"); + assertTrue(workflowSpan.getLinks().isEmpty(), "No remote-parent link when the remote context is the ancestor"); } @Test - void xrayExtraction_withParentSpanId_invocationSpanHasCorrectParent() { + void xrayExtraction_undecidedSampling_parentBasedSampler_dropsExecutionTrace() { + // With a plain ParentBased sampler and no explicit upstream Sampled, the remote parent's unset flag makes the + // whole execution trace not-sampled, so nothing is exported. + var xrayTraceId = "aabbccddee112233445566778899aabb"; + var parentSpanId = "53995c3f42cd8ad8"; + var exporter = InMemorySpanExporter.create(); + var xrayPlugin = new ExecutionOtelPlugin( + SdkTracerProvider.builder() + .setSampler(Sampler.parentBased(Sampler.alwaysOn())) + .addSpanProcessor(SimpleSpanProcessor.create(exporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> new ExtractedContext(xrayTraceId, parentSpanId)) + .enableMdc(false) + .workflowSpanName("Workflow") + .build()); + xrayPlugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + xrayPlugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); + + assertTrue( + exporter.getFinishedSpanItems().isEmpty(), + "Unset remote-parent flag drops the execution trace under a parent-based sampler"); + } + + @Test + void xrayExtraction_explicitSampled_remoteParentIsExecutionAncestor() { var xrayTraceId = "5759e988bd862e3fe1be46a994272793"; var parentSpanId = "53995c3f42cd8ad8"; var exporter = InMemorySpanExporter.create(); + // Explicit Sampled=1 with a complete parent → the remote context is the execution ancestor directly. var xrayPlugin = new ExecutionOtelPlugin( SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)), OtelPluginConfig.builder() - .contextExtractor(() -> new ExtractedContext(xrayTraceId, parentSpanId)) + .contextExtractor(() -> + new ExtractedContext(xrayTraceId, parentSpanId, ExtractedContext.Sampling.SAMPLED)) .enableMdc(false) .workflowSpanName("Workflow") .build()); @@ -993,9 +1104,12 @@ void xrayExtraction_withParentSpanId_invocationSpanHasCorrectParent() { var spans = exporter.getFinishedSpanItems(); var workflowSpan = spanByName(spans, "Workflow"); var invocationSpan = spanByName(spans, "Invocation"); - assertFalse(workflowSpan.getParentSpanContext().isValid(), "Workflow span must remain an independent root"); + assertEquals(xrayTraceId, workflowSpan.getTraceId()); assertEquals(xrayTraceId, invocationSpan.getTraceId()); - assertEquals(parentSpanId, invocationSpan.getParentSpanId()); + assertEquals(parentSpanId, workflowSpan.getParentSpanId(), "Workflow parents onto the remote span directly"); + assertEquals( + parentSpanId, invocationSpan.getParentSpanId(), "Invocation parents onto the remote span directly"); + assertTrue(workflowSpan.getLinks().isEmpty(), "No remote-parent link when the remote context is the ancestor"); } // ─── Helpers ───────────────────────────────────────────────────────── diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionTraceContextTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionTraceContextTest.java new file mode 100644 index 000000000..174897d61 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionTraceContextTest.java @@ -0,0 +1,114 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import java.time.Instant; +import java.util.function.BooleanSupplier; +import org.junit.jupiter.api.Test; + +/** Unit tests for the execution-trace resolution table shared by both plugins. */ +class ExecutionTraceContextTest { + + private static final String ARN = "arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1"; + private static final String REMOTE_TRACE_ID = "aabbccddee112233445566778899aabb"; + private static final String REMOTE_PARENT_ID = "53995c3f42cd8ad8"; + private static final Instant START = Instant.parse("2026-08-15T00:00:00Z"); + + private final DeterministicIdGenerator idGenerator = new DeterministicIdGenerator(); + + @Test + void canonicalTraceId_reusesRemoteTraceId_whenPresent() { + var extracted = new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID); + assertEquals(REMOTE_TRACE_ID, ExecutionTraceContext.canonicalTraceId(extracted, ARN, START, idGenerator)); + } + + @Test + void canonicalTraceId_derivesFromArn_whenNoRemoteTraceId() { + var canonical = ExecutionTraceContext.canonicalTraceId(null, ARN, START, idGenerator); + assertEquals(idGenerator.generateTraceIdForExecution(ARN, START), canonical); + assertTrue(canonical.matches("[0-9a-f]{32}")); + } + + @Test + void completeRemoteParentWithSampled_becomesAncestor_preservesSampled() { + // Row: valid Root, Parent, Sampled=1 -> reuse Root, remote parent, preserve sampled. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID, ExtractedContext.Sampling.SAMPLED); + var execCtx = resolve(extracted, REMOTE_TRACE_ID, () -> false); + + assertEquals(REMOTE_TRACE_ID, execCtx.executionAncestor().getTraceId()); + assertEquals(REMOTE_PARENT_ID, execCtx.executionAncestor().getSpanId()); + assertTrue(execCtx.executionAncestor().isRemote(), "The remote server span is the ancestor"); + assertTrue(execCtx.traceFlags().isSampled(), "Explicit upstream Sampled=1 wins over the supplier"); + } + + @Test + void completeRemoteParentNotSampled_becomesAncestor_preservesNotSampled() { + // Row: valid Root, Parent, Sampled=0 -> reuse Root, remote parent, preserve not-sampled. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID, ExtractedContext.Sampling.NOT_SAMPLED); + var execCtx = resolve(extracted, REMOTE_TRACE_ID, () -> true); + + assertEquals(REMOTE_PARENT_ID, execCtx.executionAncestor().getSpanId()); + assertFalse(execCtx.traceFlags().isSampled(), "Explicit upstream Sampled=0 wins over the supplier"); + } + + @Test + void completeRemoteParentUndecided_becomesAncestor_flagLeftUnset() { + // Row: valid Root, Parent, no valid Sampled -> reuse Root, remote parent, sampled flag left unset (the + // configured sampler decides at span start). The synthetic root is NOT used and the supplier is NOT consulted. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID); + var execCtx = resolve(extracted, REMOTE_TRACE_ID, () -> { + throw new AssertionError("supplier must not be consulted when a valid remote parent exists"); + }); + + assertEquals(REMOTE_TRACE_ID, execCtx.executionAncestor().getTraceId()); + assertEquals(REMOTE_PARENT_ID, execCtx.executionAncestor().getSpanId()); + assertFalse(execCtx.traceFlags().isSampled(), "Absent upstream decision leaves the flag unset"); + } + + @Test + void remoteTraceWithoutParent_synthesizesRootOnRemoteTrace() { + // Row: valid Root, missing Parent, no valid Sampled -> reuse Root, synthetic root, supplier decides. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, null); + var execCtx = resolve(extracted, REMOTE_TRACE_ID, () -> true); + + assertEquals(REMOTE_TRACE_ID, execCtx.executionAncestor().getTraceId()); + assertEquals( + idGenerator.generateExecutionRootSpanId(ARN), + execCtx.executionAncestor().getSpanId()); + assertTrue(execCtx.traceFlags().isSampled(), "The supplier decides for a synthetic root"); + } + + @Test + void remoteTraceWithoutParent_explicitSampledPreserved_overSupplier() { + // Row: valid Root, missing Parent, Sampled=1 -> reuse Root, synthetic root, preserve explicit decision. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, null, ExtractedContext.Sampling.SAMPLED); + var execCtx = resolve(extracted, REMOTE_TRACE_ID, () -> false); + + assertEquals( + idGenerator.generateExecutionRootSpanId(ARN), + execCtx.executionAncestor().getSpanId()); + assertTrue(execCtx.traceFlags().isSampled(), "Explicit decision is preserved even on a synthetic root"); + } + + @Test + void noContext_synthesizesRootOnCanonicalTrace_supplierDecidesSampling() { + // Row: missing Root -> derive trace ID from ARN and start time, synthetic root, supplier decides. + var canonical = ExecutionTraceContext.canonicalTraceId(null, ARN, START, idGenerator); + var sampledCtx = resolve(null, canonical, () -> true); + var droppedCtx = resolve(null, canonical, () -> false); + + assertEquals(canonical, sampledCtx.executionAncestor().getTraceId()); + assertEquals( + idGenerator.generateExecutionRootSpanId(ARN), + sampledCtx.executionAncestor().getSpanId()); + assertTrue(sampledCtx.traceFlags().isSampled()); + assertFalse(droppedCtx.traceFlags().isSampled()); + } + + private ExecutionTraceContext resolve( + ExtractedContext extracted, String canonicalTraceId, BooleanSupplier rootSampled) { + return ExecutionTraceContext.resolve(extracted, canonicalTraceId, ARN, idGenerator, rootSampled); + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginIntegrationTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginIntegrationTest.java index 274944c36..cb700124f 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginIntegrationTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/InvocationOtelPluginIntegrationTest.java @@ -90,10 +90,9 @@ void simpleStep_producesInvocationAndOperationAndAttemptSpans() { .findFirst() .orElseThrow() .getTraceId(); - assertNotEquals(workflowTraceId, invocationTraceId); - assertTrue(spans.stream() - .filter(span -> !span.getName().equals("Workflow")) - .allMatch(span -> span.getTraceId().equals(invocationTraceId))); + // Workflow and Invocation now share the single execution trace. + assertEquals(workflowTraceId, invocationTraceId); + assertTrue(spans.stream().allMatch(span -> span.getTraceId().equals(invocationTraceId))); } @Test 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 0d14daa15..0663ccfbd 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 @@ -293,23 +293,27 @@ void invocationOtelPluginProvider_isRegisteredAsServiceProvider() { } @Test - void invocationStart_usesCurrentSpanContext_whenExtractorReturnsNull() { - var traceId = "5759e988bd862e3fe1be46a994272793"; - var parentSpanId = "53995c3f42cd8ad8"; - var parentSpanContext = - SpanContext.create(traceId, parentSpanId, TraceFlags.getSampled(), TraceState.getDefault()); - - try (var ignored = Span.wrap(parentSpanContext).makeCurrent()) { + void invocationStart_ignoresUnrelatedAmbientTrace() { + // The extractor returns null here, so the canonical execution trace is ARN-derived. The ambient span is on an + // unrelated trace, so it is NOT used as the Invocation parent — the Invocation span parents onto the execution + // ancestor and shares the execution/Workflow trace instead. + var ambientTraceId = "5759e988bd862e3fe1be46a994272793"; + var ambientSpanId = "53995c3f42cd8ad8"; + var ambientSpanContext = + SpanContext.create(ambientTraceId, ambientSpanId, TraceFlags.getSampled(), TraceState.getDefault()); + + try (var ignored = Span.wrap(ambientSpanContext).makeCurrent()) { plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec1", true, Instant.now())); } plugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec1", true, InvocationStatus.SUCCEEDED, null)); - var invocationSpan = spanExporter.getFinishedSpanItems().stream() - .filter(span -> span.getName().equals("Invocation")) - .findFirst() - .orElseThrow(); - assertEquals(traceId, invocationSpan.getTraceId()); - assertEquals(parentSpanId, invocationSpan.getParentSpanId()); + var invocationSpan = spanByName("Invocation"); + var workflowSpan = spanByName("Workflow"); + // The unrelated ambient trace is ignored — the Invocation span is not placed on it. + assertNotEquals(ambientTraceId, invocationSpan.getTraceId()); + assertNotEquals(ambientSpanId, invocationSpan.getParentSpanId()); + // The Invocation span joins the same execution trace as the Workflow span. + assertEquals(workflowSpan.getTraceId(), invocationSpan.getTraceId()); } @Test @@ -939,32 +943,34 @@ void fullLifecycle_producesCorrectSpanHierarchy() { var workflowTraceId = spanByName("Workflow").getTraceId(); var invocationTraceId = spanByName("Invocation").getTraceId(); - assertNotEquals(workflowTraceId, invocationTraceId); - assertTrue(spans.stream() - .filter(span -> !span.getName().equals("Workflow")) - .allMatch(span -> span.getTraceId().equals(invocationTraceId))); + // Workflow and Invocation share the single execution trace. + assertEquals(workflowTraceId, invocationTraceId); + assertTrue(spans.stream().allMatch(span -> span.getTraceId().equals(invocationTraceId))); } @Test - void invocationRoots_sameExecutionReceiveFreshTraceIds() { + void invocationRoots_sameExecutionShareExecutionTrace() { var arn = "arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1"; + // Same execution start time across invocations so the ARN-derived canonical trace ID is reproducible. + var startTime = Instant.now(); - plugin.onInvocationStart(new InvocationInfo("req-1", arn, true, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-1", arn, true, startTime)); plugin.onInvocationEnd(new InvocationEndInfo("req-1", arn, true, InvocationStatus.PENDING, null)); var firstTraceId = spanByName("Invocation").getTraceId(); spanExporter.reset(); // Second invocation of same execution - plugin.onInvocationStart(new InvocationInfo("req-2", arn, false, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-2", arn, false, startTime)); plugin.onInvocationEnd(new InvocationEndInfo("req-2", arn, false, InvocationStatus.SUCCEEDED, null)); var secondTraceId = spanByName("Invocation").getTraceId(); var workflowTraceId = spanByName("Workflow").getTraceId(); - assertNotEquals(firstTraceId, secondTraceId); - assertNotEquals(firstTraceId, workflowTraceId); - assertNotEquals(secondTraceId, workflowTraceId); + // Every durable execution shares ONE trace: both invocations and the Workflow span are on it. + assertEquals(firstTraceId, secondTraceId); + assertEquals(firstTraceId, workflowTraceId); + assertEquals(secondTraceId, workflowTraceId); } @Test @@ -1105,14 +1111,19 @@ void xrayExtraction_withoutParentDoesNotForceTraceId() { assertEquals(2, spans.size()); // invocation + Workflow var invocationSpan = spanByName("Invocation"); var workflowSpan = spanByName("Workflow"); - assertFalse(invocationSpan.getParentSpanContext().isValid()); - assertNotEquals(xrayTraceId, invocationSpan.getTraceId()); - assertNotEquals(xrayTraceId, workflowSpan.getTraceId()); - assertNotEquals(invocationSpan.getTraceId(), workflowSpan.getTraceId()); + // Remote trace but no parent → a synthetic execution root on the remote trace ID anchors the execution. Both + // the Invocation and Workflow spans parent onto that synthetic root and share the extracted trace ID. + var executionRootSpanId = new DeterministicIdGenerator().generateExecutionRootSpanId("arn:exec1"); + assertTrue(invocationSpan.getParentSpanContext().isValid()); + assertEquals(xrayTraceId, invocationSpan.getTraceId()); + assertEquals(executionRootSpanId, invocationSpan.getParentSpanId()); + assertEquals(xrayTraceId, workflowSpan.getTraceId()); + assertEquals(executionRootSpanId, workflowSpan.getParentSpanId()); + assertEquals(invocationSpan.getTraceId(), workflowSpan.getTraceId()); } @Test - void xrayExtraction_invocationTreeUsesExtractedTraceId_workflowRemainsIndependent() { + void xrayExtraction_invocationTreeUsesExtractedTraceId_workflowJoinsExecutionTrace() { var xrayTraceId = "aabbccddee112233445566778899aabb"; var parentSpanId = "53995c3f42cd8ad8"; var extractedContext = new ExtractedContext(xrayTraceId, parentSpanId); @@ -1158,20 +1169,18 @@ void xrayExtraction_invocationTreeUsesExtractedTraceId_workflowRemainsIndependen xrayPlugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec1", true, InvocationStatus.SUCCEEDED, null)); var spans = spanExporter.getFinishedSpanItems(); - var workflowTraceId = spanByName("Workflow").getTraceId(); - assertNotEquals(xrayTraceId, workflowTraceId); + // The remote trace ID is the canonical execution trace, so the Workflow span joins it too. assertTrue( - spans.stream() - .filter(span -> !span.getName().equals("Workflow")) - .allMatch(span -> span.getTraceId().equals(xrayTraceId)), - "The invocation tree should inherit the extracted X-Ray trace ID"); + spans.stream().allMatch(span -> span.getTraceId().equals(xrayTraceId)), + "The whole execution shares the extracted X-Ray trace ID"); } @Test void xrayExtraction_withParentSpanId_invocationSpanHasCorrectParent() { var xrayTraceId = "5759e988bd862e3fe1be46a994272793"; var parentSpanId = "53995c3f42cd8ad8"; - var extractedContext = new ExtractedContext(xrayTraceId, parentSpanId); + // Explicit SAMPLED so the complete remote parent is the sampled ancestor and the spans export. + var extractedContext = new ExtractedContext(xrayTraceId, parentSpanId, ExtractedContext.Sampling.SAMPLED); spanExporter = InMemorySpanExporter.create(); var xrayPlugin = new InvocationOtelPlugin( @@ -1194,11 +1203,16 @@ void xrayExtraction_withParentSpanId_invocationSpanHasCorrectParent() { parentSpanId, invocationSpan.getParentSpanId(), "Invocation span should be parented to X-Ray Parent span"); - assertNotEquals(xrayTraceId, workflowSpan.getTraceId()); + // The Workflow span joins the same execution trace and parents onto the same remote ancestor. + assertEquals(xrayTraceId, workflowSpan.getTraceId()); + assertEquals( + parentSpanId, + workflowSpan.getParentSpanId(), + "Workflow span should be parented to the same X-Ray Parent span"); } @Test - void xrayExtraction_withoutParentSpanId_invocationSpanIsRoot() { + void xrayExtraction_withoutParentSpanId_invocationSpanParentsOntoSyntheticRoot() { var xrayTraceId = "5759e988bd862e3fe1be46a994272793"; var extractedContext = new ExtractedContext(xrayTraceId, null); @@ -1216,15 +1230,20 @@ void xrayExtraction_withoutParentSpanId_invocationSpanIsRoot() { var spans = spanExporter.getFinishedSpanItems(); assertEquals(2, spans.size()); // invocation + Workflow + // Remote trace, no parent → a synthetic execution root on the remote trace ID anchors the execution. The + // Invocation span has a valid parent (that synthetic root) and joins the remote trace. + var executionRootSpanId = new DeterministicIdGenerator().generateExecutionRootSpanId("arn:exec1"); var invocationSpan = spanByName("Invocation"); - assertFalse(invocationSpan.getParentSpanContext().isValid()); - assertNotEquals(xrayTraceId, invocationSpan.getTraceId()); + assertTrue(invocationSpan.getParentSpanContext().isValid()); + assertEquals(xrayTraceId, invocationSpan.getTraceId()); + assertEquals(executionRootSpanId, invocationSpan.getParentSpanId()); } @Test void xrayExtraction_multipleInvocations_sameTraceId_unifiedTrace() { var xrayTraceId = "5759e988bd862e3fe1be46a994272793"; - var extractedContext = new ExtractedContext(xrayTraceId, "53995c3f42cd8ad8"); + // Explicit SAMPLED so the complete remote parent is the sampled ancestor and spans export across invocations. + var extractedContext = new ExtractedContext(xrayTraceId, "53995c3f42cd8ad8", ExtractedContext.Sampling.SAMPLED); spanExporter = InMemorySpanExporter.create(); var xrayPlugin = new InvocationOtelPlugin( @@ -1276,17 +1295,14 @@ void xrayExtraction_multipleInvocations_sameTraceId_unifiedTrace() { var spans = spanExporter.getFinishedSpanItems(); assertTrue(spans.size() >= 4, "Should have spans from both invocations"); - var workflowTraceId = spanByName("Workflow").getTraceId(); - assertNotEquals(xrayTraceId, workflowTraceId); + // Same X-Ray Root across invocations → one unified execution trace, Workflow span included. assertTrue( - spans.stream() - .filter(span -> !span.getName().equals("Workflow")) - .allMatch(span -> span.getTraceId().equals(xrayTraceId)), - "Both invocation trees should inherit the X-Ray trace ID"); + spans.stream().allMatch(span -> span.getTraceId().equals(xrayTraceId)), + "Both invocation trees and the Workflow span share the X-Ray trace ID"); } @Test - void xrayExtraction_nullExtractor_usesIndependentValidRootIds() { + void xrayExtraction_nullExtractor_sharesArnDerivedExecutionTrace() { spanExporter = InMemorySpanExporter.create(); var noXrayPlugin = new InvocationOtelPlugin( SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(spanExporter)), @@ -1306,7 +1322,12 @@ void xrayExtraction_nullExtractor_usesIndependentValidRootIds() { var workflowTraceId = spanByName("Workflow").getTraceId(); assertTrue(invocationTraceId.matches("[0-9a-f]{32}")); assertTrue(workflowTraceId.matches("[0-9a-f]{32}")); - assertNotEquals(invocationTraceId, workflowTraceId); + // With a null extractor the whole execution shares the ARN-derived synthetic execution trace. + assertEquals(invocationTraceId, workflowTraceId); + // The Workflow span keeps its deterministic ARN-derived span ID. + assertEquals( + new DeterministicIdGenerator().generateWorkflowSpanId(arn), + spanByName("Workflow").getSpanId()); } @Test @@ -1319,8 +1340,9 @@ void xrayExtraction_extractedTraceIdMatchesXrayConversion() { var convertedId = XRayContextExtractor.xrayRootToOtelTraceId(xrayRoot); assertEquals(expectedOtelTraceId, convertedId); - // Now feed it through the plugin - var extractedContext = new ExtractedContext(convertedId, "53995c3f42cd8ad8"); + // Now feed it through the plugin. Explicit SAMPLED so the complete remote parent is the sampled ancestor and + // the spans export. + var extractedContext = new ExtractedContext(convertedId, "53995c3f42cd8ad8", ExtractedContext.Sampling.SAMPLED); spanExporter = InMemorySpanExporter.create(); var xrayPlugin = new InvocationOtelPlugin( SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(spanExporter)), @@ -1333,7 +1355,8 @@ void xrayExtraction_extractedTraceIdMatchesXrayConversion() { xrayPlugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec1", true, InvocationStatus.SUCCEEDED, null)); assertEquals(expectedOtelTraceId, spanByName("Invocation").getTraceId()); - assertNotEquals(expectedOtelTraceId, spanByName("Workflow").getTraceId()); + // The Workflow span joins the same execution trace. + assertEquals(expectedOtelTraceId, spanByName("Workflow").getTraceId()); } // ─── Cross-invocation continuation span tests ──────────────────────── @@ -1559,8 +1582,11 @@ void childOperation_parentedToParentOperationSpan() { void multiInvocation_stepWaitStep_producesCorrectSpans() { var arn = "arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1"; + // Same execution start time across invocations so the ARN-derived canonical trace ID is reproducible. + var startTime = Instant.now(); + // Invocation 1: step completes, wait starts - plugin.onInvocationStart(new InvocationInfo("req-1", arn, true, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-1", arn, true, startTime)); plugin.onOperationStart( new OperationInfo("op-1", "step-A", "STEP", "Step", null, Instant.now(), null, null, false)); plugin.onUserFunctionStart( @@ -1611,7 +1637,7 @@ void multiInvocation_stepWaitStep_producesCorrectSpans() { spanExporter.reset(); // Invocation 2: wait completed between invocations, new step runs - plugin.onInvocationStart(new InvocationInfo("req-2", arn, false, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-2", arn, false, startTime)); plugin.onOperationEnd(new OperationEndInfo( "op-2", "pause", @@ -1669,9 +1695,10 @@ void multiInvocation_stepWaitStep_producesCorrectSpans() { .filter(span -> span.getName().equals("Workflow")) .findFirst() .orElseThrow(); - assertNotEquals(inv1TraceId, inv2TraceId); - assertNotEquals(inv1TraceId, workflowSpan.getTraceId()); - assertNotEquals(inv2TraceId, workflowSpan.getTraceId()); + // Every durable execution shares ONE trace: both invocations and the Workflow span are on it. + assertEquals(inv1TraceId, inv2TraceId); + assertEquals(inv2TraceId, workflowSpan.getTraceId()); + assertTrue(inv2Spans.stream().allMatch(span -> span.getTraceId().equals(inv2TraceId))); var waitContinuation = inv2Spans.stream() .filter(s -> s.getName().contains("pause")) @@ -1679,10 +1706,17 @@ void multiInvocation_stepWaitStep_producesCorrectSpans() { .orElseThrow(); assertTrue(waitContinuation.getLinks().stream() .anyMatch(link -> link.getSpanContext().getSpanId().equals(workflowSpan.getSpanId()))); + // The continuation segment also links to the initial logical operation span on the execution trace. Because the + // initial span ID is deterministic on (arn, operationId), it matches the original wait operation span's ID. + assertEquals( + originalWaitSpanId, + new DeterministicIdGenerator().generateSpanIdForOperation(arn, "op-2"), + "Initial logical operation span ID is deterministic on (arn, operationId)"); assertTrue( waitContinuation.getLinks().stream() - .noneMatch(link -> link.getSpanContext().getSpanId().equals(originalWaitSpanId)), - "Continuation spans must not fabricate a link to an uncheckpointed prior span context"); + .anyMatch(link -> link.getSpanContext().getSpanId().equals(originalWaitSpanId) + && link.getSpanContext().getTraceId().equals(workflowSpan.getTraceId())), + "Continuation span should link to the initial logical operation span on the execution trace"); } // ─── Cross-invocation step retry scenario ──────────────────────────── @@ -1691,8 +1725,11 @@ void multiInvocation_stepWaitStep_producesCorrectSpans() { void crossInvocation_stepRetry_attemptsParentedToRespectiveInvocations() { var arn = "arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1"; + // Same execution start time across invocations so the ARN-derived canonical trace ID is reproducible. + var startTime = Instant.now(); + // Invocation 1: step starts, attempt 1 fails, invocation suspended during retry poll - plugin.onInvocationStart(new InvocationInfo("req-1", arn, true, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-1", arn, true, startTime)); plugin.onOperationStart( new OperationInfo("op-1", "process-payment", "STEP", "Step", null, Instant.now(), null, null, false)); plugin.onUserFunctionStart( @@ -1734,7 +1771,7 @@ void crossInvocation_stepRetry_attemptsParentedToRespectiveInvocations() { spanExporter.reset(); // Invocation 2: step is replayed (continuation), attempt 2 executes and succeeds - plugin.onInvocationStart(new InvocationInfo("req-2", arn, false, Instant.now())); + plugin.onInvocationStart(new InvocationInfo("req-2", arn, false, startTime)); // isReplay=true: this operation already exists in the execution state plugin.onOperationStart( new OperationInfo("op-1", "process-payment", "STEP", "Step", null, Instant.now(), null, null, true)); @@ -1808,17 +1845,24 @@ void crossInvocation_stepRetry_attemptsParentedToRespectiveInvocations() { .findFirst() .orElseThrow(); - assertNotEquals(inv1InvocationTraceId, inv2InvocationTraceId); + // Every durable execution shares ONE trace: both invocations and the Workflow span are on it. + assertEquals(inv1InvocationTraceId, inv2InvocationTraceId); + assertEquals(inv2InvocationTraceId, workflowSpan.getTraceId()); assertTrue(inv1Spans.stream().allMatch(span -> span.getTraceId().equals(inv1InvocationTraceId))); - assertTrue(inv2Spans.stream() - .filter(span -> !span.getName().equals("Workflow")) - .allMatch(span -> span.getTraceId().equals(inv2InvocationTraceId))); + assertTrue(inv2Spans.stream().allMatch(span -> span.getTraceId().equals(inv2InvocationTraceId))); assertTrue(inv2OperationSpan.getLinks().stream() .anyMatch(link -> link.getSpanContext().getSpanId().equals(workflowSpan.getSpanId()))); + // The replay segment also links to the initial logical operation span on the execution trace. Its span ID is + // deterministic on (arn, operationId), so it matches invocation 1's operation span ID. + assertEquals( + inv1OperationSpan.getSpanId(), + new DeterministicIdGenerator().generateSpanIdForOperation(arn, "op-1"), + "Initial logical operation span ID is deterministic on (arn, operationId)"); assertTrue( inv2OperationSpan.getLinks().stream() - .noneMatch(link -> link.getSpanContext().getSpanId().equals(inv1OperationSpan.getSpanId())), - "Replay spans must not fabricate a link to an uncheckpointed prior span context"); + .anyMatch(link -> link.getSpanContext().getSpanId().equals(inv1OperationSpan.getSpanId()) + && link.getSpanContext().getTraceId().equals(workflowSpan.getTraceId())), + "Replay span should link to the initial logical operation span on the execution trace"); } // ─── Workflow span + links ─────────────────────────────────────────── @@ -1845,6 +1889,50 @@ void workflowSpan_notExportedOnNonTerminal() { "Workflow span must not be exported on a non-terminal invocation"); } + @Test + void workflowSpan_notExportedOnRetrying() { + // RETRYING is non-terminal, so the deferred Workflow span is neither materialized nor abandoned. + plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec1", true, Instant.now())); + plugin.onInvocationEnd(new InvocationEndInfo( + "req-1", "arn:exec1", true, InvocationStatus.RETRYING, new RuntimeException("transient"))); + + assertTrue( + spanExporter.getFinishedSpanItems().stream() + .noneMatch(s -> s.getName().equals("Workflow")), + "Workflow span must not be exported on a RETRYING invocation"); + } + + @Test + void deferredWorkflowSpan_whenExported_isEnded_andMatchesLinkedSpanId() { + // The Workflow span is created only at the terminal invocation, but operations that ran earlier linked to its + // deterministic context. When it is finally exported it must be ended and carry that same span ID. + plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec-wf", true, Instant.now())); + plugin.onOperationStart( + new OperationInfo("op-1", "step-a", "STEP", "Step", null, Instant.now(), null, null, false)); + plugin.onOperationEnd(new OperationEndInfo( + "op-1", + "step-a", + "STEP", + "Step", + null, + Instant.now(), + Instant.now(), + "SUCCEEDED", + null, + false, + null, + null)); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", "arn:exec-wf", true, InvocationStatus.SUCCEEDED, null)); + + var workflow = spanByName("Workflow"); + var operation = spanByName("step-a"); + assertTrue(workflow.hasEnded(), "Deferred Workflow span must be ended when materialized"); + assertTrue( + operation.getLinks().stream() + .anyMatch(l -> l.getSpanContext().getSpanId().equals(workflow.getSpanId())), + "Operation's Workflow link must resolve to the materialized Workflow span ID"); + } + @Test void operationAndAttemptSpans_linkToWorkflowSpan() { plugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec-wf", true, Instant.now())); @@ -1897,8 +1985,10 @@ void operationLinksToWorkflow_withXRayContext() { var xrayPlugin = new InvocationOtelPlugin( SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)), OtelPluginConfig.builder() - .contextExtractor( - () -> new ExtractedContext("5759e988bd862e3fe1be46a994272793", "53995c3f42cd8ad8")) + .contextExtractor(() -> new ExtractedContext( + "5759e988bd862e3fe1be46a994272793", + "53995c3f42cd8ad8", + ExtractedContext.Sampling.SAMPLED)) .enableMdc(false) .build()); xrayPlugin.onInvocationStart(new InvocationInfo("req-1", "arn:exec1", true, Instant.now())); diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/WorkflowRootSamplingBridgeTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/WorkflowRootSamplingBridgeTest.java new file mode 100644 index 000000000..3ac8e3b16 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/WorkflowRootSamplingBridgeTest.java @@ -0,0 +1,74 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.otel; + +import static org.junit.jupiter.api.Assertions.*; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.util.List; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for the agent-to-application root sampling bridge. */ +class WorkflowRootSamplingBridgeTest { + + @BeforeEach + @AfterEach + void clear() { + WorkflowRootSamplingBridge.clearForTest(); + } + + @Test + void publishesNothing_whenNeverCalled() { + assertNull(WorkflowRootSamplingBridge.publishedRootSampled()); + } + + @Test + void publishesTrue_whenSamplerAlwaysSamples() { + WorkflowRootSamplingBridge.publish(Sampler.alwaysOn()); + assertEquals(Boolean.TRUE, WorkflowRootSamplingBridge.publishedRootSampled()); + } + + @Test + void publishesFalse_whenSamplerNeverSamples() { + WorkflowRootSamplingBridge.publish(Sampler.alwaysOff()); + assertEquals(Boolean.FALSE, WorkflowRootSamplingBridge.publishedRootSampled()); + } + + @Test + void publishesNothing_whenSamplerDecidesPerAttribute() { + // The probes vary attributes; a sampler keyed on an attribute makes them disagree, so nothing is published and + // callers fall back to treating the trace as sampled. + WorkflowRootSamplingBridge.publish(attributeKeyedSampler()); + assertNull(WorkflowRootSamplingBridge.publishedRootSampled()); + } + + /** Samples only when the durable execution ARN attribute is present, so empty-attribute probes disagree. */ + private static Sampler attributeKeyedSampler() { + return new Sampler() { + @Override + public SamplingResult shouldSample( + Context parentContext, + String traceId, + String name, + SpanKind spanKind, + Attributes attributes, + List parentLinks) { + return attributes.get(SpanAttributes.DURABLE_EXECUTION_ARN) != null + ? SamplingResult.recordAndSample() + : SamplingResult.drop(); + } + + @Override + public String getDescription() { + return "attributeKeyedSampler"; + } + }; + } +} diff --git a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java index f04e1c27b..295432017 100644 --- a/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java +++ b/sdk-testing/src/main/java/software/amazon/lambda/durable/testing/LocalDurableTestRunner.java @@ -43,6 +43,10 @@ public class LocalDurableTestRunner { private final SerDes serDes; private final DurableConfig customerConfig; private final Instant executionStartTime = Instant.now(); + // The execution identity is fixed for the whole execution, matching the backend: the ARN and the EXECUTION + // operation ID stay stable across reinvocations, while only per-invocation values (the checkpoint token) change. + private final String executionName = UUID.randomUUID().toString(); + private final String executionOperationId = UUID.randomUUID().toString(); private LocalDurableTestRunner( TypeToken inputType, @@ -330,14 +334,14 @@ public void stopChainedInvoke(String name, ErrorObject error) { } private DurableExecutionInput createDurableInput(I input) { - var executionName = UUID.randomUUID().toString(); - var invocationId = UUID.randomUUID().toString(); + // The last ARN segment must equal the EXECUTION operation ID (ExecutionManager parses the ARN to find it), and + // both are stable across reinvocations so the execution keeps one identity — and one derived trace ID. var executionArn = String.format( "arn:aws:lambda:us-east-1:123456789012:function:test:$LATEST/durable-execution/%s/%s", - executionName, invocationId); + executionName, executionOperationId); var inputJson = serDes.serialize(input); var executionOp = Operation.builder() - .id(invocationId) + .id(executionOperationId) .name(executionName) .type(OperationType.EXECUTION) .status(OperationStatus.STARTED)