diff --git a/.github/workflows/otel-conformance-tests.yml b/.github/workflows/otel-conformance-tests.yml index b9ad6b605..c273c61b2 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@f9998f305f1e26423baf0c58148a3bd69120d5ef + uses: aws/aws-durable-execution-conformance-tests/.github/workflows/opentelemetry-orchestrator.yml@02d6dca971a38c13d94d6233d12f687e55b2a572 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 || '02d6dca971a38c13d94d6233d12f687e55b2a572' }} 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..94f5221b5 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,47 @@ 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 + +The plugin decides sampling once per invocation and applies that single decision to every durable span (Workflow, Invocation, operation, attempt), so the configured sampler is not re-invoked per span and the full decision — including `RECORD_ONLY` — is preserved. The decision follows this precedence, highest first: + +1. **Backend decision** — `Sampled=1` / `Sampled=0` in the propagated header is authoritative and always preserved, regardless of the configured sampler. +2. **Same-trace ambient span** — when the header carries no usable `Sampled` value but a valid ambient span (for example an auto-instrumentation Lambda handler span) is already on the execution's trace, the plugin follows that span's decision: sampled → sampled; unsampled but still recording → `RECORD_ONLY`; unsampled and not recording → dropped. +3. **Configured sampler (application-owned provider)** — when you pass a `SdkTracerProvider` to the plugin, its sampler is read directly and evaluated once with the trace ID, span name, and attributes. A trace-ID-ratio sampler therefore produces a stable decision across reinvocations (the trace ID is stable). +4. **Installed sampler (Java-agent path)** — when the agent owns the provider, it is behind a classloader boundary and its *effective* sampler (which another agent extension may have wrapped or replaced) cannot be reliably read at decision time. Rather than guess, the plugin **defers**: it installs a delegating sampler through the agent's autoconfiguration and lets that wrapper consult the agent's real sampler. The delegate's decision is honored in full — if your configured policy is `always_off`, a rate limiter, or a remote sampler (`xray`, `jaeger_remote`) that returns drop, the durable spans are dropped; they are **not** force-sampled. To avoid consuming a stateful or quota-based sampler once per span, the wrapper consults the delegate once per execution (keyed by trace ID) and reuses that decision for the execution's remaining durable spans within the invocation. + +For precise, provider-independent control, set an explicit `Sampled` value upstream (for example by enabling X-Ray active tracing) — that backend decision takes precedence over everything else. + ## Span Attributes ### Invocation Span @@ -305,7 +323,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 +332,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/ContextExtractor.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ContextExtractor.java index 5f3788139..0abb0c26d 100644 --- a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ContextExtractor.java +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/ContextExtractor.java @@ -3,20 +3,33 @@ package software.amazon.lambda.durable.otel; /** - * Extracts trace context from the Lambda runtime environment. + * Extracts the durable execution's propagated trace context from the Lambda runtime environment. * *

Implementations read trace context from various sources (X-Ray trace header, W3C traceparent, etc.) and return an * {@link ExtractedContext} containing the trace ID and optional parent span ID. * - *

Plugins use a valid ambient OpenTelemetry span as the invocation parent when one is available. This extractor is - * consulted only when no ambient span context is available, providing fallback propagation context from the runtime - * environment. + *

When it is called: the plugin invokes {@link #extract()} once at the start of every invocation, + * unconditionally — including when an ambient OpenTelemetry span is active. The extracted context is the durable + * execution's identity and is resolved with the following precedence: + * + *

    + *
  1. a valid extracted backend context anchors the execution trace (this is what makes the durable spans share one + * stable trace across reinvocations, so it takes precedence over the per-invocation ambient span); + *
  2. otherwise the execution is anchored on a deterministic synthetic root derived from the execution ARN; + *
  3. the ambient span is never adopted as the execution trace. It is used only to parent the Invocation span when it + * is already on the resolved trace, and otherwise correlated with a span link. + *
+ * + *

Implementation contract: because {@code extract()} runs on every invocation, implementations must + * be side-effect-free (or idempotent) and cheap, and must return the durable execution's own context — returning stale + * or unrelated context will displace the correct execution trace. */ @FunctionalInterface public interface ContextExtractor { /** - * Extracts fallback trace context from the runtime environment. + * Extracts the durable execution's propagated trace context from the runtime environment. Called once per + * invocation, unconditionally. * * @return the extracted context, or {@code null} if no context is available */ 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/DurableSampler.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java new file mode 100644 index 000000000..7ce304a54 --- /dev/null +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSampler.java @@ -0,0 +1,137 @@ +// 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.SdkTracerProviderBuilder; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +/** + * A delegating sampler that applies the durable execution's precomputed decision to durable spans and leaves every + * other span to the wrapped sampler. + * + *

The durable plugins decide sampling for the whole execution exactly once per invocation and attach the resulting + * {@link SamplingResult} to the {@link Context} used as the parent of each durable span (via + * {@link DurableSamplingDecision}). When {@code shouldSample} sees that decision on the parent context it returns it + * verbatim, so: + * + *

+ * + *

Spans without the durable decision on their parent context (ordinary application or auto-instrumentation spans) + * are delegated to the wrapped sampler unchanged, so the customer's sampling configuration governs everything outside + * the durable execution's own spans. + */ +final class DurableSampler implements Sampler { + + // Cap on the deferred-decision cache. The sampler is long-lived (installed once), so an unbounded map would grow + // per execution on a warm container. A modest LRU cap bounds memory while keeping the "consult once per execution" + // guarantee for the handful of executions active on one Lambda instance; an evicted entry at worst re-consults the + // delegate for a later span of that execution, which is a rare, benign degradation rather than a correctness bug. + private static final int MAX_CACHED_DEFERRED_DECISIONS = 256; + + private final Sampler delegate; + // Caches the delegate's decision for a deferred durable execution, keyed by canonical trace ID, so a stateful or + // quota-based delegate is consulted once per execution rather than per span. Access-ordered LRU, size-capped and + // synchronized (contention is low: at most one miss per execution). + private final Map deferredDecisions = + Collections.synchronizedMap(new LinkedHashMap<>(16, 0.75f, true) { + @Override + protected boolean removeEldestEntry(Map.Entry eldest) { + return size() > MAX_CACHED_DEFERRED_DECISIONS; + } + }); + + private DurableSampler(Sampler delegate) { + this.delegate = delegate; + } + + /** + * Wraps {@code delegate} so durable spans use the precomputed decision. A null delegate or an already-wrapped + * delegate is handled defensively: wrapping is idempotent, and a null delegate falls back to + * {@link Sampler#parentBased(Sampler)} of {@link Sampler#alwaysOn()} (the OTel/ADOT default) so ordinary spans + * still have a sampler. + */ + static DurableSampler wrap(Sampler delegate) { + if (delegate instanceof DurableSampler durableSampler) { + return durableSampler; + } + var effectiveDelegate = delegate != null ? delegate : Sampler.parentBased(Sampler.alwaysOn()); + return new DurableSampler(effectiveDelegate); + } + + /** + * Installs the durable sampler on an application-owned {@link SdkTracerProviderBuilder} by wrapping its configured + * sampler. The builder exposes only a setter, so the configured sampler is read reflectively (mirroring how the + * deterministic ID generator wraps the builder's ID generator) and replaced with a wrapper that delegates to it. + * Installation is idempotent: a builder whose sampler is already a {@link DurableSampler} is left unchanged. + */ + static void installOn(SdkTracerProviderBuilder builder) { + var configured = configuredSampler(builder); + if (configured instanceof DurableSampler) { + return; + } + builder.setSampler(wrap(configured)); + } + + private static Sampler configuredSampler(SdkTracerProviderBuilder builder) { + for (var field : builder.getClass().getDeclaredFields()) { + if (java.lang.reflect.Modifier.isStatic(field.getModifiers()) + || !Sampler.class.isAssignableFrom(field.getType())) { + continue; + } + try { + if (!field.trySetAccessible()) { + break; + } + return (Sampler) field.get(builder); + } catch (IllegalAccessException e) { + throw new IllegalStateException("Unable to read the configured OpenTelemetry sampler", e); + } + } + throw new IllegalStateException("Unable to locate the configured OpenTelemetry sampler"); + } + + @Override + public SamplingResult shouldSample( + Context parentContext, + String traceId, + String name, + SpanKind spanKind, + Attributes attributes, + List parentLinks) { + var intent = DurableSamplingDecision.get(parentContext); + if (intent == null) { + // Not a durable span: the customer's sampler governs it unchanged. + return delegate.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); + } + if (!intent.isDeferred()) { + // A resolved decision (explicit upstream, same-trace ambient, or a locally reproduced sampler): return it + // verbatim, preserving RECORD_ONLY and never re-invoking the delegate. + return intent.resolved(); + } + // Deferred (agent path, real sampler not reproducible here): evaluate the actual delegate once per execution + // and reuse it, so an installed drop/rate-limit policy is honored and consulted only once. + return deferredDecisions.computeIfAbsent( + intent.deferredTraceId(), + key -> delegate.shouldSample(Context.root(), key, name, spanKind, attributes, Collections.emptyList())); + } + + @Override + public String getDescription() { + return "DurableSampler{" + delegate.getDescription() + "}"; + } +} diff --git a/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java new file mode 100644 index 000000000..c43d8eb02 --- /dev/null +++ b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/DurableSamplingDecision.java @@ -0,0 +1,150 @@ +// 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.context.Context; +import io.opentelemetry.context.ContextKey; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; + +/** + * Carries the durable execution's sampling intent to {@link DurableSampler} for one durable span. + * + *

Computed once per invocation (see the plugins' {@code onInvocationStart}), the intent is one of two forms: + * + *

+ * + *

Two carriers, because the plugin runs across two class loaders. Under the documented ADOT setup the plugin + * JAR is loaded twice — once by the application class loader (which computes the intent) and once by the Java-agent + * extension class loader (which installs and runs {@link DurableSampler}). A {@link ContextKey} uses reference + * identity, so a key created in one loader is not equal to the key created in the other. To bridge this: + * + *

    + *
  1. Context key — used when both sides share a class loader (an application-owned provider). It preserves + * the full {@link SamplingResult}, including any attributes a custom sampler attached. + *
  2. Thread-scoped system property — a cross-class-loader fallback modelled on + * {@link DeterministicIdGenerator}'s scoped-ID bridge. The intent is published on the thread that creates the + * durable span for the synchronous duration of {@code startSpan()} (the sampler runs on that same thread), keyed + * by thread ID under a bootstrap-visible {@link System} property so both class loaders read the same value. A + * resolved decision bridges its {@link SamplingDecision} name (the three built-in decisions carry no attributes, + * so reconstructing them is faithful); a deferral bridges a sentinel plus the canonical trace ID. + *
+ * + *

The scope is opened immediately around each durable {@code startSpan()} call and closed right after, so the + * property never leaks beyond the span it applies to. Nothing is persisted across invocations; cross-invocation + * consistency comes from recomputing the intent from stable inputs, not from sharing state. + */ +final class DurableSamplingDecision { + + /** + * The durable sampling intent for a span: either a resolved {@link SamplingResult}, or a deferral to the agent-side + * sampler's own delegate keyed by the canonical trace ID. + * + * @param resolved the resolved decision, or null when deferring + * @param deferredTraceId the canonical trace ID to key the agent-side delegate cache on, or null when resolved + */ + record Intent(SamplingResult resolved, String deferredTraceId) { + static Intent resolved(SamplingResult result) { + return new Intent(result, null); + } + + static Intent deferred(String traceId) { + return new Intent(null, traceId); + } + + boolean isDeferred() { + return resolved == null; + } + } + + private static final ContextKey KEY = + ContextKey.named("software.amazon.lambda.durable.otel.durable-sampling-decision"); + + private static final String SCOPED_PROPERTY_PREFIX = "software.amazon.lambda.durable.otel.scopedSamplingDecision."; + // Sentinel prefix distinguishing a deferral (carrying the trace ID) from a resolved decision name. + private static final String DEFERRED_PREFIX = "DEFER:"; + + private DurableSamplingDecision() {} + + /** Returns a context carrying the durable sampling intent (same-class-loader carrier), derived from the given. */ + static Context store(Context context, Intent intent) { + return context.with(KEY, intent); + } + + /** + * Publishes the intent on the current thread for the duration of the returned scope, bridging it across the + * application and Java-agent class loaders. Callers open this immediately around a durable {@code startSpan()} call + * (which runs the sampler synchronously on this thread) and close it right after. + */ + static Scope openScope(Intent intent) { + var key = scopedProperty(); + var previous = System.getProperty(key); + System.setProperty(key, encode(intent)); + return () -> { + if (previous == null) { + System.clearProperty(key); + } else { + System.setProperty(key, previous); + } + }; + } + + /** + * Returns the durable sampling intent for a span, or {@code null} when none is present. Prefers the full-fidelity + * context key (same class loader) and falls back to the thread-scoped system property (cross class loader). + */ + static Intent get(Context context) { + var fromContext = context.get(KEY); + if (fromContext != null) { + return fromContext; + } + return fromScopedProperty(); + } + + private static String encode(Intent intent) { + return intent.isDeferred() + ? DEFERRED_PREFIX + intent.deferredTraceId() + : intent.resolved().getDecision().name(); + } + + private static Intent fromScopedProperty() { + var value = System.getProperty(scopedProperty()); + if (value == null) { + return null; + } + if (value.startsWith(DEFERRED_PREFIX)) { + return Intent.deferred(value.substring(DEFERRED_PREFIX.length())); + } + return Intent.resolved( + switch (SamplingDecision.valueOf(value)) { + case RECORD_AND_SAMPLE -> SamplingResult.recordAndSample(); + case RECORD_ONLY -> SamplingResult.recordOnly(); + case DROP -> SamplingResult.drop(); + }); + } + + private static String scopedProperty() { + return SCOPED_PROPERTY_PREFIX + Thread.currentThread().getId(); + } + + static void clearSharedStateForTest() { + System.getProperties().stringPropertyNames().stream() + .filter(name -> name.startsWith(SCOPED_PROPERTY_PREFIX)) + .toList() + .forEach(System::clearProperty); + } + + /** A closeable scope that restores the previous thread-scoped intent when closed. */ + interface Scope extends AutoCloseable { + @Override + void close(); + } +} 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..9cd8ac33a 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,9 +39,10 @@ * durable-execution hierarchy: * *

* - *

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 +68,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. @@ -99,6 +91,16 @@ public class InvocationOtelPlugin implements DurableExecutionPlugin { 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; + // The execution's single sampling intent for this invocation, computed once at onInvocationStart and attached to + // every durable span's parent context so DurableSampler applies it (a resolved decision verbatim, or a deferral to + // its own delegate) without re-invoking the configured sampler per span. + private volatile DurableSamplingDecision.Intent samplingIntent; + + /** 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<>(); @@ -162,6 +164,8 @@ public InvocationOtelPlugin() { */ public InvocationOtelPlugin(SdkTracerProviderBuilder tracerProviderBuilder, OtelPluginConfig config) { this.idGenerator = DeterministicIdGenerator.installOn(tracerProviderBuilder); + // Wrap the configured sampler so durable spans use the execution's single precomputed decision. + DurableSampler.installOn(tracerProviderBuilder); this.sdkTracerProvider = tracerProviderBuilder.build(); this.tracer = sdkTracerProvider.get(config.instrumentationName()); @@ -198,40 +202,47 @@ 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 + var extracted = contextExtractor.extract(); + + // Resolve the execution ancestor the Workflow span parents onto so it joins the stable-per-execution trace. + var canonicalTraceId = ExecutionTraceContext.canonicalTraceId( + extracted, info.durableExecutionArn(), info.executionStartTime(), idGenerator); + // Resolve the execution's sampling decision once for this invocation as a full SamplingResult, then apply it to + // every durable span via DurableSampler (see below). The execution ancestor's trace flags are derived from the + // same decision so a parent-based sampler stays consistent with it. + var decision = OtelPluginSupport.resolveSamplingResult( + sdkTracerProvider, + extracted, + Span.current(), + canonicalTraceId, + workflowSpanName, + Attributes.of(DURABLE_EXECUTION_ARN, info.durableExecutionArn())); + // A null decision is unresolved on the agent path: defer to DurableSampler's own delegate (keyed by trace ID), + // rather than fabricating a decision that would bypass an installed drop/rate-limit policy. + samplingIntent = decision != null + ? DurableSamplingDecision.Intent.resolved(decision) + : DurableSamplingDecision.Intent.deferred(canonicalTraceId); + var sampled = OtelPluginSupport.isSampled(decision); + var execCtx = ExecutionTraceContext.resolve( + extracted, canonicalTraceId, info.durableExecutionArn(), idGenerator, () -> sampled); + executionTrace = new ExecutionTrace(canonicalTraceId, execCtx.traceFlags()); + + // Workflow span — one logical span per durable execution, parented onto the execution ancestor so it joins the + // execution trace. 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() + .setParent(withDurableDecision(Context.root().with(Span.wrap(execCtx.executionAncestor())))) .setAttribute(DURABLE_EXECUTION_ARN, info.durableExecutionArn()) .setStartTimestamp(info.executionStartTime()); - var workflowTraceId = - idGenerator.generateTraceIdForExecution(info.durableExecutionArn(), info.executionStartTime()); 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(); - } + // Force the span ID only; the trace ID comes from the parent so the Workflow span joins the execution trace. + workflowSpan = startDurableSpan(workflowSpanBuilder, null, workflowSpanId); + + // 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") @@ -244,7 +255,7 @@ public void onInvocationStart(InvocationInfo info) { spanBuilder.setAttribute(AttributeKey.stringKey("faas.invocation_id"), info.requestId()); } - invocationSpan = spanBuilder.startSpan(); + invocationSpan = startDurableSpan(spanBuilder); // Inject MDC on the handler thread so handler-level logs (between steps) have trace context. // This runs on the same thread as context.getLogger() calls in the handler. @@ -300,6 +311,7 @@ public void onInvocationEnd(InvocationEndInfo info) { invocationSpan.end(); invocationSpan = null; + samplingIntent = null; // End the Workflow span only on a terminal status, so it is exported exactly once per execution // (SUCCEEDED -> OK, FAILED -> ERROR; non-terminal statuses leave it un-ended / not exported this invocation). @@ -349,6 +361,12 @@ public void onOperationStart(OperationInfo info) { .setAttribute(DURABLE_OPERATION_TYPE, info.type()) .setAttribute(DURABLE_OPERATION_STATUS, info.status() != null ? info.status() : "STARTED"); + // 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 first, then to the Workflow span, so the links are ordered + // [operation, Workflow]. A non-replay operation span carries only the Workflow link. + if (info.isReplay()) { + addInitialOperationLink(spanBuilder, info.id()); + } // Link to the Workflow span for execution-level correlation (operation stays parented to the invocation span). addWorkflowLink(spanBuilder); @@ -360,8 +378,8 @@ public void onOperationStart(OperationInfo info) { } var span = info.isReplay() - ? spanBuilder.startSpan() - : idGenerator.startSpan( + ? startDurableSpan(spanBuilder) + : startDurableSpan( spanBuilder, null, idGenerator.generateSpanIdForOperation(durableExecutionArn, info.id())); // Store the open span — will be ended in onOperationEnd or onInvocationEnd @@ -405,6 +423,9 @@ public void onOperationEnd(OperationEndInfo info) { .setAttribute(DURABLE_EXECUTION_ARN, durableExecutionArn) .setAttribute(DURABLE_OPERATION_ID, info.id()) .setAttribute(DURABLE_OPERATION_TYPE, info.type()); + // This continuation segment completes an operation whose initial span ran earlier; link back to that + // initial operation span first, then to the Workflow span (the spec expects [operation, Workflow]). + addInitialOperationLink(spanBuilder, info.id()); addWorkflowLink(spanBuilder); if (info.name() != null) { @@ -414,7 +435,7 @@ public void onOperationEnd(OperationEndInfo info) { spanBuilder.setAttribute(DURABLE_OPERATION_SUBTYPE, info.subType()); } - var continuationSpan = spanBuilder.startSpan(); + var continuationSpan = startDurableSpan(spanBuilder); if (info.status() != null) { continuationSpan.setAttribute(DURABLE_OPERATION_STATUS, info.status()); @@ -485,7 +506,7 @@ public void onUserFunctionStart(UserFunctionStartInfo info) { spanBuilder.setAttribute(DURABLE_ATTEMPT_NUMBER, info.attempt().longValue()); } - var span = spanBuilder.startSpan(); + var span = startDurableSpan(spanBuilder); attemptSpans.put(key, span); // Make span current on this thread so auto-instrumented calls become children @@ -589,18 +610,67 @@ 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 withDurableDecision(Context.root().with(Span.current())); + } + return withDurableDecision(Context.root().with(Span.wrap(execCtx.executionAncestor()))); + } + private Context resolveParentContext(String parentId) { if (parentId != null) { var parentSpanContext = operationContexts.get(parentId); if (parentSpanContext != null) { - return Context.current().with(Span.wrap(parentSpanContext)); + return withDurableDecision(Context.current().with(Span.wrap(parentSpanContext))); } } // Fall back to invocation span as parent if (invocationSpan != null) { - return Context.current().with(invocationSpan); + return withDurableDecision(Context.current().with(invocationSpan)); + } + return withDurableDecision(Context.current()); + } + + /** + * Attaches the execution's sampling intent to a durable span's parent context so {@link DurableSampler} applies it + * (a resolved decision verbatim, or a deferral to its own delegate) instead of re-invoking the configured sampler + * per span. When no intent has been resolved (telemetry disabled for the invocation) the context is unchanged. + */ + private Context withDurableDecision(Context context) { + var intent = samplingIntent; + return intent != null ? DurableSamplingDecision.store(context, intent) : context; + } + + /** + * Starts a durable span with the execution's sampling intent published on the current thread for the duration of + * the sampler call, so {@link DurableSampler} applies it even when the plugin and the agent-installed sampler run + * in different class loaders (see {@link DurableSamplingDecision}). Falls back to a plain start when no intent has + * been resolved. + */ + private Span startDurableSpan(SpanBuilder spanBuilder) { + var intent = samplingIntent; + if (intent == null) { + return spanBuilder.startSpan(); + } + try (var ignored = DurableSamplingDecision.openScope(intent)) { + return spanBuilder.startSpan(); + } + } + + /** Starts a durable span with a forced span ID, publishing the sampling intent as in {@link #startDurableSpan}. */ + private Span startDurableSpan(SpanBuilder spanBuilder, String traceId, String spanId) { + var intent = samplingIntent; + if (intent == null) { + return idGenerator.startSpan(spanBuilder, traceId, spanId); + } + try (var ignored = DurableSamplingDecision.openScope(intent)) { + return idGenerator.startSpan(spanBuilder, traceId, spanId); } - return Context.current(); } /** Adds a link to the Workflow span, if one exists, for execution-level correlation. */ @@ -611,6 +681,24 @@ private void addWorkflowLink(SpanBuilder spanBuilder) { } } + /** + * 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 +724,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..d9cd38ec0 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 so durable spans get + * deterministic IDs. + */ public final class OtelPluginAutoConfigurationCustomizerProvider implements AutoConfigurationCustomizerProvider { @Override @@ -15,6 +18,9 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) { DeterministicIdGenerator.installOn(builder); return builder; }); + // Wrap the agent-configured sampler so durable spans use the execution's single precomputed decision, applied + // through the durable span's parent context, instead of re-invoking the configured sampler per span. + autoConfiguration.addSamplerCustomizer((sampler, config) -> DurableSampler.wrap(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..63fbec908 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,17 @@ package software.amazon.lambda.durable.otel; import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; +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.SamplingResult; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Collections; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -24,6 +29,116 @@ static DeterministicIdGenerator createDefaultIdGenerator() { return new DeterministicIdGenerator(); } + /** + * Resolves the durable execution's sampling decision for one invocation as a full {@link SamplingResult}, evaluated + * exactly once. The decision is applied to every durable span of the invocation (Workflow, Invocation, operation, + * attempt) via {@link DurableSampler}, so the configured sampler is not re-invoked per span and the three-way + * decision — including {@code RECORD_ONLY} — is preserved rather than reduced to a boolean. + * + *

Sampling precedence, highest first: + * + *

    + *
  1. Explicit upstream {@code Sampled} on the propagated header ({@code Sampled=1} /{@code Sampled=0}). + * An authoritative backend decision is preserved regardless of the configured sampler; + *
  2. Same-trace ambient span. When no explicit {@code Sampled} is present but a valid ambient span (an + * auto-instrumentation Lambda handler span, {@code Span.current()}) is already on the canonical execution + * trace, its established decision is followed. The ambient span is a same-trace descendant of the propagated + * parent, so its decision is representative of this trace. Its full three-way decision is preserved: a + * sampled span yields {@code RECORD_AND_SAMPLE}; an unsampled but recording span yields {@code RECORD_ONLY} + * (its spans still reach processors); only an unsampled, non-recording span yields {@code DROP}; + *
  3. Application-owned provider: configured sampler, once. When the tracer provider is reachable (the + * two-argument constructor path), its sampler is read directly and evaluated a single time with + * {@code ROOT_CONTEXT} (so a parent-based sampler applies its root policy), the canonical trace ID, span + * name, and attributes, and its full result is returned; + *
  4. Java-agent path: defer to the installed sampler. When the provider is not visible + * ({@code sdkTracerProvider == null}), this returns {@code null} to defer. It does not reconstruct + * the sampler from environment settings: the agent's effective sampler is whatever its autoconfiguration + * pipeline finally installs, and another extension's customizer can wrap or replace a recognized configured + * sampler, so a reconstruction could disagree with the real delegate. Deferring routes the decision to the + * agent-installed {@link DurableSampler}, which consults its actual delegate once per execution, caches the + * result by trace ID, and reuses it for the execution's remaining durable spans (see + * {@link DurableSampler#shouldSample}). The delegate's decision is honored in full — including a + * {@code DROP}/rate-limited outcome — so durable spans are not force-sampled. + *
+ * + *

For the app-path branches, nothing is persisted across invocations: each invocation recomputes the decision + * from stable inputs (the canonical trace ID and the upstream {@code Sampled} value), so deterministic samplers + * reach the same decision on every reinvocation. The ambient-span branch is inherently per-invocation because the + * ambient span is recreated each invocation, which the shared spec accepts for same-trace ambient membership. On + * the agent path the delegate is consulted once per execution within a process and cached; the delegate itself (for + * example a ratio sampler) governs cross-invocation consistency. + * + * @param sdkTracerProvider the resolved provider, or null when it is not visible to the application (agent path) + * @param extracted the context parsed from the propagated header, or null when none is present + * @param ambientSpan the current ambient span ({@code Span.current()}); its context may be invalid and its + * recording state distinguishes RECORD_ONLY from DROP for an unsampled same-trace parent + * @param canonicalTraceId the canonical execution trace ID (used for the ambient same-trace check and the sampler) + * @param spanName the span name passed to the sampler + * @param attributes the attributes the span is started with + * @return the resolved {@link SamplingResult}, or {@code null} when the decision is deferred to the agent-side + * {@link DurableSampler}'s own delegate + */ + static SamplingResult resolveSamplingResult( + SdkTracerProvider sdkTracerProvider, + ExtractedContext extracted, + Span ambientSpan, + String canonicalTraceId, + String spanName, + Attributes attributes) { + // 1. An explicit upstream decision is authoritative. + if (extracted != null) { + switch (extracted.sampling()) { + case SAMPLED: + return SamplingResult.recordAndSample(); + case NOT_SAMPLED: + return SamplingResult.drop(); + case UNDECIDED: + break; + } + } + // 2. No explicit decision: follow a valid same-trace ambient span's established decision. The sampled bit alone + // cannot distinguish RECORD_ONLY (recording, unsampled) from DROP (not recording, unsampled), so a sampled bit + // maps to RECORD_AND_SAMPLE, an unsampled-but-recording span maps to RECORD_ONLY (its spans still reach + // processors), and only an unsampled, non-recording span maps to DROP. + var ambient = ambientSpan != null ? ambientSpan.getSpanContext() : null; + if (ambient != null && ambient.isValid() && ambient.getTraceId().equals(canonicalTraceId)) { + if (ambient.getTraceFlags().isSampled()) { + return SamplingResult.recordAndSample(); + } + return ambientSpan.isRecording() ? SamplingResult.recordOnly() : SamplingResult.drop(); + } + // 3. An application-owned provider exposes the real sampler: evaluate it once, preserving its full result. + if (sdkTracerProvider != null) { + return sdkTracerProvider + .getSampler() + .shouldSample( + Context.root(), + canonicalTraceId, + spanName, + SpanKind.INTERNAL, + attributes, + Collections.emptyList()); + } + // 4. Agent path (provider not visible from the application class loader): defer. The agent's effective sampler + // is whatever the autoconfiguration pipeline finally installs — a recognized configured sampler can be wrapped + // or replaced by another extension's customizer, so it cannot be reliably reconstructed from environment + // settings here. Return null so the agent-installed DurableSampler consults its actual delegate once per + // execution and caches the result, honoring the customer's effective policy (including drop/rate-limit). + return null; + } + + /** + * True when the decision records and samples, used to derive the (never-exported) execution-ancestor trace flags. A + * {@code null} decision is unresolved (deferred to the agent-side sampler); it defaults the ancestor flag to + * sampled so a parent-based delegate is not biased toward dropping, while the agent-side {@link DurableSampler} + * still makes the authoritative per-span decision from its real delegate. + */ + static boolean isSampled(SamplingResult samplingResult) { + return samplingResult == null + || samplingResult.getDecision() + == io.opentelemetry.sdk.trace.samplers.SamplingDecision.RECORD_AND_SAMPLE; + } + /** The tracer provider and tracer resolved from the global OpenTelemetry instance. */ record ProviderSetup(SdkTracerProvider sdkTracerProvider, Tracer tracer) {} @@ -71,15 +186,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/XRayContextExtractor.java b/otel-plugin/src/main/java/software/amazon/lambda/durable/otel/XRayContextExtractor.java index 109f633de..a42b5ee5c 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 @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.otel; +import io.opentelemetry.api.trace.SpanId; +import io.opentelemetry.api.trace.TraceId; import java.util.regex.Pattern; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -28,11 +30,11 @@ public class XRayContextExtractor implements ContextExtractor { @Override public ExtractedContext extract() { - // Try system property first — Lambda runtime updates this per invocation - // and it avoids JVM environment variable caching issues. + // Try system property first — the Lambda runtime interface client updates this per invocation, so it reflects + // the current invocation and avoids the JVM's process-lifetime environment-variable caching. var traceHeader = System.getProperty(XRAY_SYSTEM_PROPERTY); if (traceHeader == null || traceHeader.isEmpty()) { - // Fallback to environment variable + // Fallback to the environment variable (non-agent environments, local runs). traceHeader = System.getenv(XRAY_ENV_VAR); } if (traceHeader == null || traceHeader.isEmpty()) { @@ -42,6 +44,7 @@ public ExtractedContext extract() { String root = null; String parent = null; + String sampled = null; for (var part : traceHeader.split(";")) { var eqIdx = part.indexOf('='); @@ -52,6 +55,7 @@ public ExtractedContext extract() { switch (key) { case "Root" -> root = value; case "Parent" -> parent = value; + case "Sampled" -> sampled = value; } } @@ -60,24 +64,34 @@ public ExtractedContext extract() { return null; } - // Root format: 1-5759e988-bd862e3fe1be46a994272793 - // Strip "1-" prefix, remove dashes → 32-char hex OTel trace ID + // Root format: 1-5759e988-bd862e3fe1be46a994272793 → strip "1-" and dashes → 32-char hex OTel trace ID. Reject + // an all-zero Root: it is well-formed but an invalid trace ID, and reusing it would anchor the execution on an + // invalid trace. var traceId = xrayRootToOtelTraceId(root); - if (traceId == null) { + if (traceId == null || !TraceId.isValid(traceId)) { logger.debug("Invalid X-Ray Root field: {}", root); 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). Reject an all-zero Parent + // the same way — it is well-formed but invalid. String parentSpanId = null; if (parent != null) { var normalized = parent.toLowerCase(); - if (HEX_16.matcher(normalized).matches()) { + if (HEX_16.matcher(normalized).matches() && SpanId.isValid(normalized)) { parentSpanId = normalized; } } - 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/DurableSamplerTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplerTest.java new file mode 100644 index 000000000..6fea5784e --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplerTest.java @@ -0,0 +1,273 @@ +// 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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.opentelemetry.api.GlobalOpenTelemetry; +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.SpanKind; +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer; +import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties; +import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; +import io.opentelemetry.sdk.trace.SdkTracerProvider; +import io.opentelemetry.sdk.trace.data.LinkData; +import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.function.BiFunction; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; +import software.amazon.lambda.durable.plugin.OperationEndInfo; +import software.amazon.lambda.durable.plugin.OperationInfo; +import software.amazon.lambda.durable.plugin.UserFunctionEndInfo; +import software.amazon.lambda.durable.plugin.UserFunctionOutcome; +import software.amazon.lambda.durable.plugin.UserFunctionStartInfo; + +/** + * Tests the delegating {@link DurableSampler} and its end-to-end effect through the plugin: the configured sampler is + * evaluated at most once per invocation, the full decision (including {@code RECORD_ONLY}) is preserved, and explicit + * upstream decisions are authoritative over a conflicting configured sampler. + */ +class DurableSamplerTest { + + private static final String ARN = "arn:aws:lambda:us-east-1:123:function:test:$LATEST/durable/exec1"; + private static final String TRACE_ID = "aabbccddee112233445566778899aabb"; + private static final String SPAN_ID = "1111111111111111"; + + @BeforeEach + void setUp() { + DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + } + + @AfterEach + void tearDown() { + GlobalOpenTelemetry.resetForTest(); + DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + } + + // ─── Unit tests for the wrapper ────────────────────────────────────── + + @Test + void shouldSample_returnsCarriedDecision_forDurableSpans() { + var delegate = new CountingSampler(Sampler.alwaysOn()); + var sampler = DurableSampler.wrap(delegate); + var parent = DurableSamplingDecision.store( + Context.root(), DurableSamplingDecision.Intent.resolved(SamplingResult.drop())); + + var result = sampler.shouldSample(parent, TRACE_ID, "op", SpanKind.INTERNAL, Attributes.empty(), List.of()); + + assertEquals(SamplingDecision.DROP, result.getDecision(), "The carried decision is returned verbatim"); + assertEquals(0, delegate.count(), "The delegate is not invoked when a resolved decision is present"); + } + + @Test + void shouldSample_preservesRecordOnly() { + var sampler = DurableSampler.wrap(Sampler.alwaysOn()); + var parent = DurableSamplingDecision.store( + Context.root(), DurableSamplingDecision.Intent.resolved(SamplingResult.recordOnly())); + + var result = sampler.shouldSample(parent, TRACE_ID, "op", SpanKind.INTERNAL, Attributes.empty(), List.of()); + + assertEquals(SamplingDecision.RECORD_ONLY, result.getDecision(), "RECORD_ONLY is preserved, not collapsed"); + } + + @Test + void shouldSample_delegates_forNonDurableSpans() { + var delegate = new CountingSampler(Sampler.alwaysOff()); + var sampler = DurableSampler.wrap(delegate); + + var result = sampler.shouldSample( + Context.root(), TRACE_ID, "other", SpanKind.INTERNAL, Attributes.empty(), List.of()); + + assertEquals(SamplingDecision.DROP, result.getDecision(), "Non-durable spans use the delegate"); + assertEquals(1, delegate.count(), "The delegate governs spans without a durable decision"); + } + + @Test + void wrap_isIdempotent() { + var wrapped = DurableSampler.wrap(Sampler.alwaysOn()); + assertSame(wrapped, DurableSampler.wrap(wrapped), "Wrapping an already-wrapped sampler returns it unchanged"); + } + + @Test + void deferredIntent_evaluatesRealDelegate_notBypassed() { + // Agent path: the real sampler could not be reproduced, so the plugin defers. The wrapper must consult its + // actual delegate (here always_off) rather than a fabricated sampled decision, so durable spans are dropped. + var delegate = new CountingSampler(Sampler.alwaysOff()); + var sampler = DurableSampler.wrap(delegate); + var parent = DurableSamplingDecision.store(Context.root(), DurableSamplingDecision.Intent.deferred(TRACE_ID)); + + var result = sampler.shouldSample(parent, TRACE_ID, "op", SpanKind.INTERNAL, Attributes.empty(), List.of()); + + assertEquals(SamplingDecision.DROP, result.getDecision(), "The real delegate governs a deferred decision"); + } + + @Test + void deferredIntent_evaluatesDelegateOncePerExecution() { + // A stateful/quota delegate must be consulted once per execution (trace ID), not per durable span. + var delegate = new CountingSampler(Sampler.alwaysOn()); + var sampler = DurableSampler.wrap(delegate); + var parent = DurableSamplingDecision.store(Context.root(), DurableSamplingDecision.Intent.deferred(TRACE_ID)); + + for (var i = 0; i < 4; i++) { + sampler.shouldSample(parent, TRACE_ID, "op" + i, SpanKind.INTERNAL, Attributes.empty(), List.of()); + } + + assertEquals( + 1, delegate.count(), "The deferred delegate is evaluated once per execution and cached by trace ID"); + } + + @Test + void agentCustomizer_wrapsEffectiveSampler_evenWhenAnotherExtensionChangedIt() { + // Another agent extension can wrap or replace the recognized configured sampler before ours runs. Our sampler + // customizer must wrap whatever effective sampler it receives, so a deferred durable span follows that + // effective delegate — not a reconstruction from OTEL_TRACES_SAMPLER. Here the effective sampler is always_off + // (as if a prior customizer replaced a configured always_on), so the deferred durable span is dropped. + var effectiveSampler = new CountingSampler(Sampler.alwaysOff()); + var installed = captureInstalledSampler(effectiveSampler); + + var parent = DurableSamplingDecision.store(Context.root(), DurableSamplingDecision.Intent.deferred(TRACE_ID)); + var result = installed.shouldSample(parent, TRACE_ID, "op", SpanKind.INTERNAL, Attributes.empty(), List.of()); + + assertEquals( + SamplingDecision.DROP, + result.getDecision(), + "The customizer must wrap the effective (possibly replaced) sampler, not a reconstruction"); + assertEquals(1, effectiveSampler.count(), "The effective delegate is the one consulted"); + } + + /** Runs the agent-path sampler customizer over the given effective sampler and returns what it installs. */ + @SuppressWarnings("unchecked") + private static Sampler captureInstalledSampler(Sampler effectiveSampler) { + var customizer = org.mockito.Mockito.mock(AutoConfigurationCustomizer.class, org.mockito.Mockito.RETURNS_SELF); + var samplerCustomizerCaptor = org.mockito.ArgumentCaptor.forClass(BiFunction.class); + + new OtelPluginAutoConfigurationCustomizerProvider().customize(customizer); + + org.mockito.Mockito.verify(customizer).addSamplerCustomizer(samplerCustomizerCaptor.capture()); + BiFunction samplerCustomizer = samplerCustomizerCaptor.getValue(); + return samplerCustomizer.apply(effectiveSampler, null); + } + + // ─── End-to-end tests through the plugin ───────────────────────────── + + @Test + void configuredSampler_isEvaluatedAtMostOncePerInvocation() { + var delegate = new CountingSampler(Sampler.alwaysOn()); + var exporter = InMemorySpanExporter.create(); + var plugin = new InvocationOtelPlugin( + SdkTracerProvider.builder().setSampler(delegate).addSpanProcessor(SimpleSpanProcessor.create(exporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> null) + .enableMdc(false) + .build()); + + // A full invocation with a Workflow span, Invocation span, operation span, and attempt span. + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + plugin.onOperationStart( + new OperationInfo("op-1", "step", "STEP", "Step", null, Instant.now(), null, null, false)); + plugin.onUserFunctionStart( + new UserFunctionStartInfo("op-1", "step", "STEP", "Step", null, Instant.now(), false, 1)); + plugin.onUserFunctionEnd(new UserFunctionEndInfo( + "op-1", + "step", + "STEP", + "Step", + null, + Instant.now(), + Instant.now(), + false, + 1, + UserFunctionOutcome.SUCCEEDED, + null)); + plugin.onOperationEnd(new OperationEndInfo( + "op-1", "step", "STEP", "Step", null, Instant.now(), Instant.now(), "SUCCEEDED", 1, false, null)); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); + + assertTrue( + delegate.count() <= 1, + "The configured sampler must be evaluated at most once per invocation, not per span; was " + + delegate.count()); + } + + @Test + void explicitSampled_winsOverConfiguredAlwaysOff() { + var exporter = exportedWith(Sampler.alwaysOff(), ExtractedContext.Sampling.SAMPLED); + assertFalse(exporter.getFinishedSpanItems().isEmpty(), "Explicit Sampled=1 exports spans despite always_off"); + } + + @Test + void explicitNotSampled_winsOverConfiguredAlwaysOn() { + var exporter = exportedWith(Sampler.alwaysOn(), ExtractedContext.Sampling.NOT_SAMPLED); + assertTrue(exporter.getFinishedSpanItems().isEmpty(), "Explicit Sampled=0 drops spans despite always_on"); + } + + /** + * Runs a minimal invocation with the given configured sampler and an extractor that supplies a complete remote + * parent carrying the given explicit upstream sampling decision, and returns the exporter. + */ + private InMemorySpanExporter exportedWith(Sampler configuredSampler, ExtractedContext.Sampling sampling) { + var exporter = InMemorySpanExporter.create(); + var extracted = new ExtractedContext(TRACE_ID, SPAN_ID, sampling); + var plugin = new InvocationOtelPlugin( + SdkTracerProvider.builder() + .setSampler(configuredSampler) + .addSpanProcessor(SimpleSpanProcessor.create(exporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> extracted) + .enableMdc(false) + .build()); + + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); + return exporter; + } + + /** A sampler that counts how many times {@code shouldSample} is invoked, delegating the decision. */ + private static final class CountingSampler implements Sampler { + private final Sampler delegate; + private final AtomicInteger count = new AtomicInteger(); + + CountingSampler(Sampler delegate) { + this.delegate = delegate; + } + + int count() { + return count.get(); + } + + @Override + public SamplingResult shouldSample( + Context parentContext, + String traceId, + String name, + SpanKind spanKind, + Attributes attributes, + List parentLinks) { + count.incrementAndGet(); + return delegate.shouldSample(parentContext, traceId, name, spanKind, attributes, parentLinks); + } + + @Override + public String getDescription() { + return "CountingSampler{" + delegate.getDescription() + "}"; + } + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java new file mode 100644 index 000000000..121f4c8fb --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/DurableSamplingDecisionClassLoaderTest.java @@ -0,0 +1,119 @@ +// 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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertNull; + +import io.opentelemetry.context.Context; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import io.opentelemetry.sdk.trace.samplers.SamplingResult; +import java.net.URL; +import java.net.URLClassLoader; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; + +/** + * Verifies the durable sampling decision crosses the application/Java-agent class-loader boundary. + * + *

Under the documented ADOT setup the plugin JAR is loaded twice: the application class loader computes and stores + * the decision, and a separate Java-agent extension class loader installs and runs the sampler. Because a + * {@link io.opentelemetry.context.ContextKey} uses reference identity, the two loaders hold distinct keys and the + * context carrier alone cannot bridge them. This test reproduces that topology with two child-first class loaders that + * each load {@code DurableSamplingDecision} separately while sharing the OpenTelemetry API/SDK types with the parent, + * then asserts the thread-scoped system-property bridge carries the decision from one loader to the other. + */ +class DurableSamplingDecisionClassLoaderTest { + + @AfterEach + void clearBridge() { + DurableSamplingDecision.clearSharedStateForTest(); + } + + @Test + void decisionCrossesClassLoaderBoundary_viaScopedProperty() throws Exception { + try (var appLoader = pluginClassLoader(); + var agentLoader = pluginClassLoader()) { + + var appDecision = Class.forName(DurableSamplingDecision.class.getName(), true, appLoader); + var agentDecision = Class.forName(DurableSamplingDecision.class.getName(), true, agentLoader); + + // The two loaders really did load distinct copies of the class. + assertNotSame(appDecision, agentDecision, "Each class loader must load its own DurableSamplingDecision"); + + // Build a resolved Intent from the application-side loader's own Intent type. + var appIntentClass = Class.forName(DurableSamplingDecision.class.getName() + "$Intent", true, appLoader); + var resolvedFactory = appIntentClass.getDeclaredMethod("resolved", SamplingResult.class); + resolvedFactory.setAccessible(true); + var appIntent = resolvedFactory.invoke(null, SamplingResult.drop()); + + var openScope = appDecision.getDeclaredMethod("openScope", appIntentClass); + openScope.setAccessible(true); + var get = agentDecision.getDeclaredMethod("get", Context.class); + get.setAccessible(true); + + // The application-side loader publishes the intent on this thread; the agent-side loader reads it back from + // a ROOT context (its context key would be a different instance and would miss), reconstructing its own + // Intent from the bridged value. + var scope = (AutoCloseable) openScope.invoke(null, appIntent); + try { + var crossLoaderIntent = get.invoke(null, Context.root()); + assertNotNull( + crossLoaderIntent, "The agent-side loader must read the intent published by the app side"); + // Its Intent type is the agent loader's copy; read the resolved SamplingResult reflectively. + var resolvedAccessor = crossLoaderIntent.getClass().getMethod("resolved"); + resolvedAccessor.setAccessible(true); + var resolved = (SamplingResult) resolvedAccessor.invoke(crossLoaderIntent); + assertEquals( + SamplingDecision.DROP, + resolved.getDecision(), + "The agent-side loader must read the decision published by the application-side loader"); + } finally { + scope.close(); + } + + // After the scope closes, the bridge is cleared and the agent-side read returns null (delegate applies). + assertNull(get.invoke(null, Context.root()), "Closing the scope clears the cross-loader decision"); + } + } + + /** + * A child-first class loader that loads {@code software.amazon.lambda.durable.otel.*} itself (so each instance + * holds its own copies, mirroring the two plugin class loaders) while delegating OpenTelemetry and JDK classes to + * the parent so those types are shared and interoperable across loaders. + */ + private static URLClassLoader pluginClassLoader() { + var classesDir = DurableSamplingDecisionClassLoaderTest.class + .getProtectionDomain() + .getCodeSource() + .getLocation(); + // target/test-classes -> the main classes live in target/classes alongside it. + URL mainClasses; + try { + mainClasses = new URL(classesDir.toString().replace("/test-classes/", "/classes/")); + } catch (Exception e) { + throw new IllegalStateException(e); + } + var parent = DurableSamplingDecisionClassLoaderTest.class.getClassLoader(); + return new URLClassLoader(new URL[] {mainClasses}, parent) { + @Override + protected Class loadClass(String name, boolean resolve) throws ClassNotFoundException { + if (name.startsWith("software.amazon.lambda.durable.otel.")) { + synchronized (getClassLoadingLock(name)) { + var loaded = findLoadedClass(name); + if (loaded == null) { + loaded = findClass(name); + } + if (resolve) { + resolveClass(loaded); + } + return loaded; + } + } + return super.loadClass(name, resolve); + } + }; + } +} 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..129a1b045 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionOtelPluginIntegrationTest.java @@ -0,0 +1,213 @@ +// 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(); + DurableSamplingDecision.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + 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(); + DurableSamplingDecision.clearSharedStateForTest(); + OtelPluginAutoConfigurationState.resetInstalledForTest(); + } + + @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(); + 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..1f4adcf1c 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 @@ -18,8 +18,10 @@ import io.opentelemetry.sdk.testing.exporter.InMemorySpanExporter; import io.opentelemetry.sdk.trace.SdkTracerProvider; import io.opentelemetry.sdk.trace.export.SimpleSpanProcessor; +import io.opentelemetry.sdk.trace.samplers.Sampler; import java.time.Instant; import java.util.ServiceLoader; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -38,6 +40,7 @@ class ExecutionOtelPluginTest { @BeforeEach void setUp() { DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); OtelPluginAutoConfigurationState.resetInstalledForTest(); spanExporter = InMemorySpanExporter.create(); var resource = Resource.create(Attributes.of(SERVICE_NAME, CONFIGURED_SERVICE_NAME)); @@ -56,6 +59,7 @@ void setUp() { void tearDown() { GlobalOpenTelemetry.resetForTest(); DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); OtelPluginAutoConfigurationState.resetInstalledForTest(); } @@ -247,7 +251,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 +259,139 @@ 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(); - try (var ignored = Span.wrap(parentSpanContext).makeCurrent()) { + 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)); + + 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_staysOnExecutionTrace_withoutLinkingAmbientSpan() { + // With no backend execution context (null extractor) but a valid ambient span on a different trace (for + // example a per-invocation Lambda/agent span or a custom-propagated parent), the durable spans must stay on the + // stable ARN-derived execution trace and must NOT link the ambient span: the conformance contract requires the + // Invocation span to have no links, and the ambient span on a foreign trace is not modeled as a link. + var ambientTraceId = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; + var ambientSpanId = "1111111111111111"; + var ambient = + SpanContext.create(ambientTraceId, ambientSpanId, TraceFlags.getSampled(), TraceState.getDefault()); + try (var ignored = Span.wrap(ambient).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"); + assertNotEquals(ambientTraceId, workflowSpan.getTraceId(), "Workflow stays on the execution trace"); + assertEquals(workflowSpan.getTraceId(), invocationSpan.getTraceId(), "Invocation shares the execution trace"); + assertTrue(invocationSpan.getLinks().isEmpty(), "Invocation span carries no ambient link"); + } + + @Test + void contextExtractor_isInvokedEveryInvocation_evenWithAmbientSpan_andBackendContextWins() { + // Contract: the extractor is consulted on every invocation, unconditionally — including when a valid ambient + // span is active — and a valid extracted backend context anchors the execution trace over the ambient span. + var backendTraceId = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; + var backendParentId = "2222222222222222"; + var extractCalls = new AtomicInteger(); + var exporter = InMemorySpanExporter.create(); + var extractorPlugin = new ExecutionOtelPlugin( + SdkTracerProvider.builder().addSpanProcessor(SimpleSpanProcessor.create(exporter)), + OtelPluginConfig.builder() + .contextExtractor(() -> { + extractCalls.incrementAndGet(); + return new ExtractedContext( + backendTraceId, backendParentId, ExtractedContext.Sampling.SAMPLED); + }) + .enableMdc(false) + .workflowSpanName("Workflow") + .build()); + + var ambient = SpanContext.create( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "1111111111111111", + TraceFlags.getSampled(), + TraceState.getDefault()); + try (var ignored = Span.wrap(ambient).makeCurrent()) { + extractorPlugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, Instant.now())); + } + extractorPlugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.SUCCEEDED, null)); + + assertEquals(1, extractCalls.get(), "Extractor is invoked even when a valid ambient span is active"); + var spans = exporter.getFinishedSpanItems(); + var workflowSpan = spanByName(spans, "Workflow"); + assertEquals(backendTraceId, workflowSpan.getTraceId(), "Extracted backend context anchors the trace"); + assertEquals(backendParentId, workflowSpan.getParentSpanId(), "Workflow parents onto the backend span"); + } + + @Test + void executionTrace_isStableAcrossReinvocations_withDifferentAmbientTraces() { + // Reinvocation regression: the same durable execution keeps one trace ID across invocations even when the + // ambient span differs on each invocation (as a per-invocation Lambda/agent span would). + var ambientA = SpanContext.create( + "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", + "1111111111111111", + TraceFlags.getSampled(), + TraceState.getDefault()); + var ambientB = SpanContext.create( + "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb", + "2222222222222222", + TraceFlags.getSampled(), + TraceState.getDefault()); + var startTime = Instant.now(); + + try (var ignored = Span.wrap(ambientA).makeCurrent()) { + plugin.onInvocationStart(new InvocationInfo("req-1", ARN, true, startTime)); + } + plugin.onInvocationEnd(new InvocationEndInfo("req-1", ARN, true, InvocationStatus.PENDING, null)); + var firstInvocationTrace = + spanByName(spanExporter.getFinishedSpanItems(), "Invocation").getTraceId(); + spanExporter.reset(); + + try (var ignored = Span.wrap(ambientB).makeCurrent()) { + plugin.onInvocationStart(new InvocationInfo("req-2", ARN, false, startTime)); + } + plugin.onInvocationEnd(new InvocationEndInfo("req-2", ARN, false, InvocationStatus.SUCCEEDED, null)); + var secondInvocationTrace = + spanByName(spanExporter.getFinishedSpanItems(), "Invocation").getTraceId(); + + assertEquals( + firstInvocationTrace, + secondInvocationTrace, + "The execution trace is stable across reinvocations despite different ambient traces"); } @Test @@ -819,7 +934,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 +954,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 +964,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 +1049,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 +1088,81 @@ 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_undecidedSampling_parentBasedSampler_defersToSamplerAndExports() { + // With a ParentBased(root=alwaysOn) sampler and no explicit upstream Sampled, the undecided decision is + // resolved from the configured sampler (sampled here) rather than forced unsampled. The remote parent is + // therefore built sampled, so the execution trace is exported instead of dropped. + 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)); + + var spans = exporter.getFinishedSpanItems(); + assertFalse( + spans.isEmpty(), + "Undecided upstream defers to the configured sampler (alwaysOn), so spans are exported"); + var workflowSpan = spanByName(spans, "Workflow"); + assertEquals(xrayTraceId, workflowSpan.getTraceId()); + assertEquals(parentSpanId, workflowSpan.getParentSpanId(), "Workflow parents onto the remote span"); + assertTrue(workflowSpan.getSpanContext().isSampled(), "Resolved from the sampler, the trace is sampled"); } @Test - void xrayExtraction_withParentSpanId_invocationSpanHasCorrectParent() { + void xrayExtraction_undecidedSampling_parentBasedNeverSampler_dropsExecutionTrace() { + // Symmetric case: when the configured sampler's root decision is "drop" (ParentBased(root=alwaysOff)) and the + // upstream is undecided, the resolved decision is not-sampled, so nothing is exported. This confirms the + // undecided path follows the sampler in both directions rather than being hardcoded. + var xrayTraceId = "aabbccddee112233445566778899aabb"; + var parentSpanId = "53995c3f42cd8ad8"; + var exporter = InMemorySpanExporter.create(); + var xrayPlugin = new ExecutionOtelPlugin( + SdkTracerProvider.builder() + .setSampler(Sampler.parentBased(Sampler.alwaysOff())) + .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(), + "An undecided upstream with a drop-sampler resolves to not-sampled"); + } + + @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 +1173,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..daa94ae85 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExecutionTraceContextTest.java @@ -0,0 +1,216 @@ +// 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.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +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 String ALL_ZERO_TRACE_ID = "00000000000000000000000000000000"; + private static final String ALL_ZERO_SPAN_ID = "0000000000000000"; + 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, canonicalTraceId(extracted)); + } + + @Test + void canonicalTraceId_derivesFromArn_whenNoRemoteTraceId() { + var canonical = canonicalTraceId(null); + assertEquals(idGenerator.generateTraceIdForExecution(ARN, START), canonical); + assertTrue(canonical.matches("[0-9a-f]{32}")); + } + + @Test + void canonicalTraceId_derivesFromArn_whenRemoteTraceIdIsAllZero() { + // An all-zero (invalid) Root must not become canonical; it falls back to the ARN-derived trace ID. + var extracted = new ExtractedContext(ALL_ZERO_TRACE_ID, REMOTE_PARENT_ID); + assertEquals(idGenerator.generateTraceIdForExecution(ARN, START), canonicalTraceId(extracted)); + } + + @Test + void allZeroRoot_synthesizesRootOnDerivedTrace_notOnTheZeroTrace() { + // A valid parent cannot be built from an all-zero trace ID, so the execution anchors on a synthetic root over + // the ARN-derived canonical trace rather than an invalid context. + var extracted = new ExtractedContext(ALL_ZERO_TRACE_ID, REMOTE_PARENT_ID); + var canonical = canonicalTraceId(extracted); + var execCtx = resolve(extracted, canonical, () -> true); + + assertTrue(execCtx.executionAncestor().isValid(), "The ancestor must be a valid context"); + assertEquals( + idGenerator.generateTraceIdForExecution(ARN, START), + execCtx.executionAncestor().getTraceId()); + assertEquals( + idGenerator.generateExecutionRootSpanId(ARN), + execCtx.executionAncestor().getSpanId()); + } + + @Test + void allZeroParent_isTreatedAsAbsent_synthesizesRootOnRemoteTrace() { + // A valid Root with an all-zero (invalid) Parent drops to the synthetic-root path — the invalid parent is not + // used — while still reusing the valid remote trace ID. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, ALL_ZERO_SPAN_ID); + assertFalse(extracted.hasCompleteRemoteParent(), "An all-zero parent is not a usable remote parent"); + + var execCtx = resolve(extracted, REMOTE_TRACE_ID, () -> true); + assertEquals(REMOTE_TRACE_ID, execCtx.executionAncestor().getTraceId()); + assertEquals( + idGenerator.generateExecutionRootSpanId(ARN), + execCtx.executionAncestor().getSpanId()); + } + + @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_supplierDecidesSampling() { + // Row: valid Root, Parent, no valid Sampled -> reuse Root, remote parent, and resolve the undecided decision + // from the configured sampler (via the supplier). Trace flags have no "unset" state, so a remote parent built + // unsampled would make a parent-based sampler drop every child span; the supplier's decision applies instead. + var extracted = new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID); + + var sampledCtx = resolve(extracted, REMOTE_TRACE_ID, () -> true); + assertEquals(REMOTE_TRACE_ID, sampledCtx.executionAncestor().getTraceId()); + assertEquals(REMOTE_PARENT_ID, sampledCtx.executionAncestor().getSpanId()); + assertTrue(sampledCtx.executionAncestor().isRemote(), "The remote server span is the ancestor"); + assertTrue(sampledCtx.traceFlags().isSampled(), "Undecided upstream defers to the sampler (sampled)"); + + var droppedCtx = resolve(extracted, REMOTE_TRACE_ID, () -> false); + assertEquals(REMOTE_PARENT_ID, droppedCtx.executionAncestor().getSpanId()); + assertFalse(droppedCtx.traceFlags().isSampled(), "Undecided upstream defers to the sampler (not sampled)"); + } + + @Test + void completeRemoteParentExplicitSampled_winsOverSupplier() { + // An explicit upstream Sampled=1/0 is preserved on the remote-parent path regardless of what the supplier says. + var sampledUpstream = + new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID, ExtractedContext.Sampling.SAMPLED); + assertTrue(resolve(sampledUpstream, REMOTE_TRACE_ID, () -> false) + .traceFlags() + .isSampled()); + + var notSampledUpstream = + new ExtractedContext(REMOTE_TRACE_ID, REMOTE_PARENT_ID, ExtractedContext.Sampling.NOT_SAMPLED); + assertFalse(resolve(notSampledUpstream, REMOTE_TRACE_ID, () -> true) + .traceFlags() + .isSampled()); + } + + @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 = canonicalTraceId(null); + 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()); + } + + @Test + void canonicalTraceId_reusesAmbientTraceId_whenNoRemoteContext() { + // With no backend context the canonical trace is ARN/start-time-derived and stable, NOT the ambient span's + // trace: the ambient span is per-invocation and must not destabilize the execution trace across reinvocations. + assertEquals(idGenerator.generateTraceIdForExecution(ARN, START), canonicalTraceId(null)); + } + + @Test + void ancestorIsSyntheticRoot_whenNoBackendContext_regardlessOfAmbient() { + // Ambient is not an ancestor. With no backend context the ancestor is the deterministic synthetic root on the + // ARN-derived trace, so the execution stays on one stable trace even when a per-invocation ambient span exists. + var canonical = canonicalTraceId(null); + var execCtx = resolve(null, canonical, () -> true); + + assertEquals( + idGenerator.generateTraceIdForExecution(ARN, START), + execCtx.executionAncestor().getTraceId()); + assertEquals( + idGenerator.generateExecutionRootSpanId(ARN), + execCtx.executionAncestor().getSpanId()); + } + + @Test + void canonicalTraceId_isStableAcrossReinvocations_withDifferentAmbientTraces() { + // The same execution (ARN + start time) yields the same canonical trace regardless of the ambient span, which + // may be a different per-invocation trace on each reinvocation. This is the core stability invariant. + var expected = idGenerator.generateTraceIdForExecution(ARN, START); + assertEquals(expected, canonicalTraceId(null), "reinvocation 1 (ambient trace A) -> same execution trace"); + assertEquals(expected, canonicalTraceId(null), "reinvocation 2 (ambient trace B) -> same execution trace"); + } + + private static SpanContext validAmbientSpan(String traceId) { + return SpanContext.create(traceId, "1111111111111111", TraceFlags.getSampled(), TraceState.getDefault()); + } + + private String canonicalTraceId(ExtractedContext extracted) { + return ExecutionTraceContext.canonicalTraceId(extracted, ARN, START, idGenerator); + } + + 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/ExtractedContextTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExtractedContextTest.java new file mode 100644 index 000000000..e5efb5f81 --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/ExtractedContextTest.java @@ -0,0 +1,63 @@ +// 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.assertEquals; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.jupiter.api.Test; + +/** Tests {@link ExtractedContext} construction, focusing on null-sampling normalization and legacy deserialization. */ +class ExtractedContextTest { + + private static final String TRACE_ID = "aabbccddee112233445566778899aabb"; + private static final String PARENT_SPAN_ID = "53995c3f42cd8ad8"; + + private final ObjectMapper mapper = new ObjectMapper(); + + @Test + void canonicalConstructor_normalizesNullSamplingToUndecided() { + // The canonical constructor is public; a null sampling must be normalized so downstream resolution, which + // switches on the decision, does not throw and silently disable telemetry. + var context = new ExtractedContext(TRACE_ID, PARENT_SPAN_ID, null); + assertEquals(ExtractedContext.Sampling.UNDECIDED, context.sampling()); + } + + @Test + void twoArgConstructor_defaultsToUndecided() { + var context = new ExtractedContext(TRACE_ID, PARENT_SPAN_ID); + assertEquals(ExtractedContext.Sampling.UNDECIDED, context.sampling()); + } + + @Test + void deserialize_legacyValueWithoutSamplingField_isUndecided() throws Exception { + // A value serialized before the sampling component existed omits the field; Jackson leaves it null. The + // compact constructor must normalize it to UNDECIDED rather than producing a context that throws on use. + var legacyJson = "{\"traceId\":\"" + TRACE_ID + "\",\"parentSpanId\":\"" + PARENT_SPAN_ID + "\"}"; + + var context = mapper.readValue(legacyJson, ExtractedContext.class); + + assertEquals(TRACE_ID, context.traceId()); + assertEquals(PARENT_SPAN_ID, context.parentSpanId()); + assertEquals(ExtractedContext.Sampling.UNDECIDED, context.sampling()); + } + + @Test + void deserialize_explicitNullSampling_isUndecided() throws Exception { + var json = "{\"traceId\":\"" + TRACE_ID + "\",\"parentSpanId\":null,\"sampling\":null}"; + + var context = mapper.readValue(json, ExtractedContext.class); + + assertEquals(ExtractedContext.Sampling.UNDECIDED, context.sampling()); + } + + @Test + void serialize_thenDeserialize_roundTripsSampling() throws Exception { + var original = new ExtractedContext(TRACE_ID, PARENT_SPAN_ID, ExtractedContext.Sampling.SAMPLED); + + var roundTripped = mapper.readValue(mapper.writeValueAsString(original), ExtractedContext.class); + + assertEquals(original, roundTripped); + assertEquals(ExtractedContext.Sampling.SAMPLED, roundTripped.sampling()); + } +} 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..f9ff81d20 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 @@ -41,6 +41,7 @@ class InvocationOtelPluginIntegrationTest { @BeforeEach void setUp() { DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); OtelPluginAutoConfigurationState.resetInstalledForTest(); spanExporter = InMemorySpanExporter.create(); @@ -58,6 +59,7 @@ void setUp() { void tearDown() { GlobalOpenTelemetry.resetForTest(); DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); OtelPluginAutoConfigurationState.resetInstalledForTest(); } @@ -90,10 +92,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..d46073baa 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 @@ -47,6 +47,7 @@ class InvocationOtelPluginTest { @BeforeEach void setUp() { DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); OtelPluginAutoConfigurationState.resetInstalledForTest(); spanExporter = InMemorySpanExporter.create(); @@ -62,6 +63,7 @@ void setUp() { void tearDown() { GlobalOpenTelemetry.resetForTest(); DeterministicIdGenerator.clearSharedStateForTest(); + DurableSamplingDecision.clearSharedStateForTest(); OtelPluginAutoConfigurationState.resetInstalledForTest(); } @@ -293,23 +295,26 @@ 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_staysOnExecutionTrace_withoutLinkingAmbientSpan() { + // The extractor returns null (no backend execution context), but a valid ambient span is active on a different + // trace (for example a per-invocation Lambda/agent span or a custom-propagated parent). The durable spans must + // stay on the stable ARN-derived execution trace and must NOT link the ambient span: the conformance contract + // requires the Invocation span to have no links, and the ambient span on a foreign trace is not a link. + 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"); + assertNotEquals(ambientTraceId, invocationSpan.getTraceId(), "Invocation stays on the execution trace"); + assertEquals(workflowSpan.getTraceId(), invocationSpan.getTraceId(), "Both share the execution trace"); + assertTrue(invocationSpan.getLinks().isEmpty(), "Invocation span carries no ambient link"); } @Test @@ -939,32 +944,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 +1112,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 +1170,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 +1204,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 +1231,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 +1296,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 +1323,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 +1341,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 +1356,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 +1583,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 +1638,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 +1696,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 +1707,27 @@ 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"); + // Link ORDER matters for conformance: the initial operation link comes first, then the Workflow link. + assertEquals(2, waitContinuation.getLinks().size(), "Continuation span has exactly [operation, Workflow]"); + assertEquals( + originalWaitSpanId, + waitContinuation.getLinks().get(0).getSpanContext().getSpanId(), + "First link is the initial operation span"); + assertEquals( + workflowSpan.getSpanId(), + waitContinuation.getLinks().get(1).getSpanContext().getSpanId(), + "Second link is the Workflow span"); } // ─── Cross-invocation step retry scenario ──────────────────────────── @@ -1691,8 +1736,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 +1782,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 +1856,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 ─────────────────────────────────────────── @@ -1897,8 +1952,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/OtelPluginSupportTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/OtelPluginSupportTest.java new file mode 100644 index 000000000..91ae4ceab --- /dev/null +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/OtelPluginSupportTest.java @@ -0,0 +1,156 @@ +// 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.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.opentelemetry.api.common.Attributes; +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.api.trace.TraceFlags; +import io.opentelemetry.api.trace.TraceState; +import io.opentelemetry.sdk.trace.samplers.SamplingDecision; +import org.junit.jupiter.api.Test; + +/** + * Tests the sampling decision resolved once per invocation: explicit-upstream and same-trace-ambient precedence, and + * the Java-agent path always deferring to the installed sampler (returning {@code null}). + */ +class OtelPluginSupportTest { + + private static final String TRACE_ID = "aabbccddee112233445566778899aabb"; + private static final String OTHER_TRACE_ID = "11223344556677889900aabbccddeeff"; + private static final String SPAN_ID = "1111111111111111"; + + @Test + void agentPath_defersToInstalledSampler_regardlessOfConfiguredSetting() { + // On the agent path (provider not visible) the resolver never reconstructs a sampler from configuration — the + // agent's effective sampler may have been wrapped/replaced by another extension. It always returns null so the + // agent-installed DurableSampler consults its real delegate. This holds regardless of OTEL_TRACES_SAMPLER. + assertNull(resolve(null), "No configuration: defer to the installed sampler"); + } + + @Test + void explicitSampled_isAuthoritative() { + var extracted = new ExtractedContext(TRACE_ID, SPAN_ID, ExtractedContext.Sampling.SAMPLED); + assertTrue(sampled(resolve(extracted)), "Explicit upstream Sampled=1 is authoritative"); + } + + @Test + void explicitNotSampled_isAuthoritative() { + var extracted = new ExtractedContext(TRACE_ID, SPAN_ID, ExtractedContext.Sampling.NOT_SAMPLED); + assertFalse(sampled(resolve(extracted)), "Explicit upstream Sampled=0 is authoritative"); + } + + @Test + void sameTraceAmbient_followsAmbientSampledBit_whenNoExplicitDecision() { + // No explicit Sampled, and a valid ambient span on the canonical trace: follow the ambient span's decision. + var ambientSampled = ambient(TRACE_ID, TraceFlags.getSampled(), true); + assertEquals( + SamplingDecision.RECORD_AND_SAMPLE, + resolve(null, ambientSampled).getDecision(), + "sampled -> sample"); + + var ambientNotSampledNotRecording = ambient(TRACE_ID, TraceFlags.getDefault(), false); + assertEquals( + SamplingDecision.DROP, + resolve(null, ambientNotSampledNotRecording).getDecision(), + "same-trace ambient unsampled and not recording -> DROP"); + } + + @Test + void sameTraceAmbient_recordOnly_isPreserved() { + // A RECORD_ONLY ambient span has an unsampled SpanContext but is still recording; reducing it to its sampled + // bit would DROP durable spans. It must map to RECORD_ONLY so those spans still reach processors. + var recordOnlyAmbient = ambient(TRACE_ID, TraceFlags.getDefault(), true); + assertEquals( + SamplingDecision.RECORD_ONLY, + resolve(null, recordOnlyAmbient).getDecision(), + "same-trace ambient unsampled but recording -> RECORD_ONLY"); + } + + @Test + void differentTraceAmbient_isIgnored_defersOnAgentPath() { + // An ambient span on a different trace is not representative of this execution and is ignored; with no explicit + // decision and the provider not visible, the agent path defers to the installed sampler. + var ambientOtherTrace = ambient(OTHER_TRACE_ID, TraceFlags.getSampled(), true); + assertNull(resolve(null, ambientOtherTrace), "different-trace ambient ignored -> defer on the agent path"); + } + + /** A span exposing the given context and recording state, for exercising the ambient-decision branch. */ + private static Span ambient(String traceId, TraceFlags flags, boolean recording) { + var context = SpanContext.create(traceId, SPAN_ID, flags, TraceState.getDefault()); + return new RecordingStateSpan(context, recording); + } + + private static boolean sampled(io.opentelemetry.sdk.trace.samplers.SamplingResult result) { + return OtelPluginSupport.isSampled(result); + } + + private static io.opentelemetry.sdk.trace.samplers.SamplingResult resolve(ExtractedContext extracted) { + return resolve(extracted, Span.getInvalid()); + } + + private static io.opentelemetry.sdk.trace.samplers.SamplingResult resolve( + ExtractedContext extracted, Span ambient) { + // A null sdkTracerProvider exercises the Java-agent path, where the resolver always defers (returns null). + return OtelPluginSupport.resolveSamplingResult( + null, extracted, ambient, TRACE_ID, "Workflow", Attributes.empty()); + } + + /** + * A minimal {@link Span} that only reports a fixed {@link SpanContext} and recording state; all other operations + * are no-ops. Used to drive the RECORD_ONLY-vs-DROP distinction, which {@code Span.wrap} cannot express because a + * propagated (wrapped) span is never recording. + */ + private record RecordingStateSpan(SpanContext spanContext, boolean recording) implements Span { + @Override + public SpanContext getSpanContext() { + return spanContext; + } + + @Override + public boolean isRecording() { + return recording; + } + + @Override + public Span setAttribute(io.opentelemetry.api.common.AttributeKey key, T value) { + return this; + } + + @Override + public Span addEvent(String name, Attributes attributes) { + return this; + } + + @Override + public Span addEvent(String name, Attributes attributes, long timestamp, java.util.concurrent.TimeUnit unit) { + return this; + } + + @Override + public Span setStatus(io.opentelemetry.api.trace.StatusCode statusCode, String description) { + return this; + } + + @Override + public Span recordException(Throwable exception, Attributes additionalAttributes) { + return this; + } + + @Override + public Span updateName(String name) { + return this; + } + + @Override + public void end() {} + + @Override + public void end(long timestamp, java.util.concurrent.TimeUnit unit) {} + } +} diff --git a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/XRayContextExtractorTest.java b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/XRayContextExtractorTest.java index 0f5a196d0..0a02a66e9 100644 --- a/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/XRayContextExtractorTest.java +++ b/otel-plugin/src/test/java/software/amazon/lambda/durable/otel/XRayContextExtractorTest.java @@ -46,6 +46,33 @@ void extract_implementsContextExtractor() { assertNull(extractor.extract()); } + @Test + void extract_returnsNull_whenRootIsAllZero() { + // An all-zero Root is well-formed hex but an invalid trace ID; it must not be surfaced as usable context. + var traceHeader = "Root=1-00000000-000000000000000000000000;Parent=53995c3f42cd8ad8;Sampled=1"; + System.setProperty("com.amazonaws.xray.traceHeader", traceHeader); + try { + assertNull(new XRayContextExtractor().extract(), "An all-zero Root must yield no extracted context"); + } finally { + System.clearProperty("com.amazonaws.xray.traceHeader"); + } + } + + @Test + void extract_dropsAllZeroParent_keepsValidRoot() { + // A valid Root with an all-zero Parent keeps the trace ID but treats the parent as absent. + var traceHeader = "Root=1-6a43574f-2cf3140a69b0c8fe165f9503;Parent=0000000000000000;Sampled=1"; + System.setProperty("com.amazonaws.xray.traceHeader", traceHeader); + try { + var context = new XRayContextExtractor().extract(); + assertNotNull(context); + assertEquals("6a43574f2cf3140a69b0c8fe165f9503", context.traceId()); + assertNull(context.parentSpanId(), "An all-zero Parent must be treated as absent"); + } finally { + System.clearProperty("com.amazonaws.xray.traceHeader"); + } + } + // ─── xrayRootToOtelTraceId: valid inputs ──────────────────────────── @Test diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/DurableExecutionCheckpointTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/DurableExecutionCheckpointTest.java index 826f2abbe..7bdacbb68 100644 --- a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/DurableExecutionCheckpointTest.java +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/DurableExecutionCheckpointTest.java @@ -11,6 +11,7 @@ import software.amazon.lambda.durable.execution.DurableExecutor; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; import software.amazon.lambda.durable.testing.local.LocalMemoryExecutionClient; /** Integration tests that verify checkpoint behavior using LocalMemoryExecutionClient */ @@ -105,4 +106,24 @@ void testSmallPayloadNoExtraCheckpoint() { .toList(); assertTrue(executionUpdates.isEmpty()); } + + @Test + void testLargeResultCanReplayWithoutDuplicateExecutionOperation() { + // A >6MB result checkpoints a SUCCEEDED EXECUTION operation into storage. Because the EXECUTION operation ID is + // stable across reinvocations, a second run must not collide the stored EXECUTION operation with the fresh one + // when ExecutionManager builds its ID-keyed operation map. Regression for the large-result replay path. + var largeString = "x".repeat(7 * 1024 * 1024); // 7MB string, exceeds the Lambda response limit + var runner = LocalDurableTestRunner.create(String.class, (input, ctx) -> largeString); + + // First run checkpoints the large EXECUTION result and stores the EXECUTION operation. + var first = runner.runUntilComplete("input"); + assertEquals(ExecutionStatus.SUCCEEDED, first.getStatus()); + assertEquals(largeString, first.getResult(String.class)); + + // Second run replays over the stored state; it must not throw a duplicate-key error while collecting + // operations, and the completed large result must still be recoverable on replay. + var second = runner.runUntilComplete("input"); + assertEquals(ExecutionStatus.SUCCEEDED, second.getStatus()); + assertEquals(largeString, second.getResult(String.class)); + } } 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..06d59d5d3 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,24 +334,56 @@ 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) + + // The list must contain exactly one EXECUTION operation, matching the backend, which keeps a single EXECUTION + // operation and updates it in place. Its ID is stable across reinvocations, so a stored EXECUTION operation + // (for example the SUCCEEDED one written when a >6MB result is checkpointed) shares the fresh operation's ID; + // including both would collide when ExecutionManager builds its ID-keyed operation map. Merge them into one: + // start from the fresh operation, which carries the input details (EXECUTION result payloads are not stored on + // the operation — they live on the event stream — so the fresh op is the only source of the input), and adopt + // the stored operation's terminal status and end timestamp when present, so replay of an already-completed + // execution still looks completed. + var storedExecutionOps = storage.getAllOperations().stream() + .filter(op -> op.type() == OperationType.EXECUTION) + .toList(); + if (storedExecutionOps.size() > 1) { + // The harness owns the single-EXECUTION-operation invariant; surface a violation loudly rather than hiding + // it by silently merging, since it would indicate a storage bug. + throw new IllegalStateException( + "Expected at most one stored EXECUTION operation but found " + storedExecutionOps.size()); + } + var storedExecutionOp = storedExecutionOps.isEmpty() ? null : storedExecutionOps.get(0); + var executionOpBuilder = Operation.builder() + .id(executionOperationId) .name(executionName) .type(OperationType.EXECUTION) .status(OperationStatus.STARTED) .startTimestamp(executionStartTime) .executionDetails( - ExecutionDetails.builder().inputPayload(inputJson).build()) - .build(); + ExecutionDetails.builder().inputPayload(inputJson).build()); + if (storedExecutionOp != null) { + // Preserve the persisted lifecycle state (e.g. SUCCEEDED when a large result was checkpointed) while + // keeping + // the input details from the fresh operation. Only status and end timestamp are adopted; start timestamp + // stays the (stable) fresh value. + executionOpBuilder.status(storedExecutionOp.status()); + if (storedExecutionOp.endTimestamp() != null) { + executionOpBuilder.endTimestamp(storedExecutionOp.endTimestamp()); + } + } + var executionOp = executionOpBuilder.build(); - // Load previous operations and include them in InitialExecutionState - var existingOps = storage.getAllOperations(); + // Load previous non-EXECUTION operations; the single merged EXECUTION operation above already represents the + // execution's state. + var existingOps = storage.getAllOperations().stream() + .filter(op -> op.type() != OperationType.EXECUTION) + .toList(); var allOps = new ArrayList<>(List.of(executionOp)); allOps.addAll(existingOps);