Head SDK DSC population

Note that transaction name is missing here for now

Author: sl0thentr0py

sentry_sdk/integrations/opentelemetry/consts.py

@@ -1,4 +1,5 @@
 from opentelemetry.context import create_key
+from sentry_sdk.tracing_utils import Baggage
 
 
 # propagation keys
@@ -11,7 +12,8 @@
 SENTRY_USE_CURRENT_SCOPE_KEY = create_key("sentry_use_current_scope")
 SENTRY_USE_ISOLATION_SCOPE_KEY = create_key("sentry_use_isolation_scope")
 
-SENTRY_TRACE_STATE_DROPPED = "sentry_dropped"
+TRACESTATE_SAMPLED_KEY = Baggage.SENTRY_PREFIX + "sampled"
+TRACESTATE_SAMPLE_RATE_KEY = Baggage.SENTRY_PREFIX + "sample_rate"
 
 OTEL_SENTRY_CONTEXT = "otel"
 SPAN_ORIGIN = "auto.otel"

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -1,3 +1,4 @@
+from typing import cast
 from random import random
 
 from opentelemetry import trace
@@ -6,13 +7,17 @@
 from opentelemetry.trace.span import TraceState
 
 import sentry_sdk
-from sentry_sdk.integrations.opentelemetry.consts import SENTRY_TRACE_STATE_DROPPED
 from sentry_sdk.tracing_utils import has_tracing_enabled
 from sentry_sdk.utils import is_valid_sample_rate, logger
+from sentry_sdk.integrations.opentelemetry.consts import (
+    TRACESTATE_SAMPLED_KEY,
+    TRACESTATE_SAMPLE_RATE_KEY,
+)
 
-from typing import TYPE_CHECKING, Optional, Sequence
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from typing import Optional, Sequence, Union
     from opentelemetry.context import Context
     from opentelemetry.trace import Link, SpanKind
     from opentelemetry.trace.span import SpanContext
@@ -32,30 +37,42 @@ def get_parent_sampled(parent_context, trace_id):
         if parent_context.trace_flags.sampled:
             return True
 
-        dropped = parent_context.trace_state.get(SENTRY_TRACE_STATE_DROPPED) == "true"
-        if dropped:
+        dsc_sampled = parent_context.trace_state.get(TRACESTATE_SAMPLED_KEY)
+        if dsc_sampled == "true":
+            return True
+        elif dsc_sampled == "false":
             return False
 
-        # TODO-anton: fall back to sampling decision in DSC (for this die DSC needs to be set in the trace_state)
-
     return None
 
 
-def dropped_result(span_context):
-    # type: (SpanContext) -> SamplingResult
-    trace_state = span_context.trace_state.update(SENTRY_TRACE_STATE_DROPPED, "true")
+def dropped_result(span_context, attributes, sample_rate=None):
+    # type: (SpanContext, Attributes, Optional[float]) -> SamplingResult
+    # note that trace_state.add will NOT overwrite existing entries
+    # so these will only be added the first time in a root span sampling decision
+    trace_state = span_context.trace_state.add(TRACESTATE_SAMPLED_KEY, "false")
+    if sample_rate:
+        trace_state = trace_state.add(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
 
     return SamplingResult(
         Decision.DROP,
+        attributes=attributes,
         trace_state=trace_state,
     )
 
 
-def sampled_result(span_context):
-    # type: (SpanContext) -> SamplingResult
+def sampled_result(span_context, attributes, sample_rate):
+    # type: (SpanContext, Attributes, float) -> SamplingResult
+    # note that trace_state.add will NOT overwrite existing entries
+    # so these will only be added the first time in a root span sampling decision
+    trace_state = span_context.trace_state.add(TRACESTATE_SAMPLED_KEY, "true").add(
+        TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate)
+    )
+
     return SamplingResult(
         Decision.RECORD_AND_SAMPLE,
-        trace_state=span_context.trace_state,
+        attributes=attributes,
+        trace_state=trace_state,
     )
 
 
@@ -77,7 +94,7 @@ def should_sample(
 
         # No tracing enabled, thus no sampling
         if not has_tracing_enabled(client.options):
-            return dropped_result(parent_span_context)
+            return dropped_result(parent_span_context, attributes)
 
         sample_rate = None
 
@@ -112,16 +129,16 @@ def should_sample(
             logger.warning(
                 f"[Tracing] Discarding {name} because of invalid sample rate."
             )
-            return dropped_result(parent_span_context)
+            return dropped_result(parent_span_context, attributes)
 
         # Roll the dice on sample rate
-        sampled = random() < float(sample_rate)
+        sample_rate = float(cast("Union[bool, float, int]", sample_rate))
+        sampled = random() < sample_rate
 
-        # TODO-neel-potel set sample rate as attribute for DSC
         if sampled:
-            return sampled_result(parent_span_context)
+            return sampled_result(parent_span_context, attributes, sample_rate)
         else:
-            return dropped_result(parent_span_context)
+            return dropped_result(parent_span_context, attributes, sample_rate)
 
     def get_description(self) -> str:
         return self.__class__.__name__

sentry_sdk/integrations/opentelemetry/scope.py

@@ -8,7 +8,6 @@
     SpanContext,
     NonRecordingSpan,
     TraceFlags,
-    TraceState,
     use_span,
 )
 

sentry_sdk/integrations/opentelemetry/utils.py

@@ -15,6 +15,7 @@
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.sdk.trace import ReadableSpan
 
+import sentry_sdk
 from sentry_sdk.utils import Dsn
 from sentry_sdk.consts import SPANSTATUS
 from sentry_sdk.tracing import get_span_status_from_http_code, DEFAULT_SPAN_ORIGIN
@@ -24,7 +25,7 @@
 from sentry_sdk._types import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Optional, Mapping, Sequence, Union, ItemsView
+    from typing import Any, Optional, Mapping, Sequence, Union
     from sentry_sdk._types import OtelExtractedSpanData
 
 
@@ -300,9 +301,8 @@ def get_trace_context(span, span_data=None):
     if span.attributes:
         trace_context["data"] = dict(span.attributes)
 
-    trace_context["dynamic_sampling_context"] = dsc_from_trace_state(
-        span.context.trace_state
-    )
+    trace_state = get_trace_state(span)
+    trace_context["dynamic_sampling_context"] = dsc_from_trace_state(trace_state)
 
     # TODO-neel-potel profiler thread_id, thread_name
 
@@ -319,6 +319,11 @@ def trace_state_from_baggage(baggage):
     return TraceState(items)
 
 
+def baggage_from_trace_state(trace_state):
+    # type: (TraceState) -> Baggage
+    return Baggage(dsc_from_trace_state(trace_state))
+
+
 def serialize_trace_state(trace_state):
     # type: (TraceState) -> str
     sentry_items = []
@@ -336,3 +341,55 @@ def dsc_from_trace_state(trace_state):
             key = re.sub(Baggage.SENTRY_PREFIX_REGEX, "", k)
             dsc[key] = v
     return dsc
+
+
+def has_incoming_trace(trace_state):
+    # type: (TraceState) -> bool
+    """
+    The existence a sentry-trace_id in the baggage implies we continued an upstream trace.
+    """
+    return (Baggage.SENTRY_PREFIX + "trace_id") in trace_state
+
+
+def get_trace_state(span):
+    # type: (Union[Span, ReadableSpan]) -> TraceState
+    """
+    Get the existing trace_state with sentry items
+    or populate it if we are the head SDK.
+    """
+    span_context = span.get_span_context()
+    if not span_context:
+        return TraceState()
+
+    trace_state = span_context.trace_state
+
+    if has_incoming_trace(trace_state):
+        return trace_state
+    else:
+        client = sentry_sdk.get_client()
+        if not client.is_active():
+            return trace_state
+
+        options = client.options or {}
+
+        trace_state = trace_state.update(
+            Baggage.SENTRY_PREFIX + "trace_id", format_trace_id(span_context.trace_id)
+        )
+
+        if options.get("environment"):
+            trace_state = trace_state.update(
+                Baggage.SENTRY_PREFIX + "environment", options["environment"]
+            )
+
+        if options.get("release"):
+            trace_state = trace_state.update(
+                Baggage.SENTRY_PREFIX + "release", options["release"]
+            )
+
+        if options.get("dsn"):
+            trace_state = trace_state.update(
+                Baggage.SENTRY_PREFIX + "public_key", Dsn(options["dsn"]).public_key
+            )
+
+        # TODO-neel-potel head dsc transaction name
+        return trace_state

sentry_sdk/tracing.py

@@ -5,7 +5,12 @@
 from datetime import datetime, timedelta, timezone
 
 from opentelemetry import trace as otel_trace, context
-from opentelemetry.trace import format_trace_id, format_span_id, Span as OtelSpan
+from opentelemetry.trace import (
+    format_trace_id,
+    format_span_id,
+    Span as OtelSpan,
+    TraceState,
+)
 from opentelemetry.trace.status import StatusCode
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -1453,8 +1458,7 @@ def iter_headers(self):
             serialize_trace_state,
         )
 
-        trace_state = self._otel_span.get_span_context().trace_state
-        yield BAGGAGE_HEADER_NAME, serialize_trace_state(trace_state)
+        yield BAGGAGE_HEADER_NAME, serialize_trace_state(self.trace_state)
 
     def to_traceparent(self):
         # type: () -> str
@@ -1471,10 +1475,26 @@ def to_traceparent(self):
 
         return traceparent
 
+    @property
+    def trace_state(self):
+        # type: () -> TraceState
+        from sentry_sdk.integrations.opentelemetry.utils import (
+            get_trace_state,
+        )
+
+        return get_trace_state(self._otel_span)
+
     def to_baggage(self):
-        # type: () -> Optional[Baggage]
-        # TODO-neel-potel head SDK populate baggage mess
-        pass
+        # type: () -> Baggage
+        return self.get_baggage()
+
+    def get_baggage(self):
+        # type: () -> Baggage
+        from sentry_sdk.integrations.opentelemetry.utils import (
+            baggage_from_trace_state,
+        )
+
+        return baggage_from_trace_state(self.trace_state)
 
     def set_tag(self, key, value):
         # type: (str, Any) -> None
@@ -1568,10 +1588,6 @@ def set_context(self, key, value):
         # type: (str, Any) -> None
         pass
 
-    def get_baggage(self):
-        # type: () -> Baggage
-        pass
-
 
 if TYPE_CHECKING: