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..126a3fd6 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 @@ -3,6 +3,8 @@ 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 @@ -13,6 +15,45 @@ from aws_durable_execution_sdk_python.plugin import InvocationStartInfo + +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"], "Context"] 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/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/tests/test_deterministic_id_generator.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_deterministic_id_generator.py index 25230d95..6bcd320a 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 @@ -12,6 +12,8 @@ from aws_durable_execution_sdk_python_otel.deterministic_id_generator import ( DeterministicIdGenerator, _to_otel_trace_id, + derive_execution_root_span_id, + derive_workflow_span_id, operation_id_to_span_id, ) @@ -303,3 +305,37 @@ async def main() -> tuple[tuple[int, int], tuple[int, int]]: assert result_a == (task_a_trace_id, task_a_span_id) assert result_b == (task_b_trace_id, task_b_span_id) + + +# --------------------------------------------------------------------------- +# derive_execution_root_span_id +# --------------------------------------------------------------------------- +_ROOT_ARN = "test-arn/execution-root" + + +def test_derive_execution_root_span_id_is_deterministic(): + assert derive_execution_root_span_id(_ROOT_ARN) == derive_execution_root_span_id( + _ROOT_ARN + ) + + +def test_derive_execution_root_span_id_differs_by_arn(): + assert derive_execution_root_span_id(_ROOT_ARN) != derive_execution_root_span_id( + _ROOT_ARN + "-other" + ) + + +def test_derive_execution_root_span_id_is_64_bit(): + span_id = derive_execution_root_span_id(_ROOT_ARN) + assert 0 < span_id < 2**64 + + +def test_derive_execution_root_span_id_rejects_empty_arn(): + with pytest.raises(ValueError, match="execution ARN is required"): + derive_execution_root_span_id("") + + +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 + ) diff --git a/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_trace_context.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_trace_context.py new file mode 100644 index 00000000..ae7a4f94 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_execution_trace_context.py @@ -0,0 +1,202 @@ +"""Tests for execution trace ancestry resolution.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +from opentelemetry.trace import TraceFlags + +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, +) +from aws_durable_execution_sdk_python_otel.execution_trace_context import ( + ExecutionTraceContext, + canonical_trace_id, +) + + +# --------------------------------------------------------------------------- +# ExtractedContext validity properties +# --------------------------------------------------------------------------- +def test_extracted_context_validity_properties() -> None: + complete: ExtractedContext = ExtractedContext( + trace_id=int("5759e988bd862e3fe1be46a994272793", 16), + parent_span_id=int("53995c3f42cd8ad8", 16), + ) + assert complete.has_valid_trace_id is True + assert complete.has_valid_parent_span_id is True + assert complete.has_complete_remote_parent is True + + +def test_extracted_context_rejects_missing_and_zero_ids() -> None: + missing: ExtractedContext = ExtractedContext(trace_id=None, parent_span_id=None) + assert missing.has_valid_trace_id is False + assert missing.has_valid_parent_span_id is False + assert missing.has_complete_remote_parent is False + + zero: ExtractedContext = ExtractedContext(trace_id=0, parent_span_id=0) + assert zero.has_valid_trace_id is False + assert zero.has_valid_parent_span_id is False + + +def test_extracted_context_parent_alone_is_not_complete() -> None: + parent_only: ExtractedContext = ExtractedContext( + trace_id=None, + parent_span_id=int("53995c3f42cd8ad8", 16), + ) + assert parent_only.has_valid_parent_span_id is True + assert parent_only.has_valid_trace_id is False + assert parent_only.has_complete_remote_parent is False + + +EXECUTION_ARN: str = "test-arn/execution-trace-context" +START_TIME: datetime = datetime(2026, 8, 27, 5, 11, 47, tzinfo=UTC) +REMOTE_TRACE_ID: int = int("5759e988bd862e3fe1be46a994272793", 16) +REMOTE_PARENT_ID: int = int("53995c3f42cd8ad8", 16) + + +def _complete_remote(sampling: Sampling = Sampling.SAMPLED) -> ExtractedContext: + return ExtractedContext( + trace_id=REMOTE_TRACE_ID, + parent_span_id=REMOTE_PARENT_ID, + sampling=sampling, + ) + + +# --------------------------------------------------------------------------- +# ExecutionTraceContext.resolve +# --------------------------------------------------------------------------- +def test_resolve_uses_complete_remote_parent_as_ancestor() -> None: + ctx: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=_complete_remote(), + canonical_trace_id=REMOTE_TRACE_ID, + execution_arn=EXECUTION_ARN, + root_sampled=lambda: False, + ) + + ancestor = ctx.execution_ancestor + assert ancestor.trace_id == REMOTE_TRACE_ID + assert ancestor.span_id == REMOTE_PARENT_ID + assert ancestor.is_remote is True + assert ctx.trace_id == REMOTE_TRACE_ID + + +def test_resolve_falls_back_to_synthetic_root_without_remote_parent() -> None: + trace_id: int = _to_otel_trace_id(EXECUTION_ARN, START_TIME) + ctx: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=None, + canonical_trace_id=trace_id, + execution_arn=EXECUTION_ARN, + root_sampled=lambda: True, + ) + + ancestor = ctx.execution_ancestor + assert ancestor.trace_id == trace_id + assert ancestor.span_id == derive_execution_root_span_id(EXECUTION_ARN) + assert ancestor.is_remote is False + + +def test_resolve_uses_synthetic_root_when_parent_incomplete() -> None: + incomplete: ExtractedContext = ExtractedContext( + trace_id=REMOTE_TRACE_ID, + parent_span_id=None, + sampling=Sampling.UNDECIDED, + ) + ctx: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=incomplete, + canonical_trace_id=REMOTE_TRACE_ID, + execution_arn=EXECUTION_ARN, + root_sampled=lambda: False, + ) + + assert ctx.execution_ancestor.span_id == derive_execution_root_span_id( + EXECUTION_ARN + ) + assert ctx.execution_ancestor.is_remote is False + + +# --------------------------------------------------------------------------- +# ExecutionTraceContext.resolve: trace flags +# --------------------------------------------------------------------------- +def test_resolve_backend_sampled_sets_sampled_flag() -> None: + ctx: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=_complete_remote(Sampling.SAMPLED), + canonical_trace_id=REMOTE_TRACE_ID, + execution_arn=EXECUTION_ARN, + root_sampled=lambda: False, + ) + + assert bool(ctx.trace_flags & TraceFlags.SAMPLED) is True + + +def test_resolve_backend_not_sampled_clears_sampled_flag() -> None: + ctx: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=_complete_remote(Sampling.NOT_SAMPLED), + canonical_trace_id=REMOTE_TRACE_ID, + execution_arn=EXECUTION_ARN, + root_sampled=lambda: True, + ) + + assert bool(ctx.trace_flags & TraceFlags.SAMPLED) is False + + +def test_resolve_undecided_defers_to_root_sampled_callback() -> None: + sampled: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=None, + canonical_trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME), + execution_arn=EXECUTION_ARN, + root_sampled=lambda: True, + ) + dropped: ExecutionTraceContext = ExecutionTraceContext.resolve( + extracted=None, + canonical_trace_id=_to_otel_trace_id(EXECUTION_ARN, START_TIME), + execution_arn=EXECUTION_ARN, + root_sampled=lambda: False, + ) + + assert bool(sampled.trace_flags & TraceFlags.SAMPLED) is True + assert bool(dropped.trace_flags & TraceFlags.SAMPLED) is False + + +# --------------------------------------------------------------------------- +# canonical_trace_id +# --------------------------------------------------------------------------- +def test_canonical_trace_id_prefers_valid_extracted_trace_id() -> None: + result: int = canonical_trace_id( + extracted=_complete_remote(), + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + ) + + assert result == REMOTE_TRACE_ID + + +def test_canonical_trace_id_falls_back_to_derived_id() -> None: + result: int = canonical_trace_id( + extracted=None, + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + ) + + assert result == _to_otel_trace_id(EXECUTION_ARN, START_TIME) + + +def test_canonical_trace_id_falls_back_when_extracted_trace_id_invalid() -> None: + no_trace: ExtractedContext = ExtractedContext( + trace_id=None, + parent_span_id=REMOTE_PARENT_ID, + sampling=Sampling.SAMPLED, + ) + + result: int = canonical_trace_id( + extracted=no_trace, + execution_arn=EXECUTION_ARN, + execution_start_time=START_TIME, + ) + + assert result == _to_otel_trace_id(EXECUTION_ARN, START_TIME)