Make DSC and trace_context work for outgoing traces

sentry_sdk/_types.py

@@ -197,3 +197,5 @@ class SDKInfo(TypedDict):
     )
 
     HttpStatusCodeRange = Union[int, Container[int]]
+
+    OtelExtractedSpanData = tuple[str, str, Optional[str], Optional[int], Optional[str]]

sentry_sdk/integrations/opentelemetry/potel_span_processor.py

@@ -18,6 +18,7 @@
     convert_from_otel_timestamp,
     extract_span_attributes,
     extract_span_data,
+    get_trace_context,
 )
 from sentry_sdk.integrations.opentelemetry.consts import (
     OTEL_SENTRY_CONTEXT,
@@ -136,26 +137,12 @@ def _root_span_to_transaction_event(self, span):
         if event is None:
             return None
 
-        trace_id = format_trace_id(span.context.trace_id)
-        span_id = format_span_id(span.context.span_id)
-        parent_span_id = format_span_id(span.parent.span_id) if span.parent else None
-
-        (op, description, status, _, origin) = extract_span_data(span)
-
-        trace_context = {
-            "trace_id": trace_id,
-            "span_id": span_id,
-            "origin": origin or DEFAULT_SPAN_ORIGIN,
-            "op": op,
-            "status": status,
-        }  # type: dict[str, Any]
-
-        if parent_span_id:
-            trace_context["parent_span_id"] = parent_span_id
-        if span.attributes:
-            trace_context["data"] = dict(span.attributes)
+        span_data = extract_span_data(span)
+        (_, description, _, _, _) = span_data
 
+        trace_context = get_trace_context(span, span_data=span_data)
         contexts = {"trace": trace_context}
+
         if span.resource.attributes:
             contexts[OTEL_SENTRY_CONTEXT] = {"resource": dict(span.resource.attributes)}
 

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -41,16 +41,21 @@ def get_parent_sampled(parent_context, trace_id):
     return None
 
 
-def dropped(parent_context=None):
-    # type: (Optional[SpanContext]) -> SamplingResult
-    trace_state = parent_context.trace_state if parent_context is not None else None
-    updated_trace_context = trace_state or TraceState()
-    updated_trace_context = updated_trace_context.update(
-        SENTRY_TRACE_STATE_DROPPED, "true"
-    )
+def dropped_result(span_context):
+    # type: (SpanContext) -> SamplingResult
+    trace_state = span_context.trace_state.update(SENTRY_TRACE_STATE_DROPPED, "true")
+
     return SamplingResult(
         Decision.DROP,
-        trace_state=updated_trace_context,
+        trace_state=trace_state,
+    )
+
+
+def sampled_result(span_context):
+    # type: (SpanContext) -> SamplingResult
+    return SamplingResult(
+        Decision.RECORD_AND_SAMPLE,
+        trace_state=span_context.trace_state,
     )
 
 
@@ -68,12 +73,11 @@ def should_sample(
         # type: (...) -> SamplingResult
         client = sentry_sdk.get_client()
 
-        parent_span = trace.get_current_span(parent_context)
-        parent_context = parent_span.get_span_context() if parent_span else None
+        parent_span_context = trace.get_current_span(parent_context).get_span_context()
 
         # No tracing enabled, thus no sampling
         if not has_tracing_enabled(client.options):
-            return dropped(parent_context)
+            return dropped_result(parent_span_context)
 
         sample_rate = None
 
@@ -89,14 +93,14 @@ def should_sample(
                 "transaction_context": {
                     "name": name,
                 },
-                "parent_sampled": get_parent_sampled(parent_context, trace_id),
+                "parent_sampled": get_parent_sampled(parent_span_context, trace_id),
             }
 
             sample_rate = client.options["traces_sampler"](sampling_context)
 
         else:
             # Check if there is a parent with a sampling decision
-            parent_sampled = get_parent_sampled(parent_context, trace_id)
+            parent_sampled = get_parent_sampled(parent_span_context, trace_id)
             if parent_sampled is not None:
                 sample_rate = parent_sampled
             else:
@@ -108,15 +112,16 @@ def should_sample(
             logger.warning(
                 f"[Tracing] Discarding {name} because of invalid sample rate."
             )
-            return dropped(parent_context)
+            return dropped_result(parent_span_context)
 
         # Roll the dice on sample rate
         sampled = random() < float(sample_rate)
 
+        # TODO-neel-potel set sample rate as attribute for DSC
         if sampled:
-            return SamplingResult(Decision.RECORD_AND_SAMPLE)
+            return sampled_result(parent_span_context)
         else:
-            return dropped(parent_context)
+            return dropped_result(parent_span_context)
 
     def get_description(self) -> str:
         return self.__class__.__name__

sentry_sdk/integrations/opentelemetry/scope.py

@@ -4,14 +4,15 @@
 from contextlib import contextmanager
 
 from opentelemetry.context import get_value, set_value, attach, detach, get_current
-from opentelemetry.trace import SpanContext, NonRecordingSpan, TraceFlags, use_span
+from opentelemetry.trace import SpanContext, NonRecordingSpan, TraceFlags, TraceState, use_span
 
 from sentry_sdk.integrations.opentelemetry.consts import (
     SENTRY_SCOPES_KEY,
     SENTRY_FORK_ISOLATION_SCOPE_KEY,
     SENTRY_USE_CURRENT_SCOPE_KEY,
     SENTRY_USE_ISOLATION_SCOPE_KEY,
 )
+from sentry_sdk.integrations.opentelemetry.utils import trace_state_from_baggage
 from sentry_sdk.scope import Scope, ScopeType
 from sentry_sdk.tracing import POTelSpan
 from sentry_sdk._types import TYPE_CHECKING
@@ -93,15 +94,17 @@ def _incoming_otel_span_context(self):
             else TraceFlags.DEFAULT
         )
 
-        # TODO-neel-potel tracestate
+        # TODO-neel-potel do we need parent and sampled like JS?
+        trace_state = None
+        if self._propagation_context.baggage:
+            trace_state = trace_state_from_baggage(self._propagation_context.baggage)
+
         span_context = SpanContext(
             trace_id=int(self._propagation_context.trace_id, 16),  # type: ignore
             span_id=int(self._propagation_context.parent_span_id, 16),  # type: ignore
             is_remote=True,
             trace_flags=trace_flags,
-            # TODO-anton: add trace_state (mapping[str,str]) with the parentSpanId, dsc and sampled from self._propagation_context
-            # trace_state={
-            # }
+            trace_state=trace_state,
         )
 
         return span_context

sentry_sdk/integrations/opentelemetry/utils.py

@@ -1,20 +1,24 @@
+import re
 from typing import cast
 from datetime import datetime, timezone
 
-from opentelemetry.trace import SpanKind, StatusCode
+from urllib3.util import parse_url as urlparse
+from urllib.parse import quote
+from opentelemetry.trace import Span, SpanKind, StatusCode, format_trace_id, format_span_id, TraceState
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.sdk.trace import ReadableSpan
-from sentry_sdk.consts import SPANSTATUS
-from sentry_sdk.tracing import get_span_status_from_http_code
-from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
-from urllib3.util import parse_url as urlparse
 
 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
+from sentry_sdk.tracing_utils import Baggage
+from sentry_sdk.integrations.opentelemetry.consts import SentrySpanAttribute
 
 from sentry_sdk._types import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from typing import Any, Optional, Mapping, Sequence, Union
+    from typing import Any, Optional, Mapping, Sequence, Union, ItemsView
+    from sentry_sdk._types import OtelExtractedSpanData
 
 
 GRPC_ERROR_MAP = {
@@ -87,7 +91,7 @@ def convert_to_otel_timestamp(time):
 
 
 def extract_span_data(span):
-    # type: (ReadableSpan) -> tuple[str, str, Optional[str], Optional[int], Optional[str]]
+    # type: (ReadableSpan) -> OtelExtractedSpanData
     op = span.name
     description = span.name
     status, http_status = extract_span_status(span)
@@ -125,7 +129,7 @@ def extract_span_data(span):
 
 
 def span_data_for_http_method(span):
-    # type: (ReadableSpan) -> tuple[str, str, Optional[str], Optional[int], Optional[str]]
+    # type: (ReadableSpan) -> OtelExtractedSpanData
     span_attributes = span.attributes or {}
 
     op = "http"
@@ -167,7 +171,7 @@ def span_data_for_http_method(span):
 
 
 def span_data_for_db_query(span):
-    # type: (ReadableSpan) -> tuple[str, str, Optional[str], Optional[int], Optional[str]]
+    # type: (ReadableSpan) -> OtelExtractedSpanData
     span_attributes = span.attributes or {}
 
     op = "db"
@@ -261,3 +265,65 @@ def extract_span_attributes(span, namespace):
             extracted_attrs[key] = value
 
     return extracted_attrs
+
+
+def get_trace_context(span, span_data=None):
+    # type: (ReadableSpan, Optional[OtelExtractedSpanData]) -> dict[str, Any]
+    if not span.context:
+        return {}
+
+    trace_id = format_trace_id(span.context.trace_id)
+    span_id = format_span_id(span.context.span_id)
+    parent_span_id = format_span_id(span.parent.span_id) if span.parent else None
+
+    if span_data is None:
+        span_data = extract_span_data(span)
+
+    (op, _, status, _, origin) = span_data
+
+    trace_context = {
+        "trace_id": trace_id,
+        "span_id": span_id,
+        "parent_span_id": parent_span_id,
+        "op": op,
+        "origin": origin or DEFAULT_SPAN_ORIGIN,
+        "status": status,
+    }  # type: dict[str, Any]
+
+    if span.attributes:
+        trace_context["data"] = dict(span.attributes)
+
+    trace_context["dynamic_sampling_context"] = dsc_from_trace_state(span.context.trace_state)
+
+    # TODO-neel-potel profiler thread_id, thread_name
+
+    return trace_context
+
+
+def trace_state_from_baggage(baggage):
+    # type: (Baggage) -> TraceState
+    items = []
+    for k, v in baggage.sentry_items.items():
+        key = Baggage.SENTRY_PREFIX + quote(k)
+        val = quote(str(v))
+        items.append((key, val))
+    return TraceState(items)
+
+
+def serialize_trace_state(trace_state):
+    # type: (TraceState) -> str
+    sentry_items = []
+    for k, v in trace_state.items():
+        if Baggage.SENTRY_PREFIX_REGEX.match(k):
+            sentry_items.append((k, v))
+    return ",".join(key + "=" + value for key, value in sentry_items)
+
+
+def dsc_from_trace_state(trace_state):
+    # type: (TraceState) -> dict[str, str]
+    dsc = {}
+    for k, v in trace_state.items():
+        if Baggage.SENTRY_PREFIX_REGEX.match(k):
+            key = re.sub(Baggage.SENTRY_PREFIX_REGEX, "", k)
+            dsc[key] = v
+    return dsc

sentry_sdk/tracing.py

@@ -1,7 +1,7 @@
 import uuid
 import random
 import time
-import warnings 
+import warnings
 from datetime import datetime, timedelta, timezone
 
 from opentelemetry import trace as otel_trace, context
@@ -1447,7 +1447,14 @@ def start_child(self, **kwargs):
 
     def iter_headers(self):
         # type: () -> Iterator[Tuple[str, str]]
-        pass
+        yield SENTRY_TRACE_HEADER_NAME, self.to_traceparent()
+
+        from sentry_sdk.integrations.opentelemetry.utils import (
+            serialize_trace_state,
+        )
+
+        trace_state  = self._otel_span.get_span_context().trace_state
+        yield BAGGAGE_HEADER_NAME, serialize_trace_state(trace_state)
 
     def to_traceparent(self):
         # type: () -> str
@@ -1466,6 +1473,7 @@ def to_traceparent(self):
 
     def to_baggage(self):
         # type: () -> Optional[Baggage]
+        # TODO-neel-potel head SDK populate baggage mess
         pass
 
     def set_tag(self, key, value):
@@ -1540,8 +1548,15 @@ def to_json(self):
         pass
 
     def get_trace_context(self):
-        # type: () -> Any
-        pass
+        # type: () -> dict[str, Any]
+        if not isinstance(self._otel_span, ReadableSpan):
+            return {}
+
+        from sentry_sdk.integrations.opentelemetry.utils import (
+            get_trace_context,
+        )
+
+        return get_trace_context(self._otel_span)
 
     def get_profile_context(self):
         # type: () -> Optional[ProfileContext]