Baggage creation for head of trace

Author: sl0thentr0py

sentry_sdk/tracing.py

@@ -281,6 +281,10 @@ def continue_from_headers(
 
         if sentrytrace_kwargs is not None:
             kwargs.update(sentrytrace_kwargs)
+
+            # If there's an incoming sentry-trace but no incoming baggage header,
+            # for instance in traces coming from older SDKs,
+            # baggage will be empty and immutable and won't be populated as head SDK.
             baggage.freeze()
 
         kwargs.update(extract_tracestate_data(headers.get("tracestate")))
@@ -309,8 +313,8 @@ def iter_headers(self):
         if tracestate:
             yield "tracestate", tracestate
 
-        if self.containing_transaction and self.containing_transaction._baggage:
-            baggage = self.containing_transaction._baggage.serialize()
+        if self.containing_transaction:
+            baggage = self.containing_transaction.get_baggage().serialize()
             if baggage:
                 yield "baggage", baggage
 
@@ -513,11 +517,10 @@ def get_trace_context(self):
         if sentry_tracestate:
             rv["tracestate"] = sentry_tracestate
 
-        # TODO-neel populate fresh if head SDK
-        if self.containing_transaction and self.containing_transaction._baggage:
+        if self.containing_transaction:
             rv[
                 "dynamic_sampling_context"
-            ] = self.containing_transaction._baggage.dynamic_sampling_context()
+            ] = self.containing_transaction.get_baggage().dynamic_sampling_context()
 
         return rv
 
@@ -527,6 +530,7 @@ class Transaction(Span):
         "name",
         "source",
         "parent_sampled",
+        "sample_rate",
         # the sentry portion of the `tracestate` header used to transmit
         # correlation context for server-side dynamic sampling, of the form
         # `sentry=xxxxx`, where `xxxxx` is the base64-encoded json of the
@@ -562,6 +566,7 @@ def __init__(
         Span.__init__(self, **kwargs)
         self.name = name
         self.source = source
+        self.sample_rate = None  # type: Optional[float]
         self.parent_sampled = parent_sampled
         # if tracestate isn't inherited and set here, it will get set lazily,
         # either the first time an outgoing request needs it for a header or the
@@ -570,7 +575,7 @@ def __init__(
         self._third_party_tracestate = third_party_tracestate
         self._measurements = {}  # type: Dict[str, Any]
         self._profile = None  # type: Optional[Sampler]
-        self._baggage = baggage
+        self._baggage = baggage  # type: Optional[Baggage]
 
     def __repr__(self):
         # type: () -> str
@@ -708,6 +713,17 @@ def to_json(self):
 
         return rv
 
+    def get_baggage(self):
+        # type: () -> Baggage
+        """
+        The first time a new baggage with sentry items is made,
+        it will be frozen.
+        """
+        if not self._baggage or self._baggage.mutable:
+            self._baggage = Baggage.populate_from_transaction(self)
+
+        return self._baggage
+
     def _set_initial_sampling_decision(self, sampling_context):
         # type: (SamplingContext) -> None
         """
@@ -773,6 +789,9 @@ def _set_initial_sampling_decision(self, sampling_context):
             self.sampled = False
             return
 
+        # used to create baggage value for head SDKs in dynamic sampling
+        self.sample_rate = float(sample_rate)
+
         # if the function returned 0 (or false), or if `traces_sample_rate` is
         # 0, it's a sign the transaction should be dropped
         if not sample_rate:

sentry_sdk/tracing_utils.py

@@ -470,6 +470,46 @@ def from_incoming_header(cls, header):
 
         return Baggage(sentry_items, third_party_items, mutable)
 
+    @classmethod
+    def populate_from_transaction(cls, transaction):
+        # type: (Transaction) -> Baggage
+        """
+        Populate fresh baggage entry with sentry_items and make it immutable
+        if this is the head SDK which originates traces.
+        """
+
+        hub = transaction.hub or sentry_sdk.Hub.current
+        client = hub.client
+        sentry_items = {}  # type: Dict[str, str]
+
+        if not client:
+            return Baggage(sentry_items)
+
+        options = client.options or {}
+        user = (hub.scope and hub.scope._user) or {}
+
+        sentry_items["trace_id"] = transaction.trace_id
+
+        if options.get("environment"):
+            sentry_items["environment"] = options["environment"]
+
+        if options.get("release"):
+            sentry_items["release"] = options["release"]
+
+        if options.get("dsn"):
+            sentry_items["public_key"] = Dsn(options["dsn"]).public_key
+
+        if transaction.name and transaction.source != TRANSACTION_SOURCE_URL:
+            sentry_items["transaction"] = transaction.name
+
+        if user.get("segment"):
+            sentry_items["user_segment"] = user["segment"]
+
+        if transaction.sample_rate is not None:
+            sentry_items["sample_rate"] = str(transaction.sample_rate)
+
+        return Baggage(sentry_items, mutable=False)
+
     def freeze(self):
         # type: () -> None
         self.mutable = False
@@ -500,6 +540,7 @@ def serialize(self, include_third_party=False):
 
 
 # Circular imports
+from sentry_sdk.tracing import TRANSACTION_SOURCE_URL
 
 if MYPY:
-    from sentry_sdk.tracing import Span
+    from sentry_sdk.tracing import Span, Transaction