Head SDK DSC population (#3599)

* Populates a DSC with correct values when we don't have an incoming trace.
* We rely on `trace_state.add` only adding new keys to the tracestate so these values will be populated in the first sampling decision on the root and just be propagated further.

Note that transaction name is missing here for now and will be dealt with separately as part of the transaction name PRs.


closes #3479

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: