Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"]


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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")
Comment thread
ayushiahjolia marked this conversation as resolved.
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.

Expand Down
Original file line number Diff line number Diff line change
@@ -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)
)
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)

Expand Down Expand Up @@ -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
)
Loading
Loading