From daabf76aeb9dc838dfed7d8b77a04d73d817fb21 Mon Sep 17 00:00:00 2001 From: Ayushi Ahjolia Date: Fri, 28 Aug 2026 17:00:17 -0700 Subject: [PATCH] feat(otel): add durable sampling coordination --- .../durable_sampling.py | 200 ++++++++ .../tests/test_durable_sampling.py | 459 ++++++++++++++++++ 2 files changed, 659 insertions(+) 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/tests/test_durable_sampling.py 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..f2e9f2af --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py @@ -0,0 +1,200 @@ +"""Durable execution sampling support.""" + +from __future__ import annotations + +import functools +import inspect +from dataclasses import dataclass +from typing import Any, Callable + +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 _delegate_should_sample( + self.delegate, + 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. + + Trace state from a same-trace ambient span is preserved across every + branch, so an explicit backend decision overrides only the sampling + outcome, not vendor/W3C ``tracestate`` propagation. + """ + ambient_context = ambient_span.get_span_context() + on_canonical_trace = _is_same_trace(ambient_context, canonical_trace_id) + ambient_trace_state = ambient_context.trace_state if on_canonical_trace else None + + sampling = extracted.sampling if extracted is not None else Sampling.UNDECIDED + if sampling is Sampling.SAMPLED: + return SamplingResult( + Decision.RECORD_AND_SAMPLE, + trace_state=ambient_trace_state, + ) + if sampling is Sampling.NOT_SAMPLED: + return SamplingResult(Decision.DROP, trace_state=ambient_trace_state) + + if on_canonical_trace: + 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_trace_state) + + return _delegate_should_sample( + sampler, + Context(), + canonical_trace_id, + span_name, + SpanKind.INTERNAL, + attributes, + (), + None, + ) + + +def is_sampled(result: SamplingResult) -> bool: + return result.decision is Decision.RECORD_AND_SAMPLE + + +@functools.lru_cache(maxsize=None) +def _delegate_accepts_trace_state(should_sample: Callable[..., SamplingResult]) -> bool: + """Return whether a sampler's ``should_sample`` accepts ``trace_state``. + + ``trace_state`` was added to ``Sampler.should_sample`` in OpenTelemetry SDK + 1.21. The package supports ``opentelemetry-sdk>=1.20.0``, whose samplers end + at ``links``. A parameter probe (rather than a call-time ``try/except``) + avoids masking a ``TypeError`` raised inside the sampler body and never + invokes the sampler twice. + """ + try: + parameters = inspect.signature(should_sample).parameters + except (TypeError, ValueError): + return True + if "trace_state" in parameters: + return True + return any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in parameters.values() + ) + + +def _delegate_should_sample( + sampler: Sampler, + parent_context: Context | None, + trace_id: int, + name: str, + kind: SpanKind | None, + attributes: Any, + links: Any, + trace_state: Any, +) -> SamplingResult: + """Call a delegate sampler using the signature its OTel version supports.""" + if _delegate_accepts_trace_state(sampler.should_sample): + return sampler.should_sample( + parent_context, + trace_id, + name, + kind, + attributes, + links, + trace_state=trace_state, + ) + return sampler.should_sample( + parent_context, + trace_id, + name, + kind, + attributes, + links, + ) + + +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/tests/test_durable_sampling.py b/packages/aws-durable-execution-sdk-python-otel/tests/test_durable_sampling.py new file mode 100644 index 00000000..24681b68 --- /dev/null +++ b/packages/aws-durable-execution-sdk-python-otel/tests/test_durable_sampling.py @@ -0,0 +1,459 @@ +"""Tests for durable execution sampling support.""" + +from __future__ import annotations + +from typing import Any, cast + +import pytest +from opentelemetry.context import Context +from opentelemetry.sdk.trace.sampling import ( + ALWAYS_OFF, + ALWAYS_ON, + Decision, + Sampler, + SamplingResult, +) +from opentelemetry.trace import ( + NonRecordingSpan, + Span, + SpanContext, + SpanKind, + TraceFlags, + TraceState, +) + +from aws_durable_execution_sdk_python_otel.context_extractors import ( + ExtractedContext, + Sampling, +) +from aws_durable_execution_sdk_python_otel.durable_sampling import ( + DurableSampler, + DurableSamplingIntent, + _delegate_accepts_trace_state, + is_sampled, + resolve_sampling_result, + store_sampling_intent, +) + + +TRACE_ID: int = int("5759e988bd862e3fe1be46a994272793", 16) +SPAN_ID: int = int("53995c3f42cd8ad8", 16) + + +def _span_context( + *, + trace_id: int = TRACE_ID, + sampled: bool = False, + trace_state: TraceState | None = None, +) -> SpanContext: + flags: TraceFlags = TraceFlags( + TraceFlags.SAMPLED if sampled else TraceFlags.DEFAULT + ) + return SpanContext( + trace_id=trace_id, + span_id=SPAN_ID, + is_remote=False, + trace_flags=flags, + trace_state=trace_state if trace_state is not None else TraceState(), + ) + + +def _invalid_span() -> Span: + return NonRecordingSpan(SpanContext(0, 0, is_remote=False)) + + +class _RecordingSpan(NonRecordingSpan): + """A span context wrapper that reports itself as recording.""" + + def is_recording(self) -> bool: + return True + + +def _extracted(sampling: Sampling) -> ExtractedContext: + return ExtractedContext( + trace_id=TRACE_ID, + parent_span_id=SPAN_ID, + sampling=sampling, + ) + + +# --------------------------------------------------------------------------- +# resolve_sampling_result: explicit backend decision wins +# --------------------------------------------------------------------------- +def test_backend_sampled_records_and_samples() -> None: + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.SAMPLED), + ambient_span=_invalid_span(), + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_OFF, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + + +def test_backend_not_sampled_drops() -> None: + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.NOT_SAMPLED), + ambient_span=_invalid_span(), + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_ON, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.DROP + + +def test_backend_sampled_preserves_same_trace_ambient_trace_state() -> None: + trace_state: TraceState = TraceState([("vendor", "opaque")]) + ambient: Span = NonRecordingSpan(_span_context(trace_state=trace_state)) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.SAMPLED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_OFF, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + assert result.trace_state == trace_state + + +def test_backend_not_sampled_preserves_same_trace_ambient_trace_state() -> None: + trace_state: TraceState = TraceState([("vendor", "opaque")]) + ambient: Span = NonRecordingSpan(_span_context(trace_state=trace_state)) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.NOT_SAMPLED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_ON, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.DROP + assert result.trace_state == trace_state + + +def test_backend_decision_ignores_different_trace_ambient_trace_state() -> None: + other_trace_id: int = TRACE_ID ^ 0x1 + trace_state: TraceState = TraceState([("vendor", "opaque")]) + ambient: Span = NonRecordingSpan( + _span_context(trace_id=other_trace_id, trace_state=trace_state) + ) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.SAMPLED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_OFF, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + assert result.trace_state != trace_state + + +# --------------------------------------------------------------------------- +# resolve_sampling_result: same-trace ambient span decides when undecided +# --------------------------------------------------------------------------- +def test_undecided_uses_sampled_ambient_span_on_same_trace() -> None: + trace_state: TraceState = TraceState([("vendor", "opaque")]) + ambient: Span = NonRecordingSpan( + _span_context(sampled=True, trace_state=trace_state) + ) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.UNDECIDED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_OFF, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + assert result.trace_state == trace_state + + +def test_undecided_uses_recording_ambient_span_on_same_trace() -> None: + ambient: Span = _RecordingSpan(_span_context(sampled=False)) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.UNDECIDED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_ON, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_ONLY + + +def test_undecided_drops_non_recording_ambient_span_on_same_trace() -> None: + ambient: Span = NonRecordingSpan(_span_context(sampled=False)) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.UNDECIDED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_ON, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.DROP + + +# --------------------------------------------------------------------------- +# resolve_sampling_result: falls back to the configured sampler otherwise +# --------------------------------------------------------------------------- +def test_undecided_delegates_to_sampler_when_no_extracted_context() -> None: + result: SamplingResult = resolve_sampling_result( + extracted=None, + ambient_span=_invalid_span(), + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_ON, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + + +def test_undecided_delegates_to_sampler_for_different_trace_ambient_span() -> None: + other_trace_id: int = TRACE_ID ^ 0x1 + ambient: Span = NonRecordingSpan( + _span_context(trace_id=other_trace_id, sampled=True) + ) + + result: SamplingResult = resolve_sampling_result( + extracted=_extracted(Sampling.UNDECIDED), + ambient_span=ambient, + canonical_trace_id=TRACE_ID, + sampler=ALWAYS_OFF, + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.DROP + + +# --------------------------------------------------------------------------- +# DurableSampler +# --------------------------------------------------------------------------- +def test_durable_sampler_honors_stored_intent_over_delegate() -> None: + sampler: DurableSampler = DurableSampler(ALWAYS_OFF) + intent: DurableSamplingIntent = DurableSamplingIntent( + SamplingResult(Decision.RECORD_AND_SAMPLE) + ) + parent_context: Context = store_sampling_intent(Context(), intent) + + result: SamplingResult = sampler.should_sample(parent_context, TRACE_ID, "span") + + assert result.decision is Decision.RECORD_AND_SAMPLE + + +def test_durable_sampler_merges_span_and_intent_attributes() -> None: + sampler: DurableSampler = DurableSampler(ALWAYS_OFF) + intent: DurableSamplingIntent = DurableSamplingIntent( + SamplingResult(Decision.RECORD_AND_SAMPLE, attributes={"from": "intent"}) + ) + parent_context: Context = store_sampling_intent(Context(), intent) + + result: SamplingResult = sampler.should_sample( + parent_context, + TRACE_ID, + "span", + attributes={"from": "span"}, + ) + + assert result.attributes is not None + assert result.attributes["from"] == "intent" + + +def test_durable_sampler_intent_preserves_trace_state() -> None: + trace_state: TraceState = TraceState([("vendor", "opaque")]) + sampler: DurableSampler = DurableSampler(ALWAYS_OFF) + intent: DurableSamplingIntent = DurableSamplingIntent( + SamplingResult(Decision.RECORD_AND_SAMPLE, trace_state=trace_state) + ) + parent_context: Context = store_sampling_intent(Context(), intent) + + result: SamplingResult = sampler.should_sample(parent_context, TRACE_ID, "span") + + assert result.trace_state == trace_state + + +def test_durable_sampler_delegates_without_intent() -> None: + sampler: DurableSampler = DurableSampler(ALWAYS_ON) + + result: SamplingResult = sampler.should_sample(Context(), TRACE_ID, "span") + + assert result.decision is Decision.RECORD_AND_SAMPLE + + +class _LegacySampler: + """A pre-1.21 sampler whose should_sample signature ends at ``links``. + + Deliberately not a ``Sampler`` subclass: it simulates the OpenTelemetry + SDK 1.20 signature, which predates the ``trace_state`` parameter. + """ + + def should_sample( + self, + parent_context: Context | None, + trace_id: int, + name: str, + kind: SpanKind | None = None, + attributes: Any = None, + links: Any = None, + ) -> SamplingResult: + return SamplingResult(Decision.RECORD_AND_SAMPLE) + + def get_description(self) -> str: + return "LegacySampler" + + +def test_durable_sampler_delegates_to_pre_1_21_sampler_signature() -> None: + sampler: DurableSampler = DurableSampler(cast(Sampler, _LegacySampler())) + + result: SamplingResult = sampler.should_sample(Context(), TRACE_ID, "span") + + assert result.decision is Decision.RECORD_AND_SAMPLE + + +def test_resolve_delegates_to_pre_1_21_sampler_signature() -> None: + result: SamplingResult = resolve_sampling_result( + extracted=None, + ambient_span=_invalid_span(), + canonical_trace_id=TRACE_ID, + sampler=cast(Sampler, _LegacySampler()), + span_name="Workflow", + attributes={}, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + + +class _RaisingSampler(Sampler): + """A modern-signature sampler whose body raises TypeError.""" + + 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: + raise TypeError("boom from sampler body") + + def get_description(self) -> str: + return "RaisingSampler" + + +def test_durable_sampler_does_not_swallow_sampler_body_type_error() -> None: + sampler: DurableSampler = DurableSampler(_RaisingSampler()) + + with pytest.raises(TypeError, match="boom from sampler body"): + sampler.should_sample(Context(), TRACE_ID, "span") + + +class _KwargsSampler: + """A wrapper-style sampler that captures trace_state via **kwargs only. + + Deliberately not a ``Sampler`` subclass: it models a forwarding sampler + whose ``should_sample`` accepts extra keyword arguments through ``**kwargs``. + """ + + def __init__(self) -> None: + self.received_trace_state: TraceState | None = None + + def should_sample( + self, + parent_context: Context | None, + trace_id: int, + name: str, + kind: SpanKind | None = None, + attributes: Any = None, + links: Any = None, + **kwargs: Any, + ) -> SamplingResult: + self.received_trace_state = kwargs.get("trace_state") + return SamplingResult(Decision.RECORD_AND_SAMPLE) + + def get_description(self) -> str: + return "KwargsSampler" + + +def test_durable_sampler_passes_trace_state_by_keyword_to_kwargs_delegate() -> None: + trace_state: TraceState = TraceState([("vendor", "opaque")]) + delegate: _KwargsSampler = _KwargsSampler() + sampler: DurableSampler = DurableSampler(cast(Sampler, delegate)) + + result: SamplingResult = sampler.should_sample( + Context(), + TRACE_ID, + "span", + trace_state=trace_state, + ) + + assert result.decision is Decision.RECORD_AND_SAMPLE + assert delegate.received_trace_state == trace_state + + +def test_delegate_accepts_trace_state_defaults_true_when_uninspectable() -> None: + # Some built-in callables have no inspectable signature; assume the modern + # signature rather than dropping trace_state. + uninspectable: object = iter + + assert _delegate_accepts_trace_state(uninspectable) is True # type: ignore[arg-type] + + +def test_durable_sampler_description_wraps_delegate() -> None: + sampler: DurableSampler = DurableSampler(ALWAYS_ON) + + assert ( + sampler.get_description() == f"DurableSampler{{{ALWAYS_ON.get_description()}}}" + ) + + +class _StubTracer: + def __init__(self, sampler: Sampler) -> None: + self.sampler: Sampler = sampler + + +def test_install_on_tracer_wraps_and_is_idempotent() -> None: + tracer: Any = _StubTracer(ALWAYS_ON) + + first: DurableSampler = DurableSampler.install_on_tracer(tracer) + assert isinstance(tracer.sampler, DurableSampler) + assert first.delegate is ALWAYS_ON + + second: DurableSampler = DurableSampler.install_on_tracer(tracer) + assert second is first + + +# --------------------------------------------------------------------------- +# store_sampling_intent / is_sampled +# --------------------------------------------------------------------------- +def test_store_sampling_intent_returns_context_unchanged_when_none() -> None: + parent_context: Context = Context() + + assert store_sampling_intent(parent_context, None) is parent_context + + +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 + assert is_sampled(SamplingResult(Decision.DROP)) is False