feat(otel): add durable sampling coordination

Author: ayushiahjolia

packages/aws-durable-execution-sdk-python-otel/src/aws_durable_execution_sdk_python_otel/durable_sampling.py

@@ -0,0 +1,135 @@
+"""Durable execution sampling support."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from opentelemetry import context as otel_context
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import Tracer as SdkTracer
+from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult
+from opentelemetry.trace import Span, SpanContext, SpanKind, TraceFlags
+
+from aws_durable_execution_sdk_python_otel.context_extractors import (
+    ExtractedContext,
+    Sampling,
+)
+
+
+_DURABLE_SAMPLING_INTENT_KEY = otel_context.create_key(
+    "aws_durable_execution_sampling_intent"
+)
+
+
+@dataclass(frozen=True)
+class DurableSamplingIntent:
+    """Sampling result to apply to each durable span in one invocation."""
+
+    result: SamplingResult
+
+
+class DurableSampler(Sampler):
+    """Sampler that honors a durable sampling intent carried on parent context."""
+
+    def __init__(self, delegate: Sampler) -> None:
+        self.delegate = delegate
+
+    @classmethod
+    def install_on_tracer(cls, tracer: SdkTracer) -> "DurableSampler":
+        current_sampler = tracer.sampler
+        if isinstance(current_sampler, cls):
+            return current_sampler
+        sampler = cls(current_sampler)
+        tracer.sampler = sampler
+        return sampler
+
+    def should_sample(
+        self,
+        parent_context: Context | None,
+        trace_id: int,
+        name: str,
+        kind: SpanKind | None = None,
+        attributes: Any = None,
+        links: Any = None,
+        trace_state: Any = None,
+    ) -> SamplingResult:
+        intent = otel_context.get_value(_DURABLE_SAMPLING_INTENT_KEY, parent_context)
+        if isinstance(intent, DurableSamplingIntent):
+            merged_attributes = dict(attributes or {})
+            merged_attributes.update(dict(intent.result.attributes or {}))
+            return SamplingResult(
+                intent.result.decision,
+                attributes=merged_attributes,
+                trace_state=intent.result.trace_state,
+            )
+        return self.delegate.should_sample(
+            parent_context,
+            trace_id,
+            name,
+            kind,
+            attributes,
+            links,
+            trace_state,
+        )
+
+    def get_description(self) -> str:
+        return f"DurableSampler{{{self.delegate.get_description()}}}"
+
+
+def store_sampling_intent(
+    parent_context: Context,
+    intent: DurableSamplingIntent | None,
+) -> Context:
+    """Attach a durable sampling intent to a span parent context."""
+    if intent is None:
+        return parent_context
+    return otel_context.set_value(_DURABLE_SAMPLING_INTENT_KEY, intent, parent_context)
+
+
+def resolve_sampling_result(
+    *,
+    extracted: ExtractedContext | None,
+    ambient_span: Span,
+    canonical_trace_id: int,
+    sampler: Sampler,
+    span_name: str,
+    attributes: dict[str, Any],
+) -> SamplingResult:
+    """Resolve one sampling decision for all durable spans in an invocation."""
+    sampling = extracted.sampling if extracted is not None else Sampling.UNDECIDED
+    if sampling is Sampling.SAMPLED:
+        return SamplingResult(Decision.RECORD_AND_SAMPLE)
+    if sampling is Sampling.NOT_SAMPLED:
+        return SamplingResult(Decision.DROP)
+
+    ambient_context = ambient_span.get_span_context()
+    if _is_same_trace(ambient_context, canonical_trace_id):
+        if bool(ambient_context.trace_flags & TraceFlags.SAMPLED):
+            decision = Decision.RECORD_AND_SAMPLE
+        elif ambient_span.is_recording():
+            decision = Decision.RECORD_ONLY
+        else:
+            decision = Decision.DROP
+        return SamplingResult(
+            decision,
+            trace_state=ambient_context.trace_state,
+        )
+
+    return sampler.should_sample(
+        parent_context=Context(),
+        trace_id=canonical_trace_id,
+        name=span_name,
+        kind=SpanKind.INTERNAL,
+        attributes=attributes,
+        links=(),
+        trace_state=None,
+    )
+
+
+def is_sampled(result: SamplingResult) -> bool:
+    return result.decision is Decision.RECORD_AND_SAMPLE
+
+
+def _is_same_trace(span_context: SpanContext, trace_id: int) -> bool:
+    return span_context.is_valid and span_context.trace_id == trace_id