diff --git a/packages/aws-durable-execution-sdk-python-otel/README.md b/packages/aws-durable-execution-sdk-python-otel/README.md index 2a095bf9c..187971ba5 100644 --- a/packages/aws-durable-execution-sdk-python-otel/README.md +++ b/packages/aws-durable-execution-sdk-python-otel/README.md @@ -312,8 +312,8 @@ 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 +- otherwise a deterministic `DurableExecutionRoot` span derived from the + durable execution ARN, with `durable.execution.synthetic_root=true` `InvocationOtelPlugin` keeps durable operation spans under the Invocation span and links operations to Workflow: @@ -342,10 +342,46 @@ 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. +### Synthetic fallback anchor + +Each sampled fallback invocation creates and ends `DurableExecutionRoot` at +invocation start and attempts to flush it through the normal invocation-end +lifecycle, including when it returns `PENDING` or `RETRY`. A successfully +delivered copy supplies ancestry if the execution is later stopped or times out +without resuming. Export and flush failures, or termination before flushing, +do not provide a delivery guarantee. + +The anchor uses the existing execution-scoped span ID and the stable execution +start timestamp for both its start and end. It carries no terminal outcome; +`Workflow` remains the span for execution duration and success/failure. Retries +of the first invocation and later invocations may re-export the same anchor to +recover lost telemetry; export is not exactly once. Two executions sharing a +propagated trace ID retain distinct ARN-scoped anchor IDs. +The anchor is intentionally zero-duration, so its descendants end after it; +trace viewers may flag that timing relationship even though the ancestry is valid. + +The anchor uses the same configured tracer, provider resource and processors as +other durable spans, but carries only the resolved sampling decision. Its SDK +attributes are `durable.execution.arn` and `durable.execution.synthetic_root`; +sampler-added attributes and tracestate are omitted from this repeated identity. +Ordinary durable spans retain the complete resolved sampler metadata. SDK-added +IDs, timestamps and anchor attributes stay stable across recovery exports; +provider resources and custom processor enrichment may vary. The SDK does not +replace configured service identity or freeze detector fields. + +Complete remote parents are externally owned and are never materialized by the +SDK. Unsampled invocations emit no anchor. Sampling remains resolved per +invocation: a later sampled fallback invocation also emits the anchor even if +an earlier invocation was unsampled. Missing execution start time continues to +disable telemetry; no wall-clock timestamp is invented for the anchor. + ### Sampling Sampling is resolved once per invocation and carried to every durable span in -that invocation. Precedence is: +that invocation. Each span consumes its private sampling intent before processors +run. Processor-created spans use their own configured sampling policy and normal +parent-context semantics, including when a processor forwards the supplied parent +context. Precedence is: 1. `Sampled=1` or `Sampled=0` from `_X_AMZN_TRACE_ID` 2. a same-trace ambient span's recording/sampled state @@ -446,6 +482,11 @@ them to Invocation. A custom OpenTelemetry `IdGenerator` that produces reproducible trace and span IDs from execution metadata. Exported for advanced use cases. +`use_ids` retains its existing scope-wide trace override by default. The +synthetic-root implementation uses a separate internal one-span override, +preventing processor-created roots from inheriting that internal ID. Existing +`use_ids` signatures and default behavior remain unchanged. + ### `xray_context_extractor` Default context extractor. Reads the `_X_AMZN_TRACE_ID` environment variable and 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 c3d7afc17..ff3999b9c 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 @@ -21,6 +21,8 @@ class _IdOverride: trace_id: int | None span_id: int | None + trace_id_consumed: bool = False + consume_trace_id: bool = False def _to_otel_trace_id(execution_arn: str, start_timestamp: datetime) -> int: @@ -83,9 +85,10 @@ def derive_workflow_span_id(durable_execution_arn: str) -> int: 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 + The synthetic root anchors the trace when the backend does not provide a + complete remote parent. Its zero-duration span is exported on each sampled + fallback invocation. Its ID is stable across reinvocations and uses a + namespace distinct from Workflow and operation span IDs. """ if not durable_execution_arn: @@ -143,28 +146,61 @@ def use_ids(self, *, trace_id: int | None, span_id: int | None) -> Iterator[None finally: self._id_override.reset(token) + @contextmanager + def _use_ids_for_span( + self, + *, + trace_id: int, + span_id: int, + ) -> Iterator[None]: + """Reserve IDs for one SDK-created span without changing public scope semantics.""" + token = self._id_override.set( + _IdOverride(trace_id, span_id, consume_trace_id=True) + ) + try: + yield + finally: + self._id_override.reset(token) + def generate_trace_id(self) -> int: """Generate a 128-bit trace ID.""" override = self._id_override.get() if override is not None and override.trace_id is not None: + # Consume the identity once, but keep its non-random classification + # until the SDK asks for the span ID below. + if override.consume_trace_id: + self._id_override.set(_IdOverride(None, override.span_id, True, True)) return override.trace_id + if override is not None and override.trace_id_consumed: + self._id_override.set( + _IdOverride( + None, override.span_id, consume_trace_id=override.consume_trace_id + ) + ) return self._fallback_id_generator.generate_trace_id() def generate_span_id(self) -> int: """Generate a 64-bit span ID.""" override = self._id_override.get() - if override is not None and override.span_id is not None: - span_id = override.span_id - # Consume before returning so a re-entrant call in the same span - # creation falls back instead of reusing the deterministic ID. - self._id_override.set(_IdOverride(override.trace_id, None)) - return span_id + if override is not None: + # Span ID generation completes the ID phase of SDK start_span, before + # processors run. Release the consumed trace classification too, so + # a processor creating another root sees only its own generator. + self._id_override.set( + _IdOverride( + override.trace_id, None, consume_trace_id=override.consume_trace_id + ) + ) + if override.span_id is not None: + return override.span_id return self._fallback_id_generator.generate_span_id() def is_trace_id_random(self) -> bool: """Report whether the current trace ID is randomly generated.""" override = self._id_override.get() - if override is not None and override.trace_id is not None: + if override is not None and ( + override.trace_id is not None or override.trace_id_consumed + ): return False fallback_method = getattr( self._fallback_id_generator, "is_trace_id_random", None 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 index 9d2fa6103..93256bed2 100644 --- 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 @@ -5,6 +5,7 @@ import functools import inspect from dataclasses import dataclass +from threading import Lock from typing import Any from opentelemetry import context as otel_context @@ -31,6 +32,20 @@ class DurableSamplingIntent: result: SamplingResult +class _SamplingIntentDelivery: + """Give one sampler call the intent, even if its context is forwarded.""" + + def __init__(self, intent: DurableSamplingIntent) -> None: + self._intent: DurableSamplingIntent | None = intent + self._lock = Lock() + + def consume(self) -> DurableSamplingIntent | None: + with self._lock: + intent = self._intent + self._intent = None + return intent + + class DurableSampler(Sampler): """Sampler that honors a durable sampling intent carried on parent context.""" @@ -56,8 +71,13 @@ def should_sample( links: Any = None, trace_state: Any = None, ) -> SamplingResult: - intent = otel_context.get_value(_DURABLE_SAMPLING_INTENT_KEY, parent_context) - if isinstance(intent, DurableSamplingIntent): + delivery = otel_context.get_value(_DURABLE_SAMPLING_INTENT_KEY, parent_context) + intent = ( + delivery.consume() + if isinstance(delivery, _SamplingIntentDelivery) + else None + ) + if intent is not None: merged_attributes = dict(attributes or {}) merged_attributes.update(dict(intent.result.attributes or {})) return SamplingResult( @@ -84,10 +104,17 @@ def store_sampling_intent( parent_context: Context, intent: DurableSamplingIntent | None, ) -> Context: - """Attach a durable sampling intent to a span parent context.""" + """Attach a fresh single-use delivery of the invocation's sampling result. + + The sampler consumes it before processors run. A processor forwarding this + context therefore uses its own sampling policy; another durable span gets a + new delivery without changing the reusable invocation-level intent. + """ if intent is None: return parent_context - return otel_context.set_value(_DURABLE_SAMPLING_INTENT_KEY, intent, parent_context) + return otel_context.set_value( + _DURABLE_SAMPLING_INTENT_KEY, _SamplingIntentDelivery(intent), parent_context + ) def resolve_sampling_result( 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 40a62ca84..2008615e1 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 @@ -86,6 +86,7 @@ DurableParentSpan, ensure_end_after_start, ) +from aws_durable_execution_sdk_python_otel.execution_root import ExecutionRoot from aws_durable_execution_sdk_python_otel.execution_trace_context import ( ExecutionTraceContext, canonical_trace_id, @@ -96,6 +97,7 @@ uninstall_log_filter, ) from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider +from aws_durable_execution_sdk_python_otel.timestamps import to_otel_timestamp logger = logging.getLogger(__name__) @@ -115,7 +117,7 @@ def _to_otel_timestamp(dt: datetime.datetime | None) -> int | None: """Convert a datetime to an OTel timestamp (ns since epoch), or None.""" if dt is None: return None - return int(dt.timestamp() * 1_000_000_000) + return to_otel_timestamp(dt) class ExecutionOtelPlugin(DurableInstrumentationPlugin): @@ -489,6 +491,14 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: root_sampled=lambda: is_sampled(sampling_result), ) + # Materialize the fallback anchor before user code can suspend or fail. + # Repeated invocations retain the same SDK-owned identity and timing. + ExecutionRoot( + execution_arn=self._execution_arn, + ancestor=self._execution_trace_context.execution_ancestor, + start_time=info.execution_start_time, + ).export(self._tracer, self._id_generator, self._sampling_intent) + self._start_workflow_span(info) # Keep the invocation on the shared execution trace. self._start_invocation_span(info) diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_root.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_root.py new file mode 100644 index 000000000..fb43f68a8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/execution_root.py @@ -0,0 +1,68 @@ +"""Materialize the stable SDK-owned ancestor of a fallback execution trace.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime + +from opentelemetry.context import Context +from opentelemetry.sdk.trace.sampling import SamplingResult +from opentelemetry.trace import SpanContext, SpanKind, Tracer + +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + DeterministicIdGenerator, +) +from aws_durable_execution_sdk_python_otel.durable_sampling import ( + DurableSamplingIntent, + store_sampling_intent, +) +from aws_durable_execution_sdk_python_otel.timestamps import to_otel_timestamp + + +@dataclass(frozen=True) +class ExecutionRoot: + """A zero-duration anchor, reproducible in any invocation of an execution. + + Re-exporting the same anchor permits recovery after a failed flush without + checkpointing telemetry state. The anchor carries stable SDK attributes; + configured provider resources and processor enrichment remain unchanged. + """ + + execution_arn: str + ancestor: SpanContext + start_time: datetime + + def export( + self, + tracer: Tracer, + id_generator: DeterministicIdGenerator, + sampling_intent: DurableSamplingIntent, + ) -> None: + """Export a sampled local ancestor; never replace an external parent.""" + if self.ancestor.is_remote or not self.ancestor.trace_flags.sampled: + return + + # Reuse only the resolved decision without resampling. This same-ID + # anchor must not acquire variable sampler attributes or tracestate. + # The regular tracer preserves configured resources and processors. + root_attributes: dict[str, str | bool] = { + "durable.execution.arn": self.execution_arn, + "durable.execution.synthetic_root": True, + } + root_intent = DurableSamplingIntent( + SamplingResult(sampling_intent.result.decision) + ) + root_context = store_sampling_intent(Context(), root_intent) + timestamp = to_otel_timestamp(self.start_time) + with id_generator._use_ids_for_span( + trace_id=self.ancestor.trace_id, + span_id=self.ancestor.span_id, + ): + span = tracer.start_span( + "DurableExecutionRoot", + context=root_context, + kind=SpanKind.INTERNAL, + attributes=root_attributes, + start_time=timestamp, + ) + span.end(end_time=timestamp) 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 0fb974fe1..4fc2bf02b 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 @@ -55,6 +55,7 @@ from aws_durable_execution_sdk_python_otel.durable_parent_span import ( DurableParentSpan, ) +from aws_durable_execution_sdk_python_otel.execution_root import ExecutionRoot from aws_durable_execution_sdk_python_otel.execution_trace_context import ( ExecutionTraceContext, canonical_trace_id, @@ -65,6 +66,7 @@ ) from aws_durable_execution_sdk_python_otel.otel_plugin_config import OtelPluginConfig from aws_durable_execution_sdk_python_otel.provider import create_tracer_provider +from aws_durable_execution_sdk_python_otel.timestamps import to_otel_timestamp logger = logging.getLogger(__name__) @@ -82,7 +84,7 @@ 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) + return to_otel_timestamp(dt) class InvocationOtelPlugin(DurableInstrumentationPlugin): @@ -582,6 +584,14 @@ def on_invocation_start(self, info: InvocationStartInfo) -> None: root_sampled=lambda: is_sampled(sampling_result), ) + # Materialize the fallback anchor before user code can suspend or fail. + # Repeated invocations retain the same SDK-owned identity and timing. + ExecutionRoot( + execution_arn=self._execution_arn, + ancestor=self._execution_trace_context.execution_ancestor, + start_time=info.execution_start_time, + ).export(self._tracer, self._id_generator, self._sampling_intent) + self._start_workflow_span(info) self._start_span( diff --git a/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/timestamps.py b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/timestamps.py new file mode 100644 index 000000000..d8cb716ad --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/timestamps.py @@ -0,0 +1,8 @@ +"""OpenTelemetry timestamp conversion shared by span producers.""" + +from datetime import datetime + + +def to_otel_timestamp(value: datetime) -> int: + """Convert a datetime to nanoseconds using the SDK's existing precision.""" + return int(value.timestamp() * 1_000_000_000) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_fallback_root_lifecycle.py b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_fallback_root_lifecycle.py new file mode 100644 index 000000000..a1b39fec3 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/e2e/test_fallback_root_lifecycle.py @@ -0,0 +1,225 @@ +"""Fallback anchor coverage through the decorator and local durable runner.""" + +import inspect +from typing import Any + +import pytest +from aws_durable_execution_sdk_python import DurableContext +from aws_durable_execution_sdk_python.config import Duration +from aws_durable_execution_sdk_python.execution import durable_execution +from aws_durable_execution_sdk_python.plugin import ( + DurableInstrumentationPlugin, + InvocationEndInfo, + InvocationStatus, +) +from aws_durable_execution_sdk_python_testing.runner import DurableFunctionTestRunner +from opentelemetry import context +from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult +from opentelemetry.trace import TraceState + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, +) + + +class _CaptureInvocationEnd(DurableInstrumentationPlugin): + def __init__(self, exporter: InMemorySpanExporter) -> None: + self.exporter = exporter + self.snapshots: list[tuple[InvocationStatus, tuple[ReadableSpan, ...]]] = [] + + def on_invocation_end(self, info: InvocationEndInfo) -> None: + self.snapshots.append((info.status, tuple(self.exporter.get_finished_spans()))) + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize("outcome", ["success", "failure", "timeout"]) +def test_fallback_anchor_survives_suspension( + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + outcome: str, +) -> None: + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor( + BatchSpanProcessor(exporter, schedule_delay_millis=60000) + ) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + observer = _CaptureInvocationEnd(exporter) + completed_steps: list[str] = [] + before = context.get_current() + + def step(_step_context: Any) -> str: + completed_steps.append("step") + return "stored" + + def handler(_event: Any, durable: DurableContext) -> str: + value = durable.step(step, name="completed-step") + durable.wait( + Duration.from_seconds(60 if outcome == "timeout" else 1), name="wait" + ) + if outcome == "failure": + raise ValueError("failed after resume") + return value + + wrapped = durable_execution(handler, plugins=[plugin, observer]) + try: + # Published runners use real time; newer workspace runners default to + # skipping durable waits. Exercise timeout without advancing that wait. + runner_options: dict[str, Any] = {"handler": wrapped} + if "skip_time" in inspect.signature(DurableFunctionTestRunner).parameters: + runner_options["skip_time"] = outcome != "timeout" + with DurableFunctionTestRunner(**runner_options) as runner: + arn = runner.run_async( + input="{}", timeout=3 if outcome == "timeout" else 15 + ) + result = runner.wait_for_result(arn, timeout=10) + assert completed_steps == ["step"] + assert observer.snapshots[0][0] is InvocationStatus.PENDING + first_spans = observer.snapshots[0][1] + first_roots = [s for s in first_spans if s.name == "DurableExecutionRoot"] + assert len(first_roots) == 1 + assert not any(s.name == "Workflow" for s in first_spans) + roots = [ + s for s in exporter.get_finished_spans() if s.name == "DurableExecutionRoot" + ] + assert all(s.to_json() == first_roots[0].to_json() for s in roots) + if outcome == "timeout": + assert result.status.value == "FAILED" + assert result.error is not None + assert "timed out" in (result.error.message or "") + assert len(observer.snapshots) == 1 + assert len(roots) == 1 + assert not any(s.name == "Workflow" for s in exporter.get_finished_spans()) + else: + assert result.status.value == ( + "FAILED" if outcome == "failure" else "SUCCEEDED" + ) + workflows = [ + s for s in exporter.get_finished_spans() if s.name == "Workflow" + ] + assert len(workflows) == 1 + assert workflows[0].parent is not None + assert first_roots[0].context is not None + assert workflows[0].parent.span_id == first_roots[0].context.span_id + assert context.get_current() == before + finally: + provider.shutdown() + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +def test_public_resume_keeps_anchor_stable_and_processor_sampling_independent( + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], +) -> None: + class Policy(Sampler): + def __init__(self) -> None: + self.resolutions: list[str] = [] + self.processor_calls: list[str] = [] + + def should_sample( + self, + parent_context: Any, + trace_id: int, + name: str, + kind: Any = None, + attributes: Any = None, + links: Any = None, + trace_state: Any = None, + ) -> SamplingResult: + if name == "processor-probe": + self.processor_calls.append(name) + return SamplingResult(Decision.DROP) + self.resolutions.append(name) + number = len(self.resolutions) + return SamplingResult( + Decision.RECORD_AND_SAMPLE, + attributes={"sampler.invocation": number}, + trace_state=TraceState([("vendor", str(number))]), + ) + + def get_description(self) -> str: + return "public lifecycle sampling policy" + + policy = Policy() + exporter = InMemorySpanExporter() + provider = TracerProvider(sampler=policy) + provider.add_span_processor( + BatchSpanProcessor(exporter, schedule_delay_millis=60000) + ) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: None, + enrich_logger=False, + ) + ) + observed_probes: list[bool] = [] + + class Processor(SpanProcessor): + def on_start(self, span: Span, parent_context: Any = None) -> None: + if span.name == "processor-probe": + return + with plugin._tracer.start_as_current_span( + "processor-probe", context=parent_context + ) as probe: + observed_probes.append(probe.is_recording()) + + provider.add_span_processor(Processor()) + observer = _CaptureInvocationEnd(exporter) + completed_steps: list[str] = [] + before = context.get_current() + + def step(_step_context: Any) -> str: + completed_steps.append("once") + return "stored" + + def handler(_event: Any, durable: DurableContext) -> str: + result = durable.step(step, name="completed-step") + durable.wait(Duration.from_seconds(1), name="resume") + return result + + wrapped = durable_execution(handler, plugins=[plugin, observer]) + try: + with DurableFunctionTestRunner(handler=wrapped) as runner: + arn = runner.run_async(input="{}", timeout=15) + result = runner.wait_for_result(arn, timeout=10) + assert result.status.value == "SUCCEEDED" + assert completed_steps == ["once"] + assert [status for status, _ in observer.snapshots] == [ + InvocationStatus.PENDING, + InvocationStatus.SUCCEEDED, + ] + spans = exporter.get_finished_spans() + roots = [span for span in spans if span.name == "DurableExecutionRoot"] + invocations = [span for span in spans if span.name == "Invocation"] + assert len(roots) == len(invocations) == len(policy.resolutions) == 2 + assert roots[0].to_json() == roots[1].to_json() + assert roots[0].attributes == { + "durable.execution.arn": arn, + "durable.execution.synthetic_root": True, + } + assert ( + roots[0].context is not None + and roots[0].context.trace_state == TraceState() + ) + for number, invocation in enumerate(invocations, start=1): + assert invocation.attributes is not None and invocation.context is not None + assert invocation.attributes["sampler.invocation"] == number + assert invocation.context.trace_state == TraceState( + [("vendor", str(number))] + ) + assert observed_probes and not any(observed_probes) + assert len(policy.processor_calls) == len(observed_probes) + assert not [span for span in spans if span.name == "processor-probe"] + assert context.get_current() == before + finally: + provider.shutdown() diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_deterministic_id_generator.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_deterministic_id_generator.py index 6bcd320ab..1386566a2 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_deterministic_id_generator.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_deterministic_id_generator.py @@ -339,3 +339,57 @@ def test_derive_execution_root_span_id_differs_from_workflow_span_id(): assert derive_execution_root_span_id(_ROOT_ARN) != derive_workflow_span_id( _ROOT_ARN ) + + +def test_trace_override_is_consumed_once_and_randomness_is_released() -> None: + fallback_trace = int("a" * 32, 16) + generator = DeterministicIdGenerator( + _StubIdGenerator( + trace_id=fallback_trace, + span_id=int("b" * 16, 16), + trace_id_is_random=True, + ) + ) + with generator._use_ids_for_span(trace_id=123, span_id=456): + assert generator.generate_trace_id() == 123 + assert generator.is_trace_id_random() is False + assert generator.generate_span_id() == 456 + # The SDK calls processors after generating this span's identity. + assert generator.is_trace_id_random() is True + assert generator.generate_trace_id() == fallback_trace + assert generator.generate_span_id() == int("b" * 16, 16) + + +def test_public_trace_only_scope_preserves_existing_override() -> None: + generator = DeterministicIdGenerator( + _StubIdGenerator( + trace_id=123, + span_id=456, + trace_id_is_random=True, + ) + ) + with generator.use_ids(trace_id=789, span_id=None): + assert generator.generate_trace_id() == 789 + assert generator.is_trace_id_random() is False + assert generator.generate_span_id() == 456 + assert generator.is_trace_id_random() is False + assert generator.generate_trace_id() == 789 + + +def test_default_use_ids_preserves_scope_wide_trace_override() -> None: + """Existing advanced callers can intentionally create multiple roots on one trace.""" + generator = DeterministicIdGenerator( + _StubIdGenerator( + trace_id=123, + span_id=456, + trace_id_is_random=True, + ) + ) + with generator.use_ids(trace_id=789, span_id=999): + assert generator.generate_trace_id() == 789 + assert generator.generate_span_id() == 999 + assert generator.generate_trace_id() == 789 + assert generator.generate_span_id() == 456 + assert generator.is_trace_id_random() is False + assert generator.generate_trace_id() == 123 + assert generator.is_trace_id_random() is True diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_durable_sampling.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_durable_sampling.py index 792d611f0..c1fd1cd98 100644 --- a/packages/aws-durable-execution-sdk-python-otel/tests/test_durable_sampling.py +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_durable_sampling.py @@ -4,10 +4,13 @@ import gc import weakref +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier from typing import Any, cast import pytest from opentelemetry.context import Context +from opentelemetry import context as otel_context from opentelemetry.sdk.trace.sampling import ( ALWAYS_OFF, ALWAYS_ON, @@ -491,6 +494,52 @@ def test_store_sampling_intent_returns_context_unchanged_when_none() -> None: assert store_sampling_intent(parent_context, None) is parent_context +@pytest.mark.parametrize("decision", list(Decision)) +def test_sampling_intent_delivery_is_consumed_but_template_is_reusable( + decision: Decision, +) -> None: + delegate = ALWAYS_OFF if decision is Decision.RECORD_AND_SAMPLE else ALWAYS_ON + sampler = DurableSampler(delegate) + intent = DurableSamplingIntent( + SamplingResult( + decision, + attributes={"sampler.metadata": "retained"}, + trace_state=TraceState([("vendor", "original")]), + ) + ) + parent = store_sampling_intent(Context(), intent) + first = sampler.should_sample(parent, TRACE_ID, "durable") + forwarded = sampler.should_sample(parent, TRACE_ID, "processor") + next_durable = sampler.should_sample( + store_sampling_intent(Context(), intent), TRACE_ID, "next-durable" + ) + assert first.decision is next_durable.decision is decision + assert first.attributes == next_durable.attributes == intent.result.attributes + assert first.trace_state == next_durable.trace_state == intent.result.trace_state + assert forwarded.decision is ( + Decision.DROP + if decision is Decision.RECORD_AND_SAMPLE + else Decision.RECORD_AND_SAMPLE + ) + + +def test_forwarded_contexts_share_one_atomic_sampling_delivery() -> None: + sampler = DurableSampler(ALWAYS_OFF) + intent = DurableSamplingIntent(SamplingResult(Decision.RECORD_AND_SAMPLE)) + parent = store_sampling_intent(Context(), intent) + ready = Barrier(2) + + def sample(index: int) -> Decision: + forwarded = otel_context.set_value("processor", index, parent) + ready.wait(timeout=5) + return sampler.should_sample(forwarded, TRACE_ID, "concurrent").decision + + with ThreadPoolExecutor(max_workers=2) as pool: + decisions = list(pool.map(sample, range(2))) + assert decisions.count(Decision.RECORD_AND_SAMPLE) == 1 + assert decisions.count(Decision.DROP) == 1 + + def test_is_sampled_matches_record_and_sample_only() -> None: assert is_sampled(SamplingResult(Decision.RECORD_AND_SAMPLE)) is True assert is_sampled(SamplingResult(Decision.RECORD_ONLY)) is False 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 214a3a89d..a3d79c8aa 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 @@ -551,7 +551,7 @@ def test_context_span_waits_for_terminal_operation_status( # Still a placeholder, still not exported, until the terminal operation end. assert plugin._get_span(operation_id) is active_span assert not active_span.is_recording() - assert not exporter.get_finished_spans() + assert {s.name for s in exporter.get_finished_spans()} == {"DurableExecutionRoot"} plugin.on_operation_end( OperationEndInfo( @@ -568,7 +568,9 @@ def test_context_span_waits_for_terminal_operation_status( ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert span.attributes["durable.operation.status"] == terminal_status.value assert span.status.status_code is expected_span_status @@ -616,7 +618,9 @@ def test_step_attempt_span_omits_operation_status(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert ( span.attributes["durable.attempt.outcome"] == UserFunctionOutcome.SUCCEEDED.value @@ -1551,7 +1555,9 @@ def suspend_on_worker() -> tuple[int, bool]: assert span_still_current is False assert span_key not in plugin._context_tokens assert suspended_span is not None - assert not exporter.get_finished_spans() + assert {s.name for s in exporter.get_finished_spans()} == { + "DurableExecutionRoot" + } # The timed resume lands on this thread, with nothing stale to unwind. plugin.on_user_function_start(_step_start_info("step-1")) @@ -1597,7 +1603,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): assert otel_context.get_current() == before_context assert set(plugin._context_tokens) == {"__invocation_context__"} # Neither span is ended: both operations are still in flight. - assert not exporter.get_finished_spans() + assert {s.name for s in exporter.get_finished_spans()} == {"DurableExecutionRoot"} # The timed in-process resume replays both contexts, outer first. plugin.on_user_function_start(_context_start_info("ctx-outer")) 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 df3e4b650..03522719d 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 @@ -265,6 +265,7 @@ def test_global_proxy_binds_sdk_provider_before_first_invocation( plugin.on_invocation_end(_invocation_end()) assert {span.name for span in exporter.get_finished_spans()} == { + "DurableExecutionRoot", "Invocation", "Workflow", } @@ -289,6 +290,7 @@ def test_parent_placeholder_supports_adot_style_parent_inspection() -> None: plugin.on_invocation_end(_invocation_end()) assert {span.name for span in exporter.get_finished_spans()} == { + "DurableExecutionRoot", "Invocation", "Workflow", OP_NAME, @@ -323,6 +325,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) assert {span.name for span in exporter.get_finished_spans()} == { + "DurableExecutionRoot", "Invocation", "Workflow", OP_NAME, diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_root.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_root.py new file mode 100644 index 000000000..9996bb2b8 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_root.py @@ -0,0 +1,584 @@ +"""Fallback ancestry regressions against real OTel providers and exporters.""" + +from collections.abc import Iterator +from datetime import UTC, datetime, timedelta +from typing import Any + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationEndInfo, + InvocationStartInfo, + InvocationStatus, +) +from opentelemetry import context +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import Span, SpanProcessor, TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.sampling import ( + ALWAYS_OFF, + ALWAYS_ON, + Decision, + Sampler, + SamplingResult, +) +from opentelemetry.trace import SpanKind, StatusCode, TraceState + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, + ExtractedContext, + Sampling, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + derive_execution_root_span_id, +) + + +START = datetime(2026, 10, 1, tzinfo=UTC) +ARN = "arn:aws:lambda:us-west-2:123456789012:function:root:1/durable-execution/test/id" +TRACE = 0x12345678901234567890123456789012 +PluginType = type[ExecutionOtelPlugin] | type[InvocationOtelPlugin] + + +@pytest.fixture(params=[ExecutionOtelPlugin, InvocationOtelPlugin]) +def plugin_type(request: pytest.FixtureRequest) -> PluginType: + return request.param + + +@pytest.fixture(autouse=True) +def balanced_context() -> Iterator[None]: + before = context.get_current() + yield + assert context.get_current() == before + + +def create_plugin( + plugin_type: PluginType, + extracted: ExtractedContext | None = None, + sampler: Sampler = ALWAYS_ON, + resource: Resource | None = None, + batch: bool = False, +) -> tuple[ + ExecutionOtelPlugin | InvocationOtelPlugin, TracerProvider, InMemorySpanExporter +]: + provider = TracerProvider(sampler=sampler, resource=resource) + exporter = InMemorySpanExporter() + processor = ( + BatchSpanProcessor(exporter, schedule_delay_millis=60000) + if batch + else SimpleSpanProcessor(exporter) + ) + provider.add_span_processor(processor) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=provider, + context_extractor=lambda _: extracted, + enrich_logger=False, + ) + ) + return plugin, provider, exporter + + +def start_info( + arn: str = ARN, + first: bool = True, + request: str = "request", + start_time: datetime = START, +) -> InvocationStartInfo: + return InvocationStartInfo( + request_id=request, + execution_arn=arn, + execution_start_time=start_time, + is_first_invocation=first, + ) + + +def end_info( + status: InvocationStatus, + arn: str = ARN, + first: bool = True, + request: str = "request", + start_time: datetime = START, +) -> InvocationEndInfo: + return InvocationEndInfo( + request_id=request, + execution_arn=arn, + execution_start_time=start_time, + is_first_invocation=first, + status=status, + ) + + +@pytest.mark.parametrize( + "extracted", + [ + None, + ExtractedContext(0, 0), + ExtractedContext(2**128, 2**64), + ExtractedContext(TRACE, None), + ExtractedContext(TRACE, 0), + ], +) +@pytest.mark.parametrize("status", list(InvocationStatus)) +def test_sampled_anchor_available_before_first_invocation_returns( + plugin_type: PluginType, + extracted: ExtractedContext | None, + status: InvocationStatus, +) -> None: + plugin, provider, exporter = create_plugin(plugin_type, extracted) + try: + plugin.on_invocation_start(start_info()) + roots = exporter.get_finished_spans() + assert len(roots) == 1 + root = roots[0] + assert root.name == "DurableExecutionRoot" + assert root.context is not None + assert root.context.span_id == derive_execution_root_span_id(ARN) + assert root.parent is None + assert ( + root.start_time == root.end_time == int(START.timestamp() * 1_000_000_000) + ) + assert root.status.status_code is StatusCode.UNSET + assert root.kind is SpanKind.INTERNAL + assert root.attributes == { + "durable.execution.arn": ARN, + "durable.execution.synthetic_root": True, + } + plugin.on_invocation_end(end_info(status)) + spans = exporter.get_finished_spans() + assert len([s for s in spans if s.name == "DurableExecutionRoot"]) == 1 + for span in spans: + if span.name in ("Invocation", "Workflow"): + assert span.parent is not None + assert span.parent.span_id == root.context.span_id + assert span.context is not None + assert span.context.trace_id == root.context.trace_id + finally: + provider.shutdown() + + +@pytest.mark.parametrize("status", [InvocationStatus.PENDING, InvocationStatus.RETRY]) +def test_nonterminal_end_flushes_anchor( + plugin_type: PluginType, status: InvocationStatus +) -> None: + plugin, provider, exporter = create_plugin(plugin_type, batch=True) + try: + plugin.on_invocation_start(start_info()) + plugin.on_invocation_end(end_info(status)) + assert {s.name for s in exporter.get_finished_spans()} == { + "DurableExecutionRoot", + "Invocation", + } + finally: + provider.shutdown() + + +def test_retry_and_resume_keep_sdk_fields_and_provider_resources( + plugin_type: PluginType, +) -> None: + roots = [] + for index, status in enumerate( + [InvocationStatus.RETRY, InvocationStatus.PENDING, InvocationStatus.FAILED] + ): + resource = Resource( + { + "service.name": "function", + "faas.instance": f"sandbox-{index}", + "process.pid": index, + } + ) + plugin, provider, exporter = create_plugin(plugin_type, resource=resource) + try: + # A first-invocation retry is still first; a later terminal resume is not. + plugin.on_invocation_start( + start_info(first=index < 2, request=f"request-{index}") + ) + plugin.on_invocation_end( + end_info(status, first=index < 2, request=f"request-{index}") + ) + spans = exporter.get_finished_spans() + root = next(s for s in spans if s.name == "DurableExecutionRoot") + assert root.attributes is not None + assert root.context is not None + roots.append( + ( + root.context.trace_id, + root.context.span_id, + root.parent, + root.start_time, + root.end_time, + dict(root.attributes), + root.status.status_code, + root.kind, + root.events, + root.links, + ) + ) + invocation = next(s for s in spans if s.name == "Invocation") + assert invocation.resource == provider.resource + assert invocation.resource.attributes["faas.instance"] == f"sandbox-{index}" + assert root.resource == provider.resource + assert root.resource.attributes["service.name"] == "function" + assert root.resource.attributes["faas.instance"] == f"sandbox-{index}" + finally: + provider.shutdown() + assert roots[0] == roots[1] == roots[2] + + +@pytest.mark.parametrize( + ("extracted", "sampler", "expected"), + [ + (ExtractedContext(TRACE, None, Sampling.NOT_SAMPLED), ALWAYS_ON, False), + (ExtractedContext(TRACE, None, Sampling.SAMPLED), ALWAYS_OFF, True), + (None, ALWAYS_OFF, False), + (None, ALWAYS_ON, True), + (ExtractedContext(TRACE, 42, Sampling.SAMPLED), ALWAYS_ON, False), + ], +) +def test_sampling_and_external_parent_ownership( + plugin_type: PluginType, + extracted: ExtractedContext | None, + sampler: Sampler, + expected: bool, +) -> None: + plugin, provider, exporter = create_plugin(plugin_type, extracted, sampler) + try: + plugin.on_invocation_start(start_info()) + plugin.on_invocation_end(end_info(InvocationStatus.SUCCEEDED)) + spans = exporter.get_finished_spans() + assert bool([s for s in spans if s.name == "DurableExecutionRoot"]) is expected + if extracted is not None and extracted.has_complete_remote_parent: + assert all(s.parent is not None and s.parent.span_id == 42 for s in spans) + elif not expected: + assert not spans + finally: + provider.shutdown() + + +def test_shared_trace_retains_execution_scoped_root_identity( + plugin_type: PluginType, +) -> None: + plugin, provider, exporter = create_plugin( + plugin_type, ExtractedContext(TRACE, None, Sampling.SAMPLED) + ) + try: + for arn in [ARN, ARN + "-other"]: + plugin.on_invocation_start(start_info(arn=arn)) + plugin.on_invocation_end(end_info(InvocationStatus.SUCCEEDED, arn=arn)) + roots = [ + s for s in exporter.get_finished_spans() if s.name == "DurableExecutionRoot" + ] + assert len(roots) == 2 + assert all(s.context is not None and s.context.trace_id == TRACE for s in roots) + assert {s.context.span_id for s in roots if s.context is not None} == { + derive_execution_root_span_id(ARN), + derive_execution_root_span_id(ARN + "-other"), + } + finally: + provider.shutdown() + + +def test_later_sampled_invocation_recovers_unsampled_first( + plugin_type: PluginType, +) -> None: + plugin, provider, exporter = create_plugin(plugin_type, sampler=ALWAYS_OFF) + try: + plugin.on_invocation_start(start_info()) + plugin.on_invocation_end(end_info(InvocationStatus.PENDING)) + assert not exporter.get_finished_spans() + finally: + provider.shutdown() + plugin, provider, exporter = create_plugin(plugin_type) + try: + plugin.on_invocation_start(start_info(first=False)) + plugin.on_invocation_end(end_info(InvocationStatus.SUCCEEDED, first=False)) + assert any( + s.name == "DurableExecutionRoot" for s in exporter.get_finished_spans() + ) + finally: + provider.shutdown() + + +def test_shared_trace_with_distinct_providers_and_start_times( + plugin_type: PluginType, +) -> None: + roots = [] + for index, arn in enumerate([ARN, ARN + "-other"]): + start = START + timedelta(seconds=index * 17) + resource = Resource({"service.name": f"provider-{index}"}) + plugin, provider, exporter = create_plugin( + plugin_type, + ExtractedContext(TRACE, None, Sampling.SAMPLED), + resource=resource, + ) + try: + plugin.on_invocation_start(start_info(arn=arn, start_time=start)) + plugin.on_invocation_end( + end_info(InvocationStatus.SUCCEEDED, arn=arn, start_time=start) + ) + root = next( + s + for s in exporter.get_finished_spans() + if s.name == "DurableExecutionRoot" + ) + assert root.context is not None + assert root.context.trace_id == TRACE + assert root.context.span_id == derive_execution_root_span_id(arn) + assert ( + root.start_time + == root.end_time + == int(start.timestamp() * 1_000_000_000) + ) + assert root.resource == resource + roots.append(root) + finally: + provider.shutdown() + assert roots[0].context != roots[1].context + assert roots[0].start_time != roots[1].start_time + assert roots[0].resource != roots[1].resource + + +class _CountingSampler(Sampler): + def __init__(self, decision: Decision) -> None: + self.calls = 0 + self.decision = decision + + def should_sample( + self, + parent_context: Any, + trace_id: int, + name: str, + kind: Any = None, + attributes: Any = None, + links: Any = None, + trace_state: Any = None, + ) -> SamplingResult: + self.calls += 1 + return SamplingResult( + self.decision, + attributes={"sampler.call": self.calls}, + trace_state=TraceState([("vendor", str(self.calls))]), + ) + + def get_description(self) -> str: + return "CountingSampler" + + +@pytest.mark.parametrize( + "decision", [Decision.RECORD_AND_SAMPLE, Decision.RECORD_ONLY, Decision.DROP] +) +def test_root_preserves_resolved_sampling_without_resampling( + plugin_type: PluginType, + decision: Decision, +) -> None: + sampler = _CountingSampler(decision) + plugin, provider, exporter = create_plugin(plugin_type, sampler=sampler) + try: + plugin.on_invocation_start(start_info()) + plugin.on_invocation_end(end_info(InvocationStatus.PENDING)) + assert sampler.calls == 1 + roots = [ + s for s in exporter.get_finished_spans() if s.name == "DurableExecutionRoot" + ] + if decision is Decision.RECORD_AND_SAMPLE: + assert len(roots) == 1 + root = roots[0] + assert root.attributes is not None + assert root.attributes == { + "durable.execution.arn": ARN, + "durable.execution.synthetic_root": True, + } + assert root.context is not None + assert root.context.trace_state == TraceState() + for span in exporter.get_finished_spans(): + if span.name != "DurableExecutionRoot": + assert span.attributes is not None and span.context is not None + assert span.attributes["sampler.call"] == 1 + assert span.context.trace_state == TraceState([("vendor", "1")]) + else: + assert not roots + finally: + provider.shutdown() + + +class _CollidingSampler(_CountingSampler): + def __init__(self) -> None: + super().__init__(Decision.RECORD_AND_SAMPLE) + self.results: list[SamplingResult] = [] + + def should_sample(self, *args: Any, **kwargs: Any) -> SamplingResult: + original = super().should_sample(*args, **kwargs) + result = SamplingResult( + original.decision, + attributes={ + **dict(original.attributes or {}), + "durable.execution.arn": f"sampler-arn-{self.calls}", + "durable.execution.synthetic_root": False, + }, + trace_state=original.trace_state, + ) + self.results.append(result) + return result + + +class _RootAttributeRecorder(SpanProcessor): + def __init__(self) -> None: + self.attributes: list[dict[str, Any]] = [] + + def on_start(self, span: Span, parent_context: Any = None) -> None: + if span.name == "DurableExecutionRoot": + self.attributes.append(dict(span.attributes or {})) + + +def test_root_identity_wins_over_sampler_attributes_before_processors( + plugin_type: PluginType, +) -> None: + sampler = _CollidingSampler() + plugin, provider, exporter = create_plugin(plugin_type, sampler=sampler) + recorder = _RootAttributeRecorder() + provider.add_span_processor(recorder) + try: + for index, status in enumerate( + [InvocationStatus.PENDING, InvocationStatus.SUCCEEDED] + ): + plugin.on_invocation_start(start_info(first=index == 0)) + plugin.on_invocation_end(end_info(status, first=index == 0)) + roots = [ + span + for span in exporter.get_finished_spans() + if span.name == "DurableExecutionRoot" + ] + assert len(roots) == len(recorder.attributes) == sampler.calls == 2 + assert roots[0].context is not None + assert roots[1].context is not None + assert roots[0].context.trace_id == roots[1].context.trace_id + assert roots[0].context.span_id == roots[1].context.span_id + for index, root in enumerate(roots, start=1): + expected = { + "durable.execution.arn": ARN, + "durable.execution.synthetic_root": True, + } + assert recorder.attributes[index - 1] == root.attributes == expected + assert root.context is not None + assert root.context.trace_state == TraceState() + # Protect only this anchor; do not mutate the invocation's result. + attributes = sampler.results[index - 1].attributes + assert attributes is not None + assert attributes["durable.execution.arn"] == f"sampler-arn-{index}" + assert attributes["durable.execution.synthetic_root"] is False + assert sampler.results[index - 1].trace_state == TraceState( + [("vendor", str(index))] + ) + finally: + provider.shutdown() + + +def test_reexported_anchor_has_no_variable_sampler_metadata( + plugin_type: PluginType, +) -> None: + sampler = _CountingSampler(Decision.RECORD_AND_SAMPLE) + plugin, provider, exporter = create_plugin(plugin_type, sampler=sampler) + recorder = _RootAttributeRecorder() + provider.add_span_processor(recorder) + try: + for index, status in enumerate( + [ + InvocationStatus.RETRY, + InvocationStatus.PENDING, + InvocationStatus.SUCCEEDED, + ] + ): + plugin.on_invocation_start(start_info(first=index < 2)) + plugin.on_invocation_end(end_info(status, first=index < 2)) + roots = [ + s for s in exporter.get_finished_spans() if s.name == "DurableExecutionRoot" + ] + invocations = [ + s for s in exporter.get_finished_spans() if s.name == "Invocation" + ] + assert len(roots) == len(invocations) == sampler.calls == 3 + expected = { + "durable.execution.arn": ARN, + "durable.execution.synthetic_root": True, + } + assert recorder.attributes == [expected, expected, expected] + assert all(s.attributes == expected for s in roots) + assert all( + s.context is not None and s.context.trace_state == TraceState() + for s in roots + ) + assert ( + len({(s.context.trace_id, s.context.span_id) for s in roots if s.context}) + == 1 + ) + for index, invocation in enumerate(invocations, start=1): + assert invocation.attributes is not None and invocation.context is not None + assert invocation.attributes["sampler.call"] == index + assert invocation.context.trace_state == TraceState( + [("vendor", str(index))] + ) + finally: + provider.shutdown() + + +class _ProcessorProbeSampler(_CountingSampler): + def __init__(self) -> None: + super().__init__(Decision.RECORD_AND_SAMPLE) + self.probe_calls = 0 + + def should_sample( + self, parent_context: Any, trace_id: int, name: str, *args: Any, **kwargs: Any + ) -> SamplingResult: + if name == "processor-probe": + self.probe_calls += 1 + return SamplingResult(Decision.DROP) + return super().should_sample(parent_context, trace_id, name, *args, **kwargs) + + +@pytest.mark.parametrize("forward_parent", [False, True]) +def test_processor_spans_use_their_own_sampling_policy( + plugin_type: PluginType, + forward_parent: bool, +) -> None: + sampler = _ProcessorProbeSampler() + plugin, provider, exporter = create_plugin(plugin_type, sampler=sampler) + observations: list[tuple[str, bool, bool]] = [] + + class Processor(SpanProcessor): + def on_start(self, span: Span, parent_context: Any = None) -> None: + if span.name not in {"DurableExecutionRoot", "Invocation", "Workflow"}: + return + # A processor may deliberately forward the supplied context. That + # preserves ancestry, but must not replay private sampling metadata. + with plugin._tracer.start_as_current_span( + "processor-probe", context=parent_context if forward_parent else None + ) as probe: + observations.append( + ( + span.name, + probe.is_recording(), + probe.get_span_context().trace_flags.sampled, + ) + ) + + provider.add_span_processor(Processor()) + try: + plugin.on_invocation_start(start_info()) + plugin.on_invocation_end(end_info(InvocationStatus.SUCCEEDED)) + assert {name for name, _, _ in observations} == { + "DurableExecutionRoot", + "Invocation", + "Workflow", + } + assert len(observations) == sampler.probe_calls == 3 + assert sampler.calls == 1 + assert all( + not recording and not sampled for _, recording, sampled in observations + ) + assert not [ + s for s in exporter.get_finished_spans() if s.name == "processor-probe" + ] + finally: + provider.shutdown() 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 ba373ad8d..b22d438c3 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 @@ -236,7 +236,7 @@ def test_invocation_start_and_end_emit_invocation_span(): spans = exporter.get_finished_spans() spans_by_name = {span.name: span for span in spans} # Terminal invocation also exports the Workflow span. - assert set(spans_by_name) == {"Invocation", "Workflow"} + assert set(spans_by_name) == {"DurableExecutionRoot", "Invocation", "Workflow"} invocation = spans_by_name["Invocation"] assert invocation.kind is SpanKind.INTERNAL assert invocation.attributes["durable.execution.arn"] == EXECUTION_ARN @@ -713,7 +713,9 @@ def test_operation_end_without_start_links_previous_logical_operation(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) 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} @@ -751,7 +753,9 @@ def test_continuation_span_uses_current_start_and_end_times(): ) after_callback = time.time_ns() - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert invocation_span.start_time <= span.start_time assert before_callback <= span.start_time <= span.end_time <= after_callback plugin.on_invocation_end(_invocation_end_info()) @@ -873,7 +877,9 @@ def test_retried_operation_uses_fresh_id_and_links_previous_logical_operation(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) 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} @@ -1020,7 +1026,9 @@ def test_user_function_callbacks_emit_attempt_span_attributes(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert span.name == "fetch-user attempt 1" assert span.attributes["durable.execution.arn"] == EXECUTION_ARN assert span.attributes["durable.operation.id"] == operation_id @@ -1073,7 +1081,9 @@ def test_step_attempt_span_name_includes_attempt_number(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert span.name == "fetch-user attempt 2" plugin.on_invocation_end(_invocation_end_info()) @@ -1116,7 +1126,9 @@ def test_step_attempt_span_name_defaults_to_first_attempt(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert span.name == "fetch-user attempt 1" plugin.on_invocation_end(_invocation_end_info()) @@ -1197,7 +1209,7 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( active_span.attributes["durable.operation.status"] == OperationStatus.STARTED.value ) - assert not exporter.get_finished_spans() + assert {s.name for s in exporter.get_finished_spans()} == {"DurableExecutionRoot"} plugin.on_operation_end( OperationEndInfo( @@ -1214,7 +1226,9 @@ def test_context_span_waits_for_terminal_status_and_omits_attempt_attributes( ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert span.attributes["durable.operation.type"] == OperationType.CONTEXT.value assert span.attributes["durable.operation.status"] == terminal_status.value assert "durable.attempt.number" not in span.attributes @@ -1716,7 +1730,9 @@ def test_replayed_context_span_links_previous_logical_operation(): ) ) - span = exporter.get_finished_spans()[0] + span = next( + s for s in exporter.get_finished_spans() if s.name != "DurableExecutionRoot" + ) assert span.context.span_id == random_span_id linked_span_ids = {link.context.span_id for link in span.links} assert linked_span_ids == { @@ -2127,7 +2143,9 @@ def suspend_on_worker() -> tuple[int, bool]: assert span_still_current is False assert span_key not in plugin._context_tokens assert suspended_span is not None - assert not exporter.get_finished_spans() + assert {s.name for s in exporter.get_finished_spans()} == { + "DurableExecutionRoot" + } # The timed resume lands on this thread, with nothing stale to unwind. plugin.on_user_function_start(_user_function_start_info(operation_id)) @@ -2185,7 +2203,7 @@ def test_nested_suspension_unwinds_scopes_in_reverse_order(): assert otel_context.get_current() == before_context assert set(plugin._context_tokens) == {"__invocation_context__"} # Neither span is ended: both operations are still in flight. - assert not exporter.get_finished_spans() + assert {s.name for s in exporter.get_finished_spans()} == {"DurableExecutionRoot"} # The timed in-process resume replays both contexts, outer first. plugin.on_user_function_start( 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 a50a06ac1..a407b2641 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 @@ -241,6 +241,7 @@ def test_global_proxy_binds_sdk_provider_before_first_invocation( plugin.on_invocation_end(_invocation_end()) assert {span.name for span in exporter.get_finished_spans()} == { + "DurableExecutionRoot", "Invocation", "Workflow", } @@ -273,6 +274,7 @@ def test_global_proxy_disables_entire_invocation_until_sdk_provider_is_ready( _run_step_lifecycle(plugin) plugin.on_invocation_end(_invocation_end()) assert {span.name for span in exporter.get_finished_spans()} == { + "DurableExecutionRoot", "Invocation", "Workflow", OP_NAME, diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_root_reentry.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_root_reentry.py new file mode 100644 index 000000000..6b6095f42 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_root_reentry.py @@ -0,0 +1,154 @@ +"""A synthetic anchor must not lend its ID or sampling to processor-created roots.""" + +from datetime import UTC, datetime + +import pytest +from aws_durable_execution_sdk_python.plugin import ( + InvocationStartInfo, + InvocationEndInfo, + InvocationStatus, +) +from opentelemetry import context +from opentelemetry.context import Context +from opentelemetry.sdk.trace import ( + TracerProvider, + SpanProcessor, + Span, + RandomIdGenerator, +) +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.sdk.trace.sampling import ALWAYS_ON, ALWAYS_OFF, Sampler +from opentelemetry.trace import SpanContext, Tracer + +from aws_durable_execution_sdk_python_otel import ( + ExecutionOtelPlugin, + InvocationOtelPlugin, + OtelPluginConfig, + ExtractedContext, + Sampling, +) +from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( + derive_execution_root_span_id, +) + + +TRACE = 0x11111111111111111111111111111111 +FALLBACK_TRACE = 0x22222222222222222222222222222222 +ARN = ( + "arn:aws:lambda:us-west-2:123456789012:function:reentry:1/durable-execution/test/id" +) +START = datetime(2026, 10, 1, tzinfo=UTC) + + +class _Ids(RandomIdGenerator): + def __init__(self) -> None: + self.trace_calls = 0 + + def generate_trace_id(self) -> int: + self.trace_calls += 1 + return FALLBACK_TRACE + self.trace_calls + + +class _Reentry(SpanProcessor): + def __init__(self) -> None: + self.tracer: Tracer | None = None + self.nested_context: SpanContext | None = None + self.context_unchanged = False + + def on_start(self, span: Span, parent_context: Context | None = None) -> None: + if span.name != "DurableExecutionRoot" or self.nested_context is not None: + return + assert self.tracer is not None + before = context.get_current() + nested = self.tracer.start_span("independent-root", context=Context()) + self.nested_context = nested.get_span_context() + self.context_unchanged = context.get_current() == before + nested.end() + + +@pytest.mark.parametrize("plugin_type", [ExecutionOtelPlugin, InvocationOtelPlugin]) +@pytest.mark.parametrize( + "target", ["same-scope", "same-tracer", "other-scope", "other-provider"] +) +@pytest.mark.parametrize("sampler", [ALWAYS_OFF, ALWAYS_ON]) +def test_processor_reentry_retains_independent_ids_and_sampling( + plugin_type: type[ExecutionOtelPlugin] | type[InvocationOtelPlugin], + target: str, + sampler: Sampler, +) -> None: + own = TracerProvider(sampler=sampler, id_generator=_Ids()) + exporter = InMemorySpanExporter() + own.add_span_processor(SimpleSpanProcessor(exporter)) + processor = _Reentry() + own.add_span_processor(processor) + plugin = plugin_type( + OtelPluginConfig( + tracer_provider=own, + instrument_name="durable-reentry", + enrich_logger=False, + context_extractor=lambda _: ExtractedContext(TRACE, None, Sampling.SAMPLED), + ) + ) + other = None + nested_exporter = exporter + if target == "same-scope": + processor.tracer = own.get_tracer("durable-reentry") + elif target == "same-tracer": + processor.tracer = plugin._tracer + elif target == "other-scope": + processor.tracer = own.get_tracer("other") + else: + other = TracerProvider(sampler=sampler, id_generator=_Ids()) + nested_exporter = InMemorySpanExporter() + other.add_span_processor(SimpleSpanProcessor(nested_exporter)) + processor.tracer = other.get_tracer("other") + original = context.get_current() + try: + plugin.on_invocation_start( + InvocationStartInfo( + request_id="request", + execution_arn=ARN, + execution_start_time=START, + is_first_invocation=True, + ) + ) + plugin.on_invocation_end( + InvocationEndInfo( + request_id="request", + execution_arn=ARN, + execution_start_time=START, + is_first_invocation=True, + status=InvocationStatus.PENDING, + ) + ) + root = next( + s for s in exporter.get_finished_spans() if s.name == "DurableExecutionRoot" + ) + assert root.context is not None + assert root.context.trace_id == TRACE + assert root.context.span_id == derive_execution_root_span_id(ARN) + assert root.context.trace_flags.sampled + nested = processor.nested_context + assert nested is not None + assert nested.trace_id == FALLBACK_TRACE + 1 + assert nested.span_id != root.context.span_id + assert nested.trace_flags.sampled is (sampler is ALWAYS_ON) + exported = [ + s + for s in nested_exporter.get_finished_spans() + if s.name == "independent-root" + ] + assert bool(exported) is (sampler is ALWAYS_ON) + if exported: + assert exported[0].parent is None + assert context.get_current() == original + assert processor.context_unchanged + assert processor.tracer is not None + after = processor.tracer.start_span("after", context=Context()) + assert after.get_span_context().trace_id == FALLBACK_TRACE + 2 + after.end() + finally: + own.shutdown() + if other is not None: + other.shutdown()