From e6f3f01211a81dd7a993421188ee7ddacd081225 Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Wed, 26 Aug 2026 21:23:54 -0700 Subject: [PATCH] feat(otel): parent durable spans to shared execution trace --- ...test_opentelemetry_conformance_workflow.py | 7 +- .../README.md | 107 ++++- .../__init__.py | 6 + .../context_extractors.py | 134 +++++- .../deterministic_id_generator.py | 22 +- .../durable_sampling.py | 135 ++++++ .../execution_plugin.py | 133 ++++-- .../execution_trace_context.py | 98 ++++ .../invocation_plugin.py | 222 +++++++-- .../e2e/test_invocation_wait_resume_int.py | 252 ++++++++++ .../tests/test_context_extractors.py | 102 ++-- .../tests/test_execution_plugin.py | 128 ++++- .../test_execution_plugin_integration.py | 49 +- .../tests/test_invocation_plugin.py | 443 +++++++++++++++++- .../test_invocation_plugin_integration.py | 35 +- .../tests/test_log_filter.py | 3 +- 16 files changed, 1627 insertions(+), 249 deletions(-) create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_trace_context.py create mode 100644 packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py diff --git a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py index 7b1f4535..b349892d 100644 --- a/.github/scripts/tests/test_opentelemetry_conformance_workflow.py +++ b/.github/scripts/tests/test_opentelemetry_conformance_workflow.py @@ -4,9 +4,7 @@ WORKFLOW_PATH = ( Path(__file__).parents[2] / "workflows" / "opentelemetry-conformance-tests.yml" ) -EXAMPLES_DIR = ( - ".build/durable-sdk/packages/aws-durable-execution-sdk-python-conformance-tests-otel" -) +EXAMPLES_DIR = ".build/durable-sdk/packages/aws-durable-execution-sdk-python-conformance-tests-otel" def test_opentelemetry_conformance_caller_uses_current_workflow_contract() -> None: @@ -66,8 +64,7 @@ def test_opentelemetry_conformance_runs_when_the_handlers_change() -> None: workflow = WORKFLOW_PATH.read_text() trigger_path = ( - " - " - '"packages/aws-durable-execution-sdk-python-conformance-tests-otel/**"' + ' - "packages/aws-durable-execution-sdk-python-conformance-tests-otel/**"' ) # Once for pull_request, once for push. assert workflow.count(trigger_path) == 2 diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 4adcf32b..038bb05e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -1,16 +1,16 @@ # AWS Durable Execution SDK - OpenTelemetry Plugin -OpenTelemetry instrumentation plugin for the [AWS Durable Execution SDK for Python](https://github.com/aws/aws-durable-execution-sdk-python). Emits durable execution spans with deterministic workflow and operation IDs while keeping invocation spans in the ambient Lambda trace. +OpenTelemetry instrumentation plugin for the [AWS Durable Execution SDK for Python](https://github.com/aws/aws-durable-execution-sdk-python). Emits durable execution spans on one execution trace, with deterministic workflow, synthetic-root, and operation span IDs. ## Features -- **Deterministic Workflow Traces**: Durable operations use an execution-derived trace that is independent of the ambient Lambda/X-Ray trace -- **Ambient Invocation Traces**: Invocation spans inherit the active Lambda or extracted upstream context +- **Shared Execution Trace**: Workflow and Invocation spans share one trace, anchored to a propagated backend parent when available or a deterministic synthetic execution root otherwise +- **Same-Trace Ambient Parenting**: Invocation spans use the active ambient span only when it already belongs to the execution trace - **Span-per-Operation**: Each durable operation (step, wait, invoke) gets its own span with accurate timing - **Continuation Spans**: Operations completing in another invocation produce a new correlated span without fabricating an unobserved prior span context - **Log Correlation**: Enrich application logs with trace ID and span ID for end-to-end observability - **Provider Integration**: Use the global ADOT provider or supply an explicit SDK `TracerProvider` -- **Provider-Managed Sampling**: Use standard OpenTelemetry or ADOT sampling configuration +- **Execution Sampling**: Resolve sampling once per invocation and apply it consistently to Workflow, Invocation, operation, and attempt spans ## Installation @@ -124,7 +124,7 @@ fn = lambda_.Function( ### 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 derive deterministic trace IDs that remain consistent across all invocations of the same durable execution. +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 anchor the execution trace on the propagated X-Ray `Root`/`Parent` when both are valid, and preserves `Sampled=1` or `Sampled=0` as the backend sampling decision. **AWS Console:** Lambda → Configuration → Monitoring and operations tools → Active tracing → Enable @@ -191,7 +191,7 @@ The function's execution role needs the `AWSXRayDaemonWriteAccess` managed polic | `OTEL_TRACES_SAMPLER` | Sampler to use (e.g., `traceidratio` for ratio-based sampling) | `always_on` | | `OTEL_TRACES_SAMPLER_ARG` | Argument for the sampler (e.g., `0.3` to sample 30% of traces) | — | -See the [ADOT sampling configuration](https://aws-otel.github.io/docs/getting-started/lambda#sampling-configuration) for more details. +See the [ADOT sampling configuration](https://aws-otel.github.io/docs/getting-started/lambda#sampling-configuration) for more details. When the backend header contains an explicit `Sampled` value, that backend decision takes precedence over local sampler configuration for durable spans. ## Configuration @@ -220,7 +220,15 @@ plugin = InvocationOtelPlugin( ### Context Extractors -The plugin supports multiple strategies for extracting upstream trace context: +Context extractors return an `ExtractedContext` object, or `None` when no +durable execution trace context is available. The object carries: + +- `trace_id`: 128-bit OpenTelemetry trace ID +- `parent_span_id`: 64-bit OpenTelemetry parent span ID +- `sampling`: `Sampling.SAMPLED`, `Sampling.NOT_SAMPLED`, or `Sampling.UNDECIDED` + +The plugin supports multiple strategies for extracting durable execution trace +context: ```python from aws_durable_execution_sdk_python_otel import ( @@ -230,13 +238,65 @@ from aws_durable_execution_sdk_python_otel import ( xray_context_extractor, ) -# Default: X-Ray trace header (recommended for most Lambda deployments) +# Default: X-Ray trace header (recommended for most Lambda deployments). InvocationOtelPlugin(OtelPluginConfig(context_extractor=xray_context_extractor)) -# W3C Trace Context via clientContext (requires backend propagation support) +# W3C Trace Context via clientContext (placeholder for backend propagation support). InvocationOtelPlugin(OtelPluginConfig(context_extractor=w3c_client_context_extractor)) ``` +Custom extractors should return `ExtractedContext`, not an OpenTelemetry +`Context`. + +### Trace Structure + +Both bundled plugins use the same execution ancestor: + +- a propagated backend parent when `_X_AMZN_TRACE_ID` contains a valid `Root` + and `Parent` +- otherwise a deterministic, non-recording synthetic root derived from the + durable execution ARN + +`InvocationOtelPlugin` keeps durable operation spans under the Invocation span +and links operations to Workflow: + +```text +Execution ancestor +├── Workflow +└── Invocation + └── operation + └── operation attempt 1 +``` + +`ExecutionOtelPlugin` keeps operation spans under Workflow and links operations +to the current Invocation span: + +```text +Execution ancestor +├── Workflow +│ └── operation +│ └── operation attempt 1 +└── Invocation +``` + +If an ambient Lambda span is active and already has the execution trace ID, the +Invocation span uses that ambient span as its parent. Ambient spans on a +different trace are ignored for durable parenting so Invocation remains on the +execution trace. + +### Sampling + +Sampling is resolved once per invocation and carried to every durable span in +that invocation. Precedence is: + +1. `Sampled=1` or `Sampled=0` from `_X_AMZN_TRACE_ID` +2. a same-trace ambient span's recording/sampled state +3. the configured OpenTelemetry sampler + +The resolved decision is applied to Workflow, Invocation, operation, and attempt +spans. This avoids independently querying stateful or ratio-based samplers for +each durable span in the same invocation. + ### Log Correlation When `enrich_logger=True` (the default), the plugin installs a logging filter on @@ -256,8 +316,9 @@ 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 the CloudWatch console** — Navigate to CloudWatch → Traces in the AWS Console. You should see a trace with: - - An "invocation" span per invocation +2. **Check the CloudWatch console** — Navigate to CloudWatch → Traces in the AWS Console. You should see an execution trace with: + - A "Workflow" span exported on the terminal invocation + - An "Invocation" span per invocation - Child spans for each durable operation (named after your step names) - All invocations of the same execution grouped under one trace ID @@ -272,7 +333,7 @@ After deploying your function with the plugin configured: | Symptom | Likely Cause | | --------------------------------- | --------------------------------------------------------------- | | No traces appear | ADOT layer not configured, or `AWS_LAMBDA_EXEC_WRAPPER` not set | -| Traces appear but are fragmented | X-Ray active tracing not enabled on the Lambda function | +| Traces appear but are fragmented | Backend trace context is not propagated to every invocation | | Missing spans for some operations | `OTEL_TRACES_SAMPLER_ARG` set below 1.0 | | `_X_AMZN_TRACE_ID` not populated | X-Ray active tracing not enabled | @@ -280,7 +341,7 @@ After deploying your function with the plugin configured: ### `InvocationOtelPlugin` -The main plugin class. Implements `DurableInstrumentationPlugin` from `aws_durable_execution_sdk_python`. +Invocation-rooted view. Implements `DurableInstrumentationPlugin` from `aws_durable_execution_sdk_python`. ```python InvocationOtelPlugin( @@ -297,21 +358,35 @@ InvocationOtelPlugin( Pass `tracer_provider=...` when the application owns the OpenTelemetry SDK provider. When omitted, the globally configured provider is used. +### `ExecutionOtelPlugin` + +Execution-rooted view. Uses the same execution ancestor and sampling behavior as +`InvocationOtelPlugin`, but parents operation spans under Workflow and links +them to Invocation. + ### `DeterministicIdGenerator` A custom OpenTelemetry `IdGenerator` that produces reproducible trace and span IDs from execution metadata. Exported for advanced use cases. ### `xray_context_extractor` -Default context extractor. Reads the `_X_AMZN_TRACE_ID` environment variable to derive trace context. +Default context extractor. Reads the `_X_AMZN_TRACE_ID` environment variable and +returns `ExtractedContext` containing parsed `Root`, `Parent`, and `Sampled` +fields when present. ### `w3c_client_context_extractor` -Alternative context extractor. Reads W3C `traceparent` from `context.clientContext.custom.traceparent`. Requires backend `clientContext` propagation to be enabled. +Alternative context extractor placeholder. Returns `None` until backend W3C +`traceparent` propagation is supported. ### `ContextExtractor` -Type alias for custom context extractor functions. +Type alias for custom context extractor functions: +`Callable[[InvocationStartInfo], ExtractedContext | None]`. + +### `ExtractedContext` / `Sampling` + +Structured trace context and sampling decision returned by context extractors. ### `OtelContextLogFilter` / `install_log_filter` diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py index b6552896..a8285dcb 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/__init__.py @@ -3,11 +3,14 @@ from aws_durable_execution_sdk_python_otel.__about__ import __version__ from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, + ExtractedContext, + Sampling, w3c_client_context_extractor, xray_context_extractor, ) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, + derive_execution_root_span_id, derive_workflow_span_id, operation_id_to_span_id, ) @@ -35,11 +38,14 @@ "ContextExtractor", "DeterministicIdGenerator", "ExecutionOtelPlugin", + "ExtractedContext", "OtelPluginConfig", "InvocationOtelPlugin", "OtelContextLogFilter", + "Sampling", "ProviderResult", "create_tracer_provider", + "derive_execution_root_span_id", "derive_workflow_span_id", "install_log_filter", "operation_id_to_span_id", diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_extractors.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_extractors.py index 79029fe5..6f91aee5 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_extractors.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/context_extractors.py @@ -1,40 +1,132 @@ -"""Context extractors for propagating trace context into durable executions.""" +"""Trace-context extractors for durable execution telemetry.""" from __future__ import annotations import os +from dataclasses import dataclass +from enum import Enum from typing import TYPE_CHECKING, Callable -from opentelemetry import context as otel_context, propagate - if TYPE_CHECKING: - from opentelemetry.context import Context - from aws_durable_execution_sdk_python.plugin import InvocationStartInfo -ContextExtractor = Callable[["InvocationStartInfo"], "Context"] +class Sampling(Enum): + """Sampling decision propagated by the durable execution backend.""" + + SAMPLED = "sampled" + NOT_SAMPLED = "not_sampled" + UNDECIDED = "undecided" + + +@dataclass(frozen=True) +class ExtractedContext: + """Trace context extracted from the durable execution backend. + + Attributes: + trace_id: OTel 128-bit trace ID, or ``None`` when no valid trace ID was + present. + parent_span_id: OTel 64-bit parent span ID, or ``None`` when no valid + parent was present. + sampling: Explicit backend sampling decision, or ``UNDECIDED`` when + the backend header did not include one. + """ + + trace_id: int | None + parent_span_id: int | None + sampling: Sampling = Sampling.UNDECIDED + + @property + def has_valid_trace_id(self) -> bool: + return self.trace_id is not None and 0 < self.trace_id < 2**128 + + @property + def has_valid_parent_span_id(self) -> bool: + return self.parent_span_id is not None and 0 < self.parent_span_id < 2**64 + + @property + def has_complete_remote_parent(self) -> bool: + return self.has_valid_trace_id and self.has_valid_parent_span_id + + +ContextExtractor = Callable[["InvocationStartInfo"], ExtractedContext | None] -def xray_context_extractor(info: "InvocationStartInfo") -> "Context": - """Read the X-Ray trace header from the _X_AMZN_TRACE_ID environment variable. - The durable execution backend propagates the same Root trace ID to every - invocation, so all invocations share one traceId. +def _ensure_extracted_context(extracted: object) -> ExtractedContext | None: + """Validate a context extractor result.""" + if extracted is None or isinstance(extracted, ExtractedContext): + return extracted + msg = "context extractor must return ExtractedContext or None" + raise TypeError(msg) + + +def _parse_xray_trace_id(root: str | None) -> int | None: + if root is None: + return None + parts = root.split("-") + if len(parts) != 3 or parts[0] != "1": + return None + trace_id_hex = f"{parts[1]}{parts[2]}" + if len(trace_id_hex) != 32: + return None + try: + trace_id = int(trace_id_hex, 16) + except ValueError: + return None + return trace_id if 0 < trace_id < 2**128 else None + + +def _parse_span_id(span_id_hex: str | None) -> int | None: + if span_id_hex is None or len(span_id_hex) != 16: + return None + try: + span_id = int(span_id_hex, 16) + except ValueError: + return None + return span_id if 0 < span_id < 2**64 else None + + +def _parse_sampling(value: str | None) -> Sampling: + if value == "1": + return Sampling.SAMPLED + if value == "0": + return Sampling.NOT_SAMPLED + return Sampling.UNDECIDED + + +def xray_context_extractor(info: "InvocationStartInfo") -> ExtractedContext | None: + """Read durable execution trace context from ``_X_AMZN_TRACE_ID``. + + The Lambda durable execution backend propagates an X-Ray style header. A + valid ``Root`` anchors the execution trace; a valid ``Parent`` becomes the + remote execution ancestor; and ``Sampled`` is preserved as the backend's + explicit sampling decision. """ trace_header = os.environ.get("_X_AMZN_TRACE_ID") if not trace_header: - return otel_context.get_current() - return propagate.extract( - carrier={"X-Amzn-Trace-Id": trace_header}, - context=otel_context.get_current(), - ) + return None + parts: dict[str, str] = {} + for segment in trace_header.split(";"): + key, separator, value = segment.partition("=") + if separator: + parts[key.strip()] = value.strip() -def w3c_client_context_extractor(info: "InvocationStartInfo") -> "Context": - """Read W3C traceparent from context.clientContext.custom.traceparent. + trace_id = _parse_xray_trace_id(parts.get("Root")) + parent_span_id = _parse_span_id(parts.get("Parent")) + sampling = _parse_sampling(parts.get("Sampled")) + if trace_id is None and parent_span_id is None and sampling is Sampling.UNDECIDED: + return None + return ExtractedContext( + trace_id=trace_id, + parent_span_id=parent_span_id, + sampling=sampling, + ) - Requires the backend clientContext propagation to be enabled. - This extractor is a placeholder for when backend propagation is supported. - """ - return otel_context.get_current() + +def w3c_client_context_extractor( + info: "InvocationStartInfo", +) -> ExtractedContext | None: + """Placeholder for future W3C traceparent propagation support.""" + return None diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py index 3ad3ed3e..c3d7afc1 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/deterministic_id_generator.py @@ -26,9 +26,9 @@ class _IdOverride: def _to_otel_trace_id(execution_arn: str, start_timestamp: datetime) -> int: """Build a deterministic OTel-compatible execution trace ID (128 bits). - The ID is independent of ambient Lambda or X-Ray trace context so the - parentless Workflow span remains the only root of the durable execution - trace. Invocation spans inherit ambient context separately. + The ID is used when the backend does not provide a valid trace ID. In that + case a deterministic synthetic execution root anchors the durable execution + trace across reinvocations. Raises: ValueError: If the execution start timestamp is missing. @@ -80,6 +80,22 @@ def derive_workflow_span_id(durable_execution_arn: str) -> int: return span_id or 1 +def derive_execution_root_span_id(durable_execution_arn: str) -> int: + """Derive the deterministic synthetic execution-root span ID. + + The synthetic root is a non-recording parent context used when the backend + does not provide a complete remote parent. Its ID is stable across + reinvocations and uses a namespace distinct from Workflow and operation + span IDs. + """ + if not durable_execution_arn: + raise ValueError("execution ARN is required to derive an execution root ID") + plain_value = f"execution-root:{durable_execution_arn}" + hashed = hashlib.blake2b(plain_value.encode()).hexdigest()[:16] + span_id = int(hashed, 16) + return span_id or 1 + + class DeterministicIdGenerator(RandomIdGenerator): """An ID generator with invocation-scoped deterministic ID overrides. diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py new file mode 100644 index 00000000..5f2a13a4 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py @@ -0,0 +1,135 @@ +"""Durable execution sampling support.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from opentelemetry import context as otel_context +from opentelemetry.context import Context +from opentelemetry.sdk.trace import Tracer as SdkTracer +from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult +from opentelemetry.trace import Span, SpanContext, SpanKind, TraceFlags + +from aws_durable_execution_sdk_python_otel.context_extractors import ( + ExtractedContext, + Sampling, +) + + +_DURABLE_SAMPLING_INTENT_KEY = otel_context.create_key( + "aws_durable_execution_sampling_intent" +) + + +@dataclass(frozen=True) +class DurableSamplingIntent: + """Sampling result to apply to each durable span in one invocation.""" + + result: SamplingResult + + +class DurableSampler(Sampler): + """Sampler that honors a durable sampling intent carried on parent context.""" + + def __init__(self, delegate: Sampler) -> None: + self.delegate = delegate + + @classmethod + def install_on_tracer(cls, tracer: SdkTracer) -> "DurableSampler": + current_sampler = tracer.sampler + if isinstance(current_sampler, cls): + return current_sampler + sampler = cls(current_sampler) + tracer.sampler = sampler + return sampler + + def should_sample( + self, + parent_context: Context | None, + trace_id: int, + name: str, + kind: SpanKind | None = None, + attributes: Any = None, + links: Any = None, + trace_state: Any = None, + ) -> SamplingResult: + intent = otel_context.get_value(_DURABLE_SAMPLING_INTENT_KEY, parent_context) + if isinstance(intent, DurableSamplingIntent): + merged_attributes = dict(attributes or {}) + merged_attributes.update(dict(intent.result.attributes or {})) + return SamplingResult( + intent.result.decision, + attributes=merged_attributes, + trace_state=intent.result.trace_state, + ) + return self.delegate.should_sample( + parent_context, + trace_id, + name, + kind, + attributes, + links, + trace_state, + ) + + def get_description(self) -> str: + return f"DurableSampler{{{self.delegate.get_description()}}}" + + +def store_sampling_intent( + parent_context: Context, + intent: DurableSamplingIntent | None, +) -> Context: + """Attach a durable sampling intent to a span parent context.""" + if intent is None: + return parent_context + return otel_context.set_value(_DURABLE_SAMPLING_INTENT_KEY, intent, parent_context) + + +def resolve_sampling_result( + *, + extracted: ExtractedContext | None, + ambient_span: Span, + canonical_trace_id: int, + sampler: Sampler, + span_name: str, + attributes: dict[str, Any], +) -> SamplingResult: + """Resolve one sampling decision for all durable spans in an invocation.""" + sampling = extracted.sampling if extracted is not None else Sampling.UNDECIDED + if sampling is Sampling.SAMPLED: + return SamplingResult(Decision.RECORD_AND_SAMPLE) + if sampling is Sampling.NOT_SAMPLED: + return SamplingResult(Decision.DROP) + + ambient_context = ambient_span.get_span_context() + if _is_same_trace(ambient_context, canonical_trace_id): + if bool(ambient_context.trace_flags & TraceFlags.SAMPLED): + decision = Decision.RECORD_AND_SAMPLE + elif ambient_span.is_recording(): + decision = Decision.RECORD_ONLY + else: + decision = Decision.DROP + return SamplingResult( + decision, + trace_state=ambient_context.trace_state, + ) + + return sampler.should_sample( + parent_context=Context(), + trace_id=canonical_trace_id, + name=span_name, + kind=SpanKind.INTERNAL, + attributes=attributes, + links=(), + trace_state=None, + ) + + +def is_sampled(result: SamplingResult) -> bool: + return result.decision is Decision.RECORD_AND_SAMPLE + + +def _is_same_trace(span_context: SpanContext, trace_id: int) -> bool: + return span_context.is_valid and span_context.trace_id == trace_id diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py index 3f7c10fb..0352ebfa 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_plugin.py @@ -5,11 +5,11 @@ Workflow -> Operation -> Attempt that stitches a single trace across every Lambda invocation of one durable -execution. The Workflow span is the root (created in an empty context so it -never has a parent) and is exported exactly once, when the execution reaches a -terminal status. Operations are parented under the Workflow span (or their -parent operation) and *linked* to the current Invocation span. The Invocation -span belongs to the ambient Lambda trace instead of the Workflow trace. +execution. Workflow and Invocation spans parent onto the same execution +ancestor: a propagated backend parent when present, otherwise a deterministic +synthetic root. The Workflow span is exported exactly once, when the execution +reaches a terminal status. Operations are parented under the Workflow span (or +their parent operation) and *linked* to the current Invocation span. This is the Python adaptation of the JS ``ExecutionOtelPlugin`` from aws-durable-execution-sdk-js#729. Because the Python plugin interface differs @@ -47,8 +47,10 @@ from opentelemetry import trace from opentelemetry.context import Context from opentelemetry.sdk.trace import Tracer as SdkTracer +from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import ( Link, + NonRecordingSpan, Span, SpanContext, SpanKind, @@ -58,14 +60,26 @@ from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, + ExtractedContext, + _ensure_extracted_context, xray_context_extractor, ) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, - _to_otel_trace_id, derive_workflow_span_id, operation_id_to_span_id, ) +from aws_durable_execution_sdk_python_otel.durable_sampling import ( + DurableSampler, + DurableSamplingIntent, + is_sampled, + resolve_sampling_result, + store_sampling_intent, +) +from aws_durable_execution_sdk_python_otel.execution_trace_context import ( + ExecutionTraceContext, + canonical_trace_id, +) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider @@ -113,12 +127,15 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) self._id_generator = DeterministicIdGenerator() + self._sampling_delegate: Sampler | None = None self._bind_sdk_tracer() # Per-invocation state. self._execution_arn = "" self._execution_trace_id: int | None = None - self._extracted_context: Context | None = None + self._extracted_context: ExtractedContext | None = None + self._execution_trace_context: ExecutionTraceContext | None = None + self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None self._invocation_span: Span | None = None self._operation_spans: dict[str, Span] = {} @@ -135,6 +152,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: def _bind_sdk_tracer(self) -> bool: """Bind to an SDK tracer, retrying a deferred global provider.""" + self._sampling_delegate = None tracer = self._tracer if not isinstance(tracer, SdkTracer): if self._uses_global_provider: @@ -147,6 +165,7 @@ def _bind_sdk_tracer(self) -> bool: # Deterministic stitching is scoped to this instrumentation tracer so # unrelated tracers on the same provider keep their original generator. self._id_generator = DeterministicIdGenerator.install_on_tracer(tracer) + self._sampling_delegate = DurableSampler.install_on_tracer(tracer).delegate return True # ------------------------------------------------------------------ @@ -274,14 +293,26 @@ def _resolve_parent(self, parent_id: str | None) -> Span | None: return self._workflow_span def _invocation_parent_context(self) -> Context: - """Return the active ambient context, then extracted upstream context.""" - ambient_context = otel_context.get_current() - ambient_span_context = trace.get_current_span( - ambient_context - ).get_span_context() - if ambient_span_context.is_valid: - return ambient_context - return self._extracted_context or ambient_context + """Return same-trace ambient context, else execution ancestor context.""" + execution_trace_context = self._execution_trace_context + if execution_trace_context is None: + return self._with_sampling(Context()) + + ambient_span = trace.get_current_span() + ambient_context = ambient_span.get_span_context() + if ( + ambient_context.is_valid + and ambient_context.trace_id == execution_trace_context.trace_id + ): + return self._with_sampling( + trace.set_span_in_context(ambient_span, Context()) + ) + + ancestor = NonRecordingSpan(execution_trace_context.execution_ancestor) + return self._with_sampling(trace.set_span_in_context(ancestor, Context())) + + def _with_sampling(self, parent_context: Context) -> Context: + return store_sampling_intent(parent_context, self._sampling_intent) # ------------------------------------------------------------------ # Invocation lifecycle @@ -307,13 +338,46 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: return self._execution_arn = info.execution_arn or "" - self._execution_trace_id = _to_otel_trace_id( - self._execution_arn, info.execution_start_time + if not self._execution_arn: + logger.warning( + "ExecutionOtelPlugin requires InvocationStartInfo.execution_arn " + "to derive a deterministic execution root; telemetry is disabled " + "for this invocation." + ) + self._tracing_enabled = False + return + self._extracted_context = _ensure_extracted_context( + self._context_extractor(info) + ) + self._execution_trace_id = canonical_trace_id( + extracted=self._extracted_context, + execution_arn=self._execution_arn, + execution_start_time=info.execution_start_time, + ) + if self._sampling_delegate is None: + logger.warning( + "No sampler available; telemetry is disabled for this invocation." + ) + self._tracing_enabled = False + return + sampling_result = resolve_sampling_result( + extracted=self._extracted_context, + ambient_span=trace.get_current_span(), + canonical_trace_id=self._execution_trace_id, + sampler=self._sampling_delegate, + span_name=self._workflow_span_name, + attributes={"durable.execution.arn": self._execution_arn}, + ) + self._sampling_intent = DurableSamplingIntent(sampling_result) + self._execution_trace_context = ExecutionTraceContext.resolve( + extracted=self._extracted_context, + canonical_trace_id=self._execution_trace_id, + execution_arn=self._execution_arn, + root_sampled=lambda: is_sampled(sampling_result), ) - self._extracted_context = self._context_extractor(info) self._start_workflow_span(info) - # Keep the invocation in the ambient Lambda trace in both provider modes. + # Keep the invocation on the shared execution trace. self._start_invocation_span(info) # Make the Workflow span the active span so auto-instrumented spans @@ -323,16 +387,25 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: if self._workflow_span is not None: self._attach_context( _INVOCATION_CONTEXT_KEY, - trace.set_span_in_context(self._workflow_span, self._extracted_context), + trace.set_span_in_context( + self._workflow_span, otel_context.get_current() + ), ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: if not self._execution_arn: logger.warning("No execution ARN; skipping Workflow span creation") return - # Empty context => root span with no parent. + if self._execution_trace_context is None: + return + parent_context = self._with_sampling( + trace.set_span_in_context( + NonRecordingSpan(self._execution_trace_context.execution_ancestor), + Context(), + ) + ) with self._id_generator.use_ids( - trace_id=self._execution_trace_id, + trace_id=None, span_id=derive_workflow_span_id(self._execution_arn), ): self._workflow_span = self._tracer.start_span( @@ -340,7 +413,7 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None: kind=SpanKind.INTERNAL, attributes={"durable.execution.arn": self._execution_arn}, start_time=_to_otel_timestamp(info.execution_start_time), - context=Context(), + context=parent_context, ) def _start_invocation_span(self, info: InvocationStartInfo) -> None: @@ -420,6 +493,8 @@ def _reset_state(self) -> None: self._execution_arn = "" self._execution_trace_id = None self._extracted_context = None + self._execution_trace_context = None + self._sampling_intent = None self._workflow_span = None self._invocation_span = None with self._lock: @@ -500,12 +575,12 @@ def _start_span( ) if parent is None: - parent_ctx = self._extracted_context or Context() + parent_ctx = self._with_sampling(Context()) else: - parent_ctx = trace.set_span_in_context(parent, self._extracted_context) - with self._id_generator.use_ids( - trace_id=self._execution_trace_id, span_id=span_id - ): + parent_ctx = self._with_sampling( + trace.set_span_in_context(parent, Context()) + ) + with self._id_generator.use_ids(trace_id=None, span_id=span_id): span = self._tracer.start_span( name=name, attributes=self._operation_attributes(info), @@ -552,7 +627,7 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: start_time=info.start_time, ) self._attach_context( - key, trace.set_span_in_context(span, self._extracted_context) + key, trace.set_span_in_context(span, otel_context.get_current()) ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_trace_context.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_trace_context.py new file mode 100644 index 00000000..2875abd8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_trace_context.py @@ -0,0 +1,98 @@ +"""Execution trace ancestry for durable execution telemetry.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from datetime import datetime + +from opentelemetry.trace import SpanContext, TraceFlags, TraceState + +from aws_durable_execution_sdk_python_otel.context_extractors import ( + ExtractedContext, + Sampling, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + _to_otel_trace_id, + derive_execution_root_span_id, +) + + +@dataclass(frozen=True) +class ExecutionTraceContext: + """Common ancestor for Workflow and Invocation spans.""" + + execution_ancestor: SpanContext + + @property + def trace_id(self) -> int: + return self.execution_ancestor.trace_id + + @property + def trace_flags(self) -> TraceFlags: + return self.execution_ancestor.trace_flags + + @classmethod + def resolve( + cls, + *, + extracted: ExtractedContext | None, + canonical_trace_id: int, + execution_arn: str, + root_sampled: Callable[[], bool], + ) -> "ExecutionTraceContext": + """Resolve the execution ancestor. + + A complete extracted remote parent is authoritative. Otherwise a + deterministic synthetic root anchors all invocations of the execution on + the same trace. + """ + sampling = extracted.sampling if extracted is not None else Sampling.UNDECIDED + trace_flags = _trace_flags(sampling, root_sampled) + if extracted is not None and extracted.has_complete_remote_parent: + return cls( + SpanContext( + trace_id=canonical_trace_id, + span_id=extracted.parent_span_id or 0, + is_remote=True, + trace_flags=trace_flags, + trace_state=TraceState(), + ) + ) + + return cls( + SpanContext( + trace_id=canonical_trace_id, + span_id=derive_execution_root_span_id(execution_arn), + is_remote=False, + trace_flags=trace_flags, + trace_state=TraceState(), + ) + ) + + +def canonical_trace_id( + *, + extracted: ExtractedContext | None, + execution_arn: str, + execution_start_time: datetime, +) -> int: + """Return the stable trace ID for this durable execution.""" + if extracted is not None and extracted.has_valid_trace_id: + return extracted.trace_id or 0 + return _to_otel_trace_id(execution_arn, execution_start_time) + + +def _trace_flags( + sampling: Sampling, + root_sampled: Callable[[], bool], +) -> TraceFlags: + if sampling is Sampling.SAMPLED: + return TraceFlags(TraceFlags.SAMPLED) + if sampling is Sampling.NOT_SAMPLED: + return TraceFlags(TraceFlags.DEFAULT) + return ( + TraceFlags(TraceFlags.SAMPLED) + if root_sampled() + else TraceFlags(TraceFlags.DEFAULT) + ) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py index 3b95a5d5..f9e390c4 100644 --- a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/invocation_plugin.py @@ -22,8 +22,10 @@ from opentelemetry import context, trace from opentelemetry.context import Context from opentelemetry.sdk.trace import Tracer as SdkTracer +from opentelemetry.sdk.trace.sampling import Sampler from opentelemetry.trace import ( Link, + NonRecordingSpan, Span, SpanContext, SpanKind, @@ -33,14 +35,26 @@ from aws_durable_execution_sdk_python_otel.context_extractors import ( ContextExtractor, + ExtractedContext, + _ensure_extracted_context, xray_context_extractor, ) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, - _to_otel_trace_id, derive_workflow_span_id, operation_id_to_span_id, ) +from aws_durable_execution_sdk_python_otel.durable_sampling import ( + DurableSampler, + DurableSamplingIntent, + is_sampled, + resolve_sampling_result, + store_sampling_intent, +) +from aws_durable_execution_sdk_python_otel.execution_trace_context import ( + ExecutionTraceContext, + canonical_trace_id, +) from aws_durable_execution_sdk_python_otel.log_filter import install_log_filter from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider @@ -51,12 +65,13 @@ _TERMINAL_INVOCATION_STATUSES = frozenset( {InvocationStatus.SUCCEEDED, InvocationStatus.FAILED} ) +_TIMESTAMP_STEP_NANOS = 1_000 _SpanAttributes = dict[str, str | bool | int] -def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: - """Convert a datetime to OTel timestamp (nanoseconds since epoch), or None.""" +def _to_otel_timestamp(dt: datetime.datetime | None) -> int: + """Convert a datetime to OTel timestamp (nanoseconds since epoch).""" if dt is None: dt = datetime.datetime.now(datetime.UTC) return int(dt.timestamp() * 1_000_000_000) @@ -66,10 +81,12 @@ class InvocationOtelPlugin(DurableInstrumentationPlugin): """OpenTelemetry instrumentation plugin for durable executions. The plugin creates spans for Lambda invocations, durable operations, and - user-function attempts. The Workflow trace ID is derived from the durable - execution ARN and start time. Invocation spans inherit ambient or extracted - upstream context, and operation spans are correlated with the Workflow by a - span link. + user-function attempts. Workflow and Invocation spans share one execution + trace, parented to a propagated backend parent when present or a + deterministic synthetic execution root otherwise. Invocation spans use the + active ambient span only when it is already on that execution trace. + Operation spans remain parented to the Invocation span and link to the + Workflow span. Operation IDs are converted into deterministic span IDs. The first observed span for an operation uses that deterministic ID; later continuation spans @@ -114,13 +131,17 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: self._uses_global_provider = result.uses_global_provider self._tracer: Tracer = self._provider.get_tracer(self._config.instrument_name) self._id_generator = DeterministicIdGenerator() + self._sampling_delegate: Sampler | None = None self._bind_sdk_tracer() # per invocation status: self._execution_arn = "" self._execution_trace_id: int | None = None - self._extracted_context: Context | None = None + self._extracted_context: ExtractedContext | None = None + self._execution_trace_context: ExecutionTraceContext | None = None + self._sampling_intent: DurableSamplingIntent | None = None self._workflow_span: Span | None = None + self._span_time_floor_ns: int | None = None # Maps operation ID (None for root) to the active span. self._operation_spans: dict[str | None, Span] = {} # Tokens returned by context.attach(), keyed by the span registry key, @@ -140,6 +161,7 @@ def __init__(self, config: OtelPluginConfig | None = None) -> None: def _bind_sdk_tracer(self) -> bool: """Bind to an SDK tracer, retrying a deferred global provider.""" + self._sampling_delegate = None tracer = self._tracer if not isinstance(tracer, SdkTracer): if self._uses_global_provider: @@ -150,6 +172,7 @@ def _bind_sdk_tracer(self) -> bool: return False self._id_generator = DeterministicIdGenerator.install_on_tracer(tracer) + self._sampling_delegate = DurableSampler.install_on_tracer(tracer).delegate return True def _set_span(self, operation_id: str | None, span: Span) -> None: @@ -248,12 +271,11 @@ def get_current_span_context(self) -> SpanContext | None: """Return the span context to use for log correlation. Resolution order: - 1. The span attached to the OTel thread-local context. Inside a step - this is the active attempt span, and inside a child context this is - the active context span (attached in - on_user_function_start), and between the steps of a child context it - is the enclosing context span, restored when on_user_function_end - detaches the inner scope. + 1. The same-trace span attached to the OTel thread-local context. + Inside a step this is the active attempt span, and inside a child + context this is the active context span (attached in + on_user_function_start). Unrelated ambient spans are ignored so logs + stay correlated to the durable execution trace. 2. The invocation span from the plugin registry. This is the path used for top-level handler code: the invocation span is never attached to the worker thread's context, so the registry is the only way to @@ -265,7 +287,11 @@ def get_current_span_context(self) -> SpanContext | None: A valid SpanContext, or None if no span is active. """ span_context = trace.get_current_span().get_span_context() - if span_context and span_context.is_valid: + if ( + span_context + and span_context.is_valid + and span_context.trace_id == self._execution_trace_id + ): return span_context invocation_span = self._get_span(None) @@ -298,14 +324,52 @@ def _resolve_parent_span(self, parent_id: str | None = None) -> Span: raise ValueError("No parent span found") def _invocation_parent_context(self) -> Context: - """Return the active ambient context, then extracted upstream context.""" - ambient_context = context.get_current() - ambient_span_context = trace.get_current_span( - ambient_context - ).get_span_context() - if ambient_span_context.is_valid: - return ambient_context - return self._extracted_context or ambient_context + """Return same-trace ambient context, else execution ancestor context.""" + execution_trace_context = self._execution_trace_context + if execution_trace_context is None: + return self._with_sampling(Context()) + + ambient_span = trace.get_current_span() + ambient_context = ambient_span.get_span_context() + if ( + ambient_context.is_valid + and ambient_context.trace_id == execution_trace_context.trace_id + ): + return self._with_sampling( + trace.set_span_in_context(ambient_span, Context()) + ) + + ancestor = NonRecordingSpan(execution_trace_context.execution_ancestor) + return self._with_sampling(trace.set_span_in_context(ancestor, Context())) + + def _with_sampling(self, parent_context: Context) -> Context: + return store_sampling_intent(parent_context, self._sampling_intent) + + def _next_ordered_timestamp( + self, + timestamp: datetime.datetime | None = None, + ) -> int: + """Return an invocation-local timestamp that does not move backward.""" + candidate = _to_otel_timestamp(timestamp) + with self._operation_spans_lock: + floor = self._span_time_floor_ns + if floor is not None and candidate <= floor: + candidate = floor + _TIMESTAMP_STEP_NANOS + self._span_time_floor_ns = candidate + return candidate + + def _operation_link_context(self, operation_id: str) -> SpanContext | None: + """Return the deterministic logical operation context for links.""" + execution_trace_context = self._execution_trace_context + if execution_trace_context is None: + return None + return SpanContext( + trace_id=execution_trace_context.trace_id, + span_id=operation_id_to_span_id(self._execution_arn, operation_id), + is_remote=False, + trace_flags=execution_trace_context.trace_flags, + trace_state=execution_trace_context.execution_ancestor.trace_state, + ) def _start_span( self, @@ -357,6 +421,10 @@ def _start_span( if operation_id else None ) + if existed and operation_id is not None: + operation_context = self._operation_link_context(operation_id) + if operation_context is not None and operation_context.is_valid: + links = [*links, Link(context=operation_context)] # Operation and attempt spans link to the execution-scoped Workflow # span (the invocation span itself, operation_id=None, does not). if self._workflow_span is not None and operation_id is not None: @@ -366,20 +434,26 @@ def _start_span( if parent_span is None: parent_context = self._invocation_parent_context() else: - parent_context = trace.set_span_in_context( - parent_span, self._extracted_context + parent_context = self._with_sampling( + trace.set_span_in_context(parent_span, Context()) ) - trace_id = self._execution_trace_id if operation_id is not None else None - with self._id_generator.use_ids(trace_id=trace_id, span_id=span_id): + span_start_time = ( + self._next_ordered_timestamp(start_time) + if operation_id is not None + else _to_otel_timestamp(start_time) + ) + with self._id_generator.use_ids(trace_id=None, span_id=span_id): span = self._tracer.start_span( name=name, kind=SpanKind.INTERNAL, attributes=attributes, - start_time=_to_otel_timestamp(start_time), + start_time=span_start_time, context=parent_context, links=links, ) self._operation_spans[registry_key] = span + if operation_id is None: + self._span_time_floor_ns = span_start_time logger.debug("Started OTel span: %s", span) return span @@ -400,7 +474,7 @@ def _end_span( span = self._operation_spans.pop(operation_id, None) if span: # the span is not going to be populated if it has the same end_time and start_time - end_time = _to_otel_timestamp(end_timestamp) if end_timestamp else None + end_time = self._next_ordered_timestamp(end_timestamp) span.end(end_time=end_time) logger.debug("Ended OTel span: %s", span) @@ -429,10 +503,43 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: return self._execution_arn = info.execution_arn or "" - self._execution_trace_id = _to_otel_trace_id( - self._execution_arn, info.execution_start_time + if not self._execution_arn: + logger.warning( + "InvocationOtelPlugin requires InvocationStartInfo.execution_arn " + "to derive a deterministic execution root; telemetry is disabled " + "for this invocation." + ) + self._tracing_enabled = False + return + self._extracted_context = _ensure_extracted_context( + self._context_extractor(info) + ) + self._execution_trace_id = canonical_trace_id( + extracted=self._extracted_context, + execution_arn=self._execution_arn, + execution_start_time=info.execution_start_time, + ) + if self._sampling_delegate is None: + logger.warning( + "No sampler available; telemetry is disabled for this invocation." + ) + self._tracing_enabled = False + return + sampling_result = resolve_sampling_result( + extracted=self._extracted_context, + ambient_span=trace.get_current_span(), + canonical_trace_id=self._execution_trace_id, + sampler=self._sampling_delegate, + span_name=self._workflow_span_name, + attributes={"durable.execution.arn": self._execution_arn}, + ) + self._sampling_intent = DurableSamplingIntent(sampling_result) + self._execution_trace_context = ExecutionTraceContext.resolve( + extracted=self._extracted_context, + canonical_trace_id=self._execution_trace_id, + execution_arn=self._execution_arn, + root_sampled=lambda: is_sampled(sampling_result), ) - self._extracted_context = self._context_extractor(info) self._start_workflow_span(info) @@ -443,23 +550,28 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: ) def _start_workflow_span(self, info: InvocationStartInfo) -> None: - """Create the deterministic, execution-scoped Workflow root span. - - The Workflow span is a parentless root keyed to a deterministic span ID - derived from the execution ARN, so every invocation of the same durable - execution contributes to one Workflow span. It is exported once, on a - terminal invocation. Operation and attempt spans link to it while - remaining parented to the invocation span. It is created unconditionally - -- InvocationOtelPlugin has no default/owned tracer-provider distinction, - so the span is emitted whether the provider is the ambient (ADOT/global) - one or an explicitly supplied one. + """Create the deterministic, execution-scoped Workflow span. + + The Workflow span is keyed to a deterministic span ID derived from the + execution ARN, so every invocation of the same durable execution + contributes to one Workflow span. It is parented to the shared execution + ancestor and exported once, on a terminal invocation. Operation and + attempt spans link to it while remaining parented to the invocation + span. """ if not self._execution_arn: logger.warning("No execution ARN; skipping Workflow span creation") return - # Empty context => root span with no parent. + if self._execution_trace_context is None: + return + parent_context = self._with_sampling( + trace.set_span_in_context( + NonRecordingSpan(self._execution_trace_context.execution_ancestor), + Context(), + ) + ) with self._id_generator.use_ids( - trace_id=self._execution_trace_id, + trace_id=None, span_id=derive_workflow_span_id(self._execution_arn), ): self._workflow_span = self._tracer.start_span( @@ -467,7 +579,7 @@ def _start_workflow_span(self, info: InvocationStartInfo) -> None: kind=SpanKind.INTERNAL, attributes={"durable.execution.arn": self._execution_arn}, start_time=_to_otel_timestamp(info.execution_start_time), - context=Context(), + context=parent_context, ) def on_invocation_end(self, info: InvocationEndInfo) -> None: @@ -535,7 +647,10 @@ def _reset_state(self) -> None: self._execution_arn = "" self._execution_trace_id = None self._extracted_context = None + self._execution_trace_context = None + self._sampling_intent = None self._workflow_span = None + self._span_time_floor_ns = None with self._operation_spans_lock: self._operation_spans = {} self._tracing_enabled = False @@ -573,9 +688,9 @@ def on_operation_end(self, info: OperationEndInfo) -> None: return span = self._get_span(info.operation_id) if span is None: - # The operation started in a prior invocation. The prior SpanContext - # is not checkpointed, so create a new correlated segment without a - # fabricated link. + # The operation started in a prior invocation. Create a new + # correlated segment and link it to the deterministic logical + # operation context shared across invocations. parent_span = self._resolve_parent_span(info.parent_id) attributes = self._extract_attributes(info) span = self._start_span( @@ -597,7 +712,7 @@ def on_operation_end(self, info: OperationEndInfo) -> None: else: span.set_status(StatusCode.OK) - self._end_span(info.operation_id) + self._end_span(info.operation_id, info.end_time) def on_user_function_start(self, info: UserFunctionStartInfo) -> None: """Called when a context or step operation starts user code. @@ -630,18 +745,23 @@ def on_user_function_start(self, info: UserFunctionStartInfo) -> None: if info.operation_type is OperationType.STEP: span_name = f"{span_name} attempt {info.attempt or 1}" span_key = self._user_function_span_key(info) + span_start_time = ( + datetime.datetime.now(datetime.UTC) + if info.operation_type is OperationType.CONTEXT + else info.start_time + ) span = self._start_span( operation_id=info.operation_id, name=span_name, attributes=attributes, - start_time=info.start_time, + start_time=span_start_time, parent_span=parent_span, existed=info.attempt != 1 and info.operation_type is not OperationType.STEP, span_key=span_key, deterministic_span_id=info.operation_type is not OperationType.STEP, ) self._attach_context( - span_key, trace.set_span_in_context(span, self._extracted_context) + span_key, trace.set_span_in_context(span, context.get_current()) ) def on_user_function_end(self, info: UserFunctionEndInfo) -> None: diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py new file mode 100644 index 00000000..07513662 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_invocation_wait_resume_int.py @@ -0,0 +1,252 @@ +"""End-to-end invocation-view OTel coverage for wait/resume.""" + +from __future__ import annotations + +from dataclasses import replace +from datetime import UTC, datetime +from typing import Any +from unittest.mock import Mock, patch + +import pytest +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.context import DurableContext, durable_step +from aws_durable_execution_sdk_python.execution import ( + InvocationStatus, + durable_execution, +) +from aws_durable_execution_sdk_python.lambda_service import ( + CheckpointOutput, + CheckpointUpdatedExecutionState, + ExecutionDetails, + Operation, + OperationAction, + OperationStatus, + OperationSubType, + OperationType, + StepDetails, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + derive_workflow_span_id, +) +from aws_durable_execution_sdk_python_otel.execution_plugin import ExecutionOtelPlugin +from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + +EXECUTION_ARN = "test-arn/execution-otel-wait-resume" +EXECUTION_START = datetime(2026, 8, 27, 5, 11, 47, tzinfo=UTC) +XRAY_TRACE_HEADER = ( + "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1" +) +XRAY_TRACE_ID = int("5759e988bd862e3fe1be46a994272793", 16) +XRAY_PARENT_SPAN_ID = int("53995c3f42cd8ad8", 16) + + +def _lambda_context() -> Mock: + context = Mock() + context.aws_request_id = "test-request-id" + context.client_context = None + context.identity = None + context._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + context.invoked_function_arn = "test-arn" + context.tenant_id = None + return context + + +def _event( + operations: list[Operation], + updated_operation_ids: list[str] | None = None, +) -> dict[str, Any]: + event: dict[str, Any] = { + "DurableExecutionArn": EXECUTION_ARN, + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [operation.to_json_dict() for operation in operations], + "NextMarker": "", + }, + "LocalRunner": True, + } + if updated_operation_ids is not None: + event["UpdatedOperationIds"] = updated_operation_ids + return event + + +def _execution_operation() -> Operation: + return Operation( + operation_id="execution-otel-wait-resume", + operation_type=OperationType.EXECUTION, + status=OperationStatus.STARTED, + start_timestamp=EXECUTION_START, + execution_details=ExecutionDetails(input_payload="{}"), + ) + + +def _checkpoint_store(initial_operations: list[Operation]): + operations = {operation.operation_id: operation for operation in initial_operations} + + def checkpoint( + durable_execution_arn, # noqa: ARG001 + checkpoint_token, # noqa: ARG001 + updates, + client_token="token", # noqa: S107, ARG001 + ) -> CheckpointOutput: + for update in updates: + now = datetime.now(UTC) + previous = operations.get(update.operation_id) + if update.action is OperationAction.START: + operations[update.operation_id] = Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=now, + ) + elif update.action is OperationAction.SUCCEED: + base = previous or Operation( + operation_id=update.operation_id, + operation_type=update.operation_type, + status=OperationStatus.STARTED, + parent_id=update.parent_id, + name=update.name, + sub_type=update.sub_type, + start_timestamp=now, + ) + operations[update.operation_id] = replace( + base, + status=OperationStatus.SUCCEEDED, + end_timestamp=now, + step_details=( + StepDetails(result=update.payload, attempt=1) + if update.operation_type is OperationType.STEP + else base.step_details + ), + ) + + return CheckpointOutput( + checkpoint_token="new-token", + new_execution_state=CheckpointUpdatedExecutionState( + operations=list(operations.values()) + ), + ) + + return checkpoint, operations + + +@pytest.mark.parametrize( + "plugin_type", + [InvocationOtelPlugin, ExecutionOtelPlugin], +) +def test_otel_wait_resume_spans_share_default_xray_execution_trace( + monkeypatch: pytest.MonkeyPatch, + plugin_type: type[InvocationOtelPlugin] | type[ExecutionOtelPlugin], +) -> None: + monkeypatch.setenv("_X_AMZN_TRACE_ID", XRAY_TRACE_HEADER) + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + enrich_logger=False, + ) + ) + + @durable_step + def complete_after_resume(_step_context) -> str: + return "resumed" + + def handler_impl(_event: Any, context: DurableContext) -> str: + context.wait(Duration.from_seconds(1), name="otel-wait") + return context.step(complete_after_resume(), name="otel-after-resume") + + handler = durable_execution(handler_impl, plugins=[plugin]) + + initial_operations = [_execution_operation()] + first_checkpoint, first_operations = _checkpoint_store(initial_operations) + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = first_checkpoint + mock_client_class.initialize_client.return_value = mock_client + + first_result = handler(_event(initial_operations), _lambda_context()) + + assert first_result["Status"] == InvocationStatus.PENDING.value + wait_operation = next( + operation + for operation in first_operations.values() + if operation.name == "otel-wait" + ) + completed_wait = replace( + wait_operation, + status=OperationStatus.SUCCEEDED, + end_timestamp=datetime.now(UTC), + sub_type=OperationSubType.WAIT, + ) + + second_operations = [_execution_operation(), completed_wait] + second_checkpoint, _ = _checkpoint_store(second_operations) + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client.checkpoint = second_checkpoint + mock_client_class.initialize_client.return_value = mock_client + + second_result = handler( + _event( + second_operations, + updated_operation_ids=[completed_wait.operation_id], + ), + _lambda_context(), + ) + + assert second_result["Status"] == InvocationStatus.SUCCEEDED.value + + spans = exporter.get_finished_spans() + durable_spans = [ + span + for span in spans + if span.name in {"Workflow", "Invocation", "otel-wait", "otel-after-resume"} + ] + trace_ids = {span.context.trace_id for span in durable_spans} + assert trace_ids == {XRAY_TRACE_ID} + + workflow = next(span for span in spans if span.name == "Workflow") + invocations = [span for span in spans if span.name == "Invocation"] + waits = [span for span in spans if span.name == "otel-wait"] + after_resume = next(span for span in spans if span.name == "otel-after-resume") + + assert len(invocations) >= 2 + assert len(waits) >= 2 + assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) + assert workflow.parent is not None + assert workflow.parent.span_id == XRAY_PARENT_SPAN_ID + assert {span.parent.span_id for span in invocations if span.parent} == { + XRAY_PARENT_SPAN_ID + } + + assert after_resume.parent is not None + if plugin_type is InvocationOtelPlugin: + assert after_resume.parent.span_id in { + span.context.span_id for span in invocations + } + else: + assert after_resume.parent.span_id == workflow.context.span_id + completed_wait_span = next( + span + for span in waits + if span.parent is not None + and span.parent.span_id == after_resume.parent.span_id + ) + assert completed_wait_span.end_time is not None + assert after_resume.start_time is not None + assert completed_wait_span.end_time <= after_resume.start_time diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_extractors.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_extractors.py index d150af92..5ed3a8db 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_context_extractors.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_context_extractors.py @@ -2,65 +2,79 @@ from __future__ import annotations -from opentelemetry.context import Context +import pytest from aws_durable_execution_sdk_python_otel import context_extractors +from aws_durable_execution_sdk_python_otel.context_extractors import ( + ExtractedContext, + Sampling, +) -def test_xray_context_extractor_returns_current_context_without_trace_header( - monkeypatch, -): - """Verify absent X-Ray trace headers leave the active context unchanged.""" - current_context = Context({"durable": "current"}) +def test_xray_context_extractor_returns_none_without_trace_header(monkeypatch): monkeypatch.delenv("_X_AMZN_TRACE_ID", raising=False) - monkeypatch.setattr( - context_extractors.otel_context, - "get_current", - lambda: current_context, + + assert context_extractors.xray_context_extractor(object()) is None + + +def test_xray_context_extractor_extracts_trace_parent_and_sampling(monkeypatch): + monkeypatch.setenv( + "_X_AMZN_TRACE_ID", + "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1", ) - assert context_extractors.xray_context_extractor(object()) is current_context + extracted = context_extractors.xray_context_extractor(object()) + assert extracted is not None + assert extracted.trace_id == int("5759e988bd862e3fe1be46a994272793", 16) + assert extracted.parent_span_id == int("53995c3f42cd8ad8", 16) + assert extracted.sampling is Sampling.SAMPLED + assert extracted.has_complete_remote_parent -def test_xray_context_extractor_extracts_trace_header_from_environment( - monkeypatch, -): - """Verify X-Ray trace headers are passed through OpenTelemetry propagation.""" - trace_header = ( - "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=53995c3f42cd8ad8;Sampled=1" + +def test_xray_context_extractor_preserves_valid_root_without_parent(monkeypatch): + monkeypatch.setenv( + "_X_AMZN_TRACE_ID", + "Root=1-5759e988-bd862e3fe1be46a994272793;Parent=0000000000000000;Sampled=0", ) - current_context = Context({"durable": "current"}) - extracted_context = Context({"durable": "extracted"}) - extract_calls = [] - monkeypatch.setenv("_X_AMZN_TRACE_ID", trace_header) - monkeypatch.setattr( - context_extractors.otel_context, - "get_current", - lambda: current_context, + + extracted = context_extractors.xray_context_extractor(object()) + + assert extracted is not None + assert extracted.trace_id == int("5759e988bd862e3fe1be46a994272793", 16) + assert extracted.parent_span_id is None + assert extracted.sampling is Sampling.NOT_SAMPLED + assert not extracted.has_complete_remote_parent + + +def test_xray_context_extractor_drops_all_zero_root(monkeypatch): + monkeypatch.setenv( + "_X_AMZN_TRACE_ID", + "Root=1-00000000-000000000000000000000000;Parent=53995c3f42cd8ad8;Sampled=1", ) - def extract(*, carrier, context): - extract_calls.append({"carrier": carrier, "context": context}) - return extracted_context + extracted = context_extractors.xray_context_extractor(object()) + + assert extracted is not None + assert extracted.trace_id is None + assert extracted.parent_span_id == int("53995c3f42cd8ad8", 16) + assert extracted.sampling is Sampling.SAMPLED - monkeypatch.setattr(context_extractors.propagate, "extract", extract) - assert context_extractors.xray_context_extractor(object()) is extracted_context - assert extract_calls == [ - { - "carrier": {"X-Amzn-Trace-Id": trace_header}, - "context": current_context, - } - ] +def test_ensure_extracted_context_rejects_invalid_extractor_result(): + with pytest.raises(TypeError, match="ExtractedContext or None"): + context_extractors._ensure_extracted_context(object()) -def test_w3c_client_context_extractor_returns_current_context(monkeypatch): - """Verify the placeholder W3C extractor leaves the active context unchanged.""" - current_context = Context({"durable": "current"}) - monkeypatch.setattr( - context_extractors.otel_context, - "get_current", - lambda: current_context, +def test_ensure_extracted_context_accepts_structured_context(): + extracted = ExtractedContext( + trace_id=int("5759e988bd862e3fe1be46a994272793", 16), + parent_span_id=int("53995c3f42cd8ad8", 16), + sampling=Sampling.SAMPLED, ) - assert context_extractors.w3c_client_context_extractor(object()) is current_context + assert context_extractors._ensure_extracted_context(extracted) is extracted + + +def test_w3c_client_context_extractor_returns_none(): + assert context_extractors.w3c_client_context_extractor(object()) is None diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py index dc4a699c..6d8ad688 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin.py @@ -25,13 +25,13 @@ UserFunctionOutcome, UserFunctionStartInfo, ) -from opentelemetry import trace -from opentelemetry.context import Context +from opentelemetry import baggage, trace from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + derive_execution_root_span_id, derive_workflow_span_id, operation_id_to_span_id, ) @@ -59,7 +59,9 @@ def _assert_otel_context_balanced(): ) -def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: +def _create_plugin( + context_extractor=lambda _: None, +) -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: """Create an ExecutionOtelPlugin wired to an in-memory exporter.""" exporter = InMemorySpanExporter() provider = TracerProvider() @@ -67,7 +69,7 @@ def _create_plugin() -> tuple[ExecutionOtelPlugin, InMemorySpanExporter]: plugin = ExecutionOtelPlugin( OtelPluginConfig( tracer_provider=provider, - context_extractor=lambda _: Context(), + context_extractor=context_extractor, enrich_logger=False, ) ) @@ -148,7 +150,7 @@ def test_derive_workflow_span_id_rejects_empty_arn(): # --------------------------------------------------------------------------- # Workflow + invocation span hierarchy # --------------------------------------------------------------------------- -def test_workflow_and_invocation_are_separate_roots_without_ambient_parent(): +def test_workflow_and_invocation_share_execution_trace_without_ambient_parent(): plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -161,18 +163,21 @@ def test_workflow_and_invocation_are_separate_roots_without_ambient_parent(): workflow = spans["Workflow"] invocation = spans["Invocation"] - # Workflow is a root span (no parent) with the deterministic span ID. - assert workflow.parent is None + # Workflow is parented to the synthetic execution root with a deterministic + # Workflow span ID. + assert workflow.parent is not None assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) + assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) assert workflow.attributes["durable.execution.arn"] == EXECUTION_ARN assert ( workflow.attributes["durable.execution.status"] == InvocationStatus.SUCCEEDED.value ) - # Without ambient context, Invocation starts a separate provider trace. - assert invocation.parent is None - assert invocation.context.trace_id != workflow.context.trace_id + # Without extracted context, Invocation shares the synthetic execution root. + assert invocation.parent is not None + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id def test_invocation_start_without_execution_start_time_disables_tracing( @@ -193,7 +198,68 @@ def test_invocation_start_without_execution_start_time_disables_tracing( assert exporter.get_finished_spans() == () -def test_explicit_mode_invocation_span_parented_to_ambient_span(): +def test_invocation_start_without_execution_arn_disables_tracing( + caplog: pytest.LogCaptureFixture, +) -> None: + plugin, exporter = _create_plugin() + info = InvocationStartInfo( + request_id="request-1", + execution_arn=None, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + + plugin.on_invocation_start(info) + plugin.on_operation_start( + OperationStartInfo( + operation_id="after-rejected-start", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="after rejected start", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + assert "requires InvocationStartInfo.execution_arn" in caplog.text + assert exporter.get_finished_spans() == () + + +def test_invocation_start_without_sampler_disables_tracing( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + plugin, exporter = _create_plugin() + + def bind_without_sampler() -> bool: + plugin._sampling_delegate = None + return True + + monkeypatch.setattr(plugin, "_bind_sdk_tracer", bind_without_sampler) + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id="after-rejected-start", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="after rejected start", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + assert "No sampler available" in caplog.text + assert exporter.get_finished_spans() == () + + +def test_explicit_mode_invocation_span_ignores_different_trace_ambient_span(): plugin, exporter = _create_plugin() ambient = plugin._provider.get_tracer("ambient").start_span("lambda-invocation") @@ -209,9 +275,11 @@ def test_explicit_mode_invocation_span_parented_to_ambient_span(): workflow = spans["Workflow"] invocation = spans["Invocation"] assert invocation.parent is not None - assert invocation.parent.span_id == ambient.get_span_context().span_id - assert invocation.context.trace_id == ambient.get_span_context().trace_id - assert workflow.context.trace_id != ambient.get_span_context().trace_id + assert workflow.parent is not None + assert invocation.parent.span_id != ambient.get_span_context().span_id + assert invocation.context.trace_id != ambient.get_span_context().trace_id + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id def test_workflow_span_dropped_on_non_terminal_status(): @@ -473,7 +541,7 @@ def _create_default_mode_plugin( monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) plugin = ExecutionOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -492,11 +560,13 @@ def test_default_mode_creates_invocation_span(monkeypatch): invocation = spans["Invocation"] assert invocation.attributes["durable.execution.arn"] == EXECUTION_ARN assert invocation.attributes["durable.invocation.first"] is True - assert invocation.parent is None - assert invocation.context.trace_id != spans["Workflow"].context.trace_id + assert invocation.parent is not None + assert spans["Workflow"].parent is not None + assert invocation.context.trace_id == spans["Workflow"].context.trace_id + assert invocation.parent.span_id == spans["Workflow"].parent.span_id -def test_default_mode_invocation_span_parented_to_ambient_span(monkeypatch): +def test_default_mode_invocation_span_ignores_different_trace_ambient_span(monkeypatch): plugin, exporter = _create_default_mode_plugin(monkeypatch) # Simulate the ambient Lambda invocation span from the ADOT layer. @@ -511,8 +581,8 @@ def test_default_mode_invocation_span_parented_to_ambient_span(monkeypatch): invocation = {s.name: s for s in exporter.get_finished_spans()}["Invocation"] assert invocation.parent is not None - assert invocation.parent.span_id == ambient.get_span_context().span_id - assert invocation.context.trace_id == ambient.get_span_context().trace_id + assert invocation.parent.span_id != ambient.get_span_context().span_id + assert invocation.context.trace_id != ambient.get_span_context().trace_id def test_open_operation_span_not_exported_at_invocation_end(): @@ -753,6 +823,24 @@ def test_step_scope_is_released_at_user_function_end(outcome): assert plugin._context_tokens == {} +def test_user_function_start_preserves_baggage_in_current_context(): + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + baggage_context = baggage.set_baggage( + "durable-test-key", "durable-test-value", otel_context.get_current() + ) + token = otel_context.attach(baggage_context) + try: + plugin.on_user_function_start(_step_start_info("step-baggage")) + + assert baggage.get_baggage("durable-test-key") == "durable-test-value" + + plugin.on_user_function_end(_step_end_info("step-baggage")) + finally: + otel_context.detach(token) + plugin.on_invocation_end(_invocation_end_info()) + + def test_sequential_steps_do_not_accumulate_scopes(): """Verify repeated steps unwind to the same enclosing context each time.""" plugin, _ = _create_plugin() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py index 52000518..90493613 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_plugin_integration.py @@ -3,11 +3,11 @@ Drives the full plugin lifecycle against a real TracerProvider + InMemorySpanExporter for the two deployment shapes: -* Community collector layer: the caller supplies a provider; the Workflow and - Invocation spans root separate traces when no ambient parent exists. -* ADOT layer: the ADOT Lambda layer supplies the global provider and the ambient - Lambda invocation span; the plugin's Invocation span parents to that ambient - span. +* Community collector layer: the caller supplies a provider. +* ADOT layer: the ADOT Lambda layer supplies the global provider. + +Both paths keep Workflow and Invocation on one execution trace, parented to a +shared execution ancestor. """ from __future__ import annotations @@ -43,6 +43,7 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, + derive_execution_root_span_id, derive_workflow_span_id, operation_id_to_span_id, ) @@ -171,7 +172,7 @@ def _config_for_provider( monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) return OtelPluginConfig( tracer_provider=None if uses_global_provider else provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) @@ -225,7 +226,7 @@ def test_global_proxy_binds_sdk_provider_before_first_invocation( monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) plugin = ExecutionOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -250,7 +251,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) plugin = ExecutionOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -283,7 +284,7 @@ def test_community_layer_full_lifecycle_is_workflow_rooted(): plugin = ExecutionOtelPlugin( OtelPluginConfig( tracer_provider=provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -299,17 +300,20 @@ def test_community_layer_full_lifecycle_is_workflow_rooted(): operation = spans[OP_NAME] attempt = spans[f"{OP_NAME} attempt 1"] - # Workflow is the trace root with the deterministic workflow span id. - assert workflow.parent is None + # Workflow is parented to the synthetic execution ancestor with the + # deterministic workflow span id. + assert workflow.parent is not None + assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) assert ( workflow.attributes["durable.execution.status"] == InvocationStatus.SUCCEEDED.value ) - # Without ambient context, Invocation roots a separate provider trace. - assert invocation.parent is None - assert invocation.context.trace_id != workflow.context.trace_id + # Without extracted context, Invocation shares the synthetic execution root. + assert invocation.parent is not None + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id # Operation span: deterministic id, parented under Workflow, linked to invocation. assert operation.context.span_id == operation_id_to_span_id(EXECUTION_ARN, OP_ID) @@ -328,13 +332,13 @@ def test_community_layer_full_lifecycle_is_workflow_rooted(): # --------------------------------------------------------------------------- # ADOT layer (default provider; ambient invocation span) # --------------------------------------------------------------------------- -def test_adot_layer_full_lifecycle_parents_to_ambient_span(monkeypatch): +def test_adot_layer_full_lifecycle_ignores_different_trace_ambient_span(monkeypatch): provider, exporter = _provider() # Simulate the ADOT layer having configured the global TracerProvider. monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) plugin = ExecutionOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -356,11 +360,14 @@ def test_adot_layer_full_lifecycle_parents_to_ambient_span(monkeypatch): invocation = spans["Invocation"] operation = spans[OP_NAME] - # Invocation span parents to the ambient ADOT span and carries the first flag. + # Invocation ignores the different-trace ambient ADOT span and stays on the + # execution trace. assert invocation.parent is not None - assert invocation.parent.span_id == ambient.get_span_context().span_id - assert invocation.context.trace_id == ambient.get_span_context().trace_id - assert workflow.context.trace_id != ambient.get_span_context().trace_id + assert workflow.parent is not None + assert invocation.parent.span_id != ambient.get_span_context().span_id + assert invocation.context.trace_id != ambient.get_span_context().trace_id + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id assert invocation.attributes["durable.invocation.first"] is True # Operation span still uses the deterministic id and links to the durable @@ -379,7 +386,7 @@ def test_second_plugin_uses_execution_trace_id_independent_of_xray(monkeypatch): provider, exporter = _provider() monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) config = OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) first_plugin = ExecutionOtelPlugin(config) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py index 0f3bb62d..66a6300f 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin.py @@ -2,6 +2,7 @@ from __future__ import annotations +import logging import threading import time from concurrent.futures import ThreadPoolExecutor @@ -26,18 +27,33 @@ UserFunctionOutcome, UserFunctionStartInfo, ) -from opentelemetry import trace +from opentelemetry import baggage, trace from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.sampling import ALWAYS_OFF, ALWAYS_ON, Sampler from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from opentelemetry.trace import SpanKind, StatusCode +from opentelemetry.trace import ( + NonRecordingSpan, + SpanContext, + SpanKind, + StatusCode, + TraceFlags, + TraceState, +) +from aws_durable_execution_sdk_python_otel.context_extractors import ( + ExtractedContext, + Sampling, +) from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + _to_otel_trace_id, + derive_execution_root_span_id, derive_workflow_span_id, operation_id_to_span_id, ) from aws_durable_execution_sdk_python_otel.invocation_plugin import InvocationOtelPlugin +from aws_durable_execution_sdk_python_otel.log_filter import OtelContextLogFilter from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig @@ -62,14 +78,22 @@ def _assert_otel_context_balanced(): def _create_plugin() -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: + """Create a plugin wired to an in-memory span exporter.""" + return _create_plugin_with_sampler() + + +def _create_plugin_with_sampler( + sampler: Sampler | None = None, + context_extractor=lambda _: None, +) -> tuple[InvocationOtelPlugin, InMemorySpanExporter]: """Create a plugin wired to an in-memory span exporter.""" exporter = InMemorySpanExporter() - trace_provider = TracerProvider() + trace_provider = TracerProvider(sampler=sampler) trace_provider.add_span_processor(SimpleSpanProcessor(exporter)) plugin = InvocationOtelPlugin( OtelPluginConfig( tracer_provider=trace_provider, - context_extractor=lambda _: Context(), + context_extractor=context_extractor, ) ) return plugin, exporter @@ -201,7 +225,7 @@ def test_extract_attributes_uses_structural_event_attributes(): def test_invocation_start_and_end_emit_invocation_span(): - """Verify invocation lifecycle callbacks create and finish the root span.""" + """Verify invocation lifecycle callbacks create and finish the span.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) @@ -222,8 +246,11 @@ def test_invocation_start_and_end_emit_invocation_span(): == InvocationStatus.SUCCEEDED.value ) workflow = spans_by_name["Workflow"] - assert invocation.parent is None - assert invocation.context.trace_id != workflow.context.trace_id + assert invocation.parent is not None + assert workflow.parent is not None + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id + assert invocation.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) assert plugin._get_span(None) is None @@ -245,7 +272,68 @@ def test_invocation_start_without_execution_start_time_disables_tracing( assert exporter.get_finished_spans() == () -def test_invocation_span_parents_to_ambient_span(): +def test_invocation_start_without_execution_arn_disables_tracing( + caplog: pytest.LogCaptureFixture, +) -> None: + plugin, exporter = _create_plugin() + info = InvocationStartInfo( + request_id="request-1", + execution_arn=None, + execution_start_time=START_TIME, + is_first_invocation=True, + ) + + plugin.on_invocation_start(info) + plugin.on_operation_start( + OperationStartInfo( + operation_id="after-rejected-start", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="after rejected start", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + assert "requires InvocationStartInfo.execution_arn" in caplog.text + assert exporter.get_finished_spans() == () + + +def test_invocation_start_without_sampler_disables_tracing( + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + plugin, exporter = _create_plugin() + + def bind_without_sampler() -> bool: + plugin._sampling_delegate = None + return True + + monkeypatch.setattr(plugin, "_bind_sdk_tracer", bind_without_sampler) + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_operation_start( + OperationStartInfo( + operation_id="after-rejected-start", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="after rejected start", + parent_id=None, + start_time=START_TIME, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_invocation_end(_invocation_end_info()) + + assert "No sampler available" in caplog.text + assert exporter.get_finished_spans() == () + + +def test_invocation_span_ignores_different_trace_ambient_span(): plugin, exporter = _create_plugin() ambient = plugin._provider.get_tracer("ambient").start_span("lambda-invocation") @@ -261,9 +349,136 @@ def test_invocation_span_parents_to_ambient_span(): invocation = spans["Invocation"] workflow = spans["Workflow"] assert invocation.parent is not None - assert invocation.parent.span_id == ambient.get_span_context().span_id - assert invocation.context.trace_id == ambient.get_span_context().trace_id - assert workflow.context.trace_id != ambient.get_span_context().trace_id + assert workflow.parent is not None + assert invocation.parent.span_id != ambient.get_span_context().span_id + assert invocation.context.trace_id != ambient.get_span_context().trace_id + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id + + +def test_log_filter_uses_invocation_trace_when_ambient_trace_is_rejected(): + plugin, _ = _create_plugin() + ambient = plugin._provider.get_tracer("ambient").start_span("lambda-invocation") + token = otel_context.attach(trace.set_span_in_context(ambient)) + try: + plugin.on_invocation_start(_invocation_start_info()) + record = logging.LogRecord( + name="test", + level=logging.INFO, + pathname=__file__, + lineno=1, + msg="message", + args=(), + exc_info=None, + ) + + OtelContextLogFilter(plugin).filter(record) + + invocation_span = plugin._get_span(None) + assert invocation_span is not None + invocation_context = invocation_span.get_span_context() + assert record.traceId == format(invocation_context.trace_id, "032x") + assert record.spanId == format(invocation_context.span_id, "016x") + assert record.traceId != format(ambient.get_span_context().trace_id, "032x") + finally: + plugin.on_invocation_end(_invocation_end_info()) + otel_context.detach(token) + ambient.end() + + +def test_invocation_span_parents_to_same_trace_ambient_span(): + plugin, exporter = _create_plugin() + canonical_trace_id = _to_otel_trace_id(EXECUTION_ARN, START_TIME) + trace_state = TraceState([("vendor", "opaque")]) + ambient_context = SpanContext( + trace_id=canonical_trace_id, + span_id=int("1234567890abcdef", 16), + is_remote=False, + trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_state=trace_state, + ) + ambient = NonRecordingSpan(ambient_context) + token = otel_context.attach(trace.set_span_in_context(ambient, Context())) + try: + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + finally: + otel_context.detach(token) + + spans = {span.name: span for span in exporter.get_finished_spans()} + invocation = spans["Invocation"] + workflow = spans["Workflow"] + assert invocation.parent is not None + assert workflow.parent is not None + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.context.trace_state == trace_state + assert workflow.context.trace_state == trace_state + assert invocation.parent.span_id == ambient_context.span_id + assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) + + +def test_extracted_remote_parent_is_execution_ancestor(): + remote_trace_id = int("5759e988bd862e3fe1be46a994272793", 16) + remote_parent_id = int("53995c3f42cd8ad8", 16) + plugin, exporter = _create_plugin_with_sampler( + context_extractor=lambda _: ExtractedContext( + trace_id=remote_trace_id, + parent_span_id=remote_parent_id, + sampling=Sampling.SAMPLED, + ) + ) + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + + spans = {span.name: span for span in exporter.get_finished_spans()} + workflow = spans["Workflow"] + invocation = spans["Invocation"] + assert workflow.context.trace_id == remote_trace_id + assert invocation.context.trace_id == remote_trace_id + assert workflow.parent is not None + assert invocation.parent is not None + assert workflow.parent.span_id == remote_parent_id + assert invocation.parent.span_id == remote_parent_id + + +def test_backend_sampled_overrides_local_always_off_sampler(): + remote_trace_id = int("5759e988bd862e3fe1be46a994272793", 16) + remote_parent_id = int("53995c3f42cd8ad8", 16) + plugin, exporter = _create_plugin_with_sampler( + sampler=ALWAYS_OFF, + context_extractor=lambda _: ExtractedContext( + trace_id=remote_trace_id, + parent_span_id=remote_parent_id, + sampling=Sampling.SAMPLED, + ), + ) + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + + assert {span.name for span in exporter.get_finished_spans()} == { + "Invocation", + "Workflow", + } + + +def test_backend_not_sampled_overrides_local_always_on_sampler(): + remote_trace_id = int("5759e988bd862e3fe1be46a994272793", 16) + remote_parent_id = int("53995c3f42cd8ad8", 16) + plugin, exporter = _create_plugin_with_sampler( + sampler=ALWAYS_ON, + context_extractor=lambda _: ExtractedContext( + trace_id=remote_trace_id, + parent_span_id=remote_parent_id, + sampling=Sampling.NOT_SAMPLED, + ), + ) + + plugin.on_invocation_start(_invocation_start_info()) + plugin.on_invocation_end(_invocation_end_info()) + + assert exporter.get_finished_spans() == () def test_invocation_span_records_subsequent_invocation(): @@ -422,8 +637,8 @@ def test_operation_callbacks_emit_child_span_with_deterministic_span_id(): ) -def test_operation_end_without_start_omits_unobserved_previous_span_link(): - """A continuation cannot link to a SpanContext that was not checkpointed.""" +def test_operation_end_without_start_links_previous_logical_operation(): + """A continuation links to the deterministic logical operation context.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) operation_id = "wait-existing" @@ -451,8 +666,10 @@ def test_operation_end_without_start_omits_unobserved_previous_span_link(): assert span.name == "existing-wait" assert span.context.span_id == random_span_id linked_span_ids = {link.context.span_id for link in span.links} - assert linked_span_ids == {derive_workflow_span_id(EXECUTION_ARN)} - assert operation_id_to_span_id(EXECUTION_ARN, operation_id) not in linked_span_ids + assert linked_span_ids == { + derive_workflow_span_id(EXECUTION_ARN), + operation_id_to_span_id(EXECUTION_ARN, operation_id), + } assert ( span.attributes["durable.operation.status"] == OperationStatus.SUCCEEDED.value ) @@ -487,8 +704,86 @@ def test_continuation_span_uses_current_start_and_end_times(): assert before_callback <= span.start_time <= span.end_time <= after_callback -def test_retried_operation_uses_fresh_id_without_unobserved_previous_span_link(): - """Retried segments use fresh IDs without fabricating a prior context.""" +def test_resume_operation_timestamps_do_not_precede_current_invocation(): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + old_start_time = START_TIME + old_end_time = START_TIME + + plugin.on_operation_end( + OperationEndInfo( + operation_id="wait-resume", + operation_type=OperationType.WAIT, + sub_type=OperationSubType.WAIT, + name="otel-wait", + parent_id=None, + start_time=old_start_time, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=old_end_time, + error=None, + ) + ) + plugin.on_operation_start( + OperationStartInfo( + operation_id="after-resume", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="otel-after-resume", + parent_id=None, + start_time=old_start_time, + is_replayed=False, + status=OperationStatus.STARTED, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id="after-resume", + operation_type=OperationType.STEP, + sub_type=OperationSubType.STEP, + name="otel-after-resume", + parent_id=None, + start_time=old_start_time, + is_replayed=False, + status=OperationStatus.SUCCEEDED, + end_time=old_end_time, + error=None, + ) + ) + + spans = {span.name: span for span in exporter.get_finished_spans()} + wait_span = spans["otel-wait"] + after_resume_span = spans["otel-after-resume"] + assert invocation_span.start_time <= wait_span.start_time <= wait_span.end_time + assert wait_span.end_time <= after_resume_span.start_time + assert invocation_span.start_time <= after_resume_span.start_time + assert after_resume_span.parent is not None + assert after_resume_span.parent.span_id == invocation_span.context.span_id + + +def test_ordered_timestamps_are_thread_safe(): + plugin, _ = _create_plugin() + base_time = START_TIME + + with ThreadPoolExecutor(max_workers=8) as executor: + timestamps = list( + executor.map( + lambda _: plugin._next_ordered_timestamp(base_time), + range(100), + ) + ) + + assert len(set(timestamps)) == len(timestamps) + assert sorted(timestamps) == [ + int(base_time.timestamp() * 1_000_000_000) + index * 1_000 + for index in range(100) + ] + + +def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): + """Retried segments use fresh IDs and link the logical operation context.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) operation_id = "step-retried" @@ -528,8 +823,10 @@ def test_retried_operation_uses_fresh_id_without_unobserved_previous_span_link() assert span.name == "retried-step" assert span.context.span_id == random_span_id linked_span_ids = {link.context.span_id for link in span.links} - assert linked_span_ids == {derive_workflow_span_id(EXECUTION_ARN)} - assert operation_id_to_span_id(EXECUTION_ARN, operation_id) not in linked_span_ids + assert linked_span_ids == { + derive_workflow_span_id(EXECUTION_ARN), + operation_id_to_span_id(EXECUTION_ARN, operation_id), + } def test_step_operation_span_parents_attempt_span(): @@ -913,6 +1210,24 @@ def test_user_function_end_restores_enclosing_context(): assert plugin.get_current_span_context().span_id == invocation_span_id +def test_user_function_start_preserves_baggage_in_current_context(): + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + baggage_context = baggage.set_baggage( + "durable-test-key", "durable-test-value", otel_context.get_current() + ) + token = otel_context.attach(baggage_context) + try: + plugin.on_user_function_start(_user_function_start_info("step-baggage")) + + assert baggage.get_baggage("durable-test-key") == "durable-test-value" + + plugin.on_user_function_end(_user_function_end_info("step-baggage")) + finally: + otel_context.detach(token) + plugin.on_invocation_end(_invocation_end_info()) + + def test_user_function_end_restores_enclosing_context_on_failure(): """Verify the enclosing context is restored even when the step fails.""" plugin, _ = _create_plugin() @@ -1046,6 +1361,24 @@ def test_user_function_end_restores_parent_context_span_for_nested_step(): plugin.on_invocation_end(_invocation_end_info()) +def test_child_context_start_uses_invocation_time_not_durable_start_timestamp(): + plugin, _ = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + invocation_span = plugin._get_span(None) + assert invocation_span is not None + + context_id = "ctx-1" + plugin.on_user_function_start( + _user_function_start_info(context_id, operation_type=OperationType.CONTEXT) + ) + + context_span = plugin._get_span(context_id) + assert context_span is not None + assert context_span.start_time > invocation_span.start_time + + plugin.on_invocation_end(_invocation_end_info()) + + def test_top_level_step_end_falls_back_to_invocation_for_correlation(): """Verify a top-level step (parent_id=None) correlates to the invocation.""" plugin, _ = _create_plugin() @@ -1131,14 +1464,14 @@ def test_nested_steps_restore_context_span_across_multiple_iterations(): ], ) def test_workflow_span_exported_on_terminal(status, expected_code): - """A terminal invocation exports a deterministic Workflow root span.""" + """A terminal invocation exports a deterministic Workflow span.""" plugin, exporter = _create_plugin() plugin.on_invocation_start(_invocation_start_info()) plugin.on_invocation_end(_invocation_end_info(status)) workflow = next(s for s in exporter.get_finished_spans() if s.name == "Workflow") - # Root span: no parent. - assert workflow.parent is None + assert workflow.parent is not None + assert workflow.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) assert workflow.kind is SpanKind.INTERNAL # Deterministic span id derived from the execution ARN. assert workflow.context.span_id == derive_workflow_span_id(EXECUTION_ARN) @@ -1204,6 +1537,70 @@ def test_operation_span_links_to_workflow_span(): assert op_span.parent.span_id == spans_by_name["Invocation"].context.span_id +def test_replayed_context_span_links_previous_logical_operation(): + plugin, exporter = _create_plugin() + plugin.on_invocation_start(_invocation_start_info()) + operation_id = "callback-context" + random_span_id = int("fedcba9876543210", 16) + plugin._id_generator._fallback_id_generator.generate_span_id = lambda: ( + random_span_id + ) + + plugin.on_user_function_start( + UserFunctionStartInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.WAIT_FOR_CALLBACK, + name="wait-for-callback", + parent_id=None, + start_time=START_TIME, + is_replayed=True, + status=OperationStatus.STARTED, + is_replay_children=True, + attempt=2, + ) + ) + plugin.on_user_function_end( + UserFunctionEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.WAIT_FOR_CALLBACK, + name="wait-for-callback", + parent_id=None, + start_time=START_TIME, + is_replayed=True, + status=OperationStatus.STARTED, + is_replay_children=True, + attempt=2, + outcome=UserFunctionOutcome.INCOMPLETE, + end_time=None, + error=None, + ) + ) + plugin.on_operation_end( + OperationEndInfo( + operation_id=operation_id, + operation_type=OperationType.CONTEXT, + sub_type=OperationSubType.WAIT_FOR_CALLBACK, + name="wait-for-callback", + parent_id=None, + start_time=START_TIME, + is_replayed=True, + status=OperationStatus.SUCCEEDED, + end_time=END_TIME, + error=None, + ) + ) + + span = exporter.get_finished_spans()[0] + assert span.context.span_id == random_span_id + linked_span_ids = {link.context.span_id for link in span.links} + assert linked_span_ids == { + derive_workflow_span_id(EXECUTION_ARN), + operation_id_to_span_id(EXECUTION_ARN, operation_id), + } + + def test_workflow_span_name_is_configurable(): """The Workflow span name can be overridden via constructor kwarg.""" exporter = InMemorySpanExporter() @@ -1212,7 +1609,7 @@ def test_workflow_span_name_is_configurable(): plugin = InvocationOtelPlugin( OtelPluginConfig( tracer_provider=trace_provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, workflow_span_name="MyExecution", ) ) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py index b7f69152..a50a06ac 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_invocation_plugin_integration.py @@ -47,6 +47,7 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, + derive_execution_root_span_id, derive_workflow_span_id, operation_id_to_span_id, ) @@ -175,7 +176,7 @@ def _config_for_provider( monkeypatch.setattr(trace, "get_tracer_provider", lambda: provider) return OtelPluginConfig( tracer_provider=None if uses_global_provider else provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) @@ -229,7 +230,7 @@ def test_global_proxy_binds_sdk_provider_before_first_invocation( monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) plugin = InvocationOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -254,7 +255,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( monkeypatch.setattr(trace, "get_tracer_provider", lambda: current_provider[0]) plugin = InvocationOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -286,10 +287,14 @@ def _assert_hierarchy(exporter: InMemorySpanExporter) -> None: operation = spans[OP_NAME] attempt = spans[f"{OP_NAME} attempt 1"] - # Invocation span is a root (empty extracted context) and records status. - assert invocation.parent is None + # Invocation shares the synthetic execution root with Workflow and records + # status. + assert invocation.parent is not None + assert workflow.parent is not None assert invocation.kind is SpanKind.INTERNAL - assert invocation.context.trace_id != workflow.context.trace_id + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id + assert invocation.parent.span_id == derive_execution_root_span_id(EXECUTION_ARN) assert invocation.attributes is not None assert ( invocation.attributes["durable.invocation.status"] @@ -311,7 +316,7 @@ def _assert_hierarchy(exporter: InMemorySpanExporter) -> None: [False, True], ids=["explicit", "global"], ) -def test_invocation_span_parents_to_ambient_for_all_provider_modes( +def test_invocation_span_ignores_different_trace_ambient_for_all_provider_modes( uses_global_provider: bool, monkeypatch: pytest.MonkeyPatch ) -> None: provider, exporter = _provider() @@ -334,11 +339,13 @@ def test_invocation_span_parents_to_ambient_for_all_provider_modes( invocation = spans["Invocation"] operation = spans[OP_NAME] - assert workflow.parent is None + assert workflow.parent is not None assert workflow.context.trace_id != ambient.get_span_context().trace_id assert invocation.parent is not None - assert invocation.parent.span_id == ambient.get_span_context().span_id - assert invocation.context.trace_id == ambient.get_span_context().trace_id + assert invocation.parent.span_id != ambient.get_span_context().span_id + assert invocation.context.trace_id != ambient.get_span_context().trace_id + assert invocation.context.trace_id == workflow.context.trace_id + assert invocation.parent.span_id == workflow.parent.span_id assert operation.parent is not None assert operation.parent.span_id == invocation.context.span_id assert {link.context.span_id for link in operation.links} == { @@ -354,7 +361,7 @@ def test_community_layer_full_lifecycle_uses_supplied_provider(): plugin = InvocationOtelPlugin( OtelPluginConfig( tracer_provider=provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -376,7 +383,7 @@ def test_adot_layer_full_lifecycle_uses_global_provider(monkeypatch): plugin = InvocationOtelPlugin( OtelPluginConfig( - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) @@ -395,14 +402,14 @@ def test_second_plugin_uses_execution_trace_id_independent_of_xray(monkeypatch): first_plugin = InvocationOtelPlugin( OtelPluginConfig( tracer_provider=provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) target_plugin = InvocationOtelPlugin( OtelPluginConfig( tracer_provider=provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=False, ) ) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py index c7daa077..0cce516e 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_log_filter.py @@ -13,7 +13,6 @@ OperationType, UserFunctionStartInfo, ) -from opentelemetry.context import Context from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter @@ -40,7 +39,7 @@ def _create_plugin( plugin = InvocationOtelPlugin( OtelPluginConfig( tracer_provider=trace_provider, - context_extractor=lambda _: Context(), + context_extractor=lambda _: None, enrich_logger=enrich_logger, ) )