From e273b4f9c6482fb82cf4e385dc14391fd3ae19b1 Mon Sep 17 00:00:00 2001
From: Neel Shah <neel.shah@sentry.io>
Date: Wed, 11 Sep 2024 15:41:51 +0200
Subject: [PATCH] Store Baggage object on PropagationContext instead of DSC
 hash

---
 sentry_sdk/scope.py         | 35 +++++++++++++----------------------
 sentry_sdk/tracing_utils.py | 20 ++++++++++++--------
 2 files changed, 25 insertions(+), 30 deletions(-)

diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py
index f9ed9e4474..75a39ab9ba 100644
--- a/sentry_sdk/scope.py
+++ b/sentry_sdk/scope.py
@@ -476,18 +476,12 @@ def get_dynamic_sampling_context(self):
         # type: () -> Optional[Dict[str, str]]
         """
         Returns the Dynamic Sampling Context from the Propagation Context.
-        If not existing, creates a new one.
         """
-        if self._propagation_context is None:
-            return None
-
-        baggage = self.get_baggage()
-        if baggage is not None:
-            self._propagation_context.dynamic_sampling_context = (
-                baggage.dynamic_sampling_context()
-            )
-
-        return self._propagation_context.dynamic_sampling_context
+        return (
+            self._propagation_context.dynamic_sampling_context
+            if self._propagation_context
+            else None
+        )
 
     def get_traceparent(self, *args, **kwargs):
         # type: (Any, Any) -> Optional[str]
@@ -517,6 +511,7 @@ def get_baggage(self, *args, **kwargs):
         """
         Returns the Sentry "baggage" header containing trace information from the
         currently active span or the scopes Propagation Context.
+        If not existing, creates a new one.
         """
         client = self.get_client()
 
@@ -525,14 +520,11 @@ def get_baggage(self, *args, **kwargs):
             return self.span.to_baggage()
 
         # If this scope has a propagation context, return baggage from there
+        # populate a fresh one if it doesn't exist
         if self._propagation_context is not None:
-            dynamic_sampling_context = (
-                self._propagation_context.dynamic_sampling_context
-            )
-            if dynamic_sampling_context is None:
-                return Baggage.from_options(self)
-            else:
-                return Baggage(dynamic_sampling_context)
+            if self._propagation_context.baggage is None:
+                self._propagation_context.baggage = Baggage.from_options(self)
+            return self._propagation_context.baggage
 
         # Fall back to isolation scope's baggage. It always has one
         return self.get_isolation_scope().get_baggage()
@@ -594,10 +586,9 @@ def iter_headers(self):
             if traceparent is not None:
                 yield SENTRY_TRACE_HEADER_NAME, traceparent
 
-            dsc = self.get_dynamic_sampling_context()
-            if dsc is not None:
-                baggage = Baggage(dsc).serialize()
-                yield BAGGAGE_HEADER_NAME, baggage
+            baggage = self.get_baggage()
+            if baggage is not None:
+                yield BAGGAGE_HEADER_NAME, baggage.serialize()
 
     def iter_trace_propagation_headers(self, *args, **kwargs):
         # type: (Any, Any) -> Generator[Tuple[str, str], None, None]
diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py
index 8af17559e9..a221ee140e 100644
--- a/sentry_sdk/tracing_utils.py
+++ b/sentry_sdk/tracing_utils.py
@@ -337,7 +337,7 @@ class PropagationContext:
         "_span_id",
         "parent_span_id",
         "parent_sampled",
-        "dynamic_sampling_context",
+        "baggage",
     )
 
     def __init__(
@@ -346,7 +346,7 @@ def __init__(
         span_id=None,  # type: Optional[str]
         parent_span_id=None,  # type: Optional[str]
         parent_sampled=None,  # type: Optional[bool]
-        dynamic_sampling_context=None,  # type: Optional[Dict[str, str]]
+        baggage=None,  # type: Optional[Baggage]
     ):
         # type: (...) -> None
         self._trace_id = trace_id
@@ -364,8 +364,13 @@ def __init__(
         Important when the parent span originated in an upstream service,
         because we watn to sample the whole trace, or nothing from the trace."""
 
-        self.dynamic_sampling_context = dynamic_sampling_context
-        """Data that is used for dynamic sampling decisions."""
+        self.baggage = baggage
+        """Baggage object used for dynamic sampling decisions."""
+
+    @property
+    def dynamic_sampling_context(self):
+        # type: () -> Optional[Dict[str, str]]
+        return self.baggage.dynamic_sampling_context() if self.baggage else None
 
     @classmethod
     def from_incoming_data(cls, incoming_data):
@@ -376,9 +381,7 @@ def from_incoming_data(cls, incoming_data):
         baggage_header = normalized_data.get(BAGGAGE_HEADER_NAME)
         if baggage_header:
             propagation_context = PropagationContext()
-            propagation_context.dynamic_sampling_context = Baggage.from_incoming_header(
-                baggage_header
-            ).dynamic_sampling_context()
+            propagation_context.baggage = Baggage.from_incoming_header(baggage_header)
 
         sentry_trace_header = normalized_data.get(SENTRY_TRACE_HEADER_NAME)
         if sentry_trace_header:
@@ -431,11 +434,12 @@ def update(self, other_dict):
 
     def __repr__(self):
         # type: (...) -> str
-        return "<PropagationContext _trace_id={} _span_id={} parent_span_id={} parent_sampled={} dynamic_sampling_context={}>".format(
+        return "<PropagationContext _trace_id={} _span_id={} parent_span_id={} parent_sampled={} baggage={} dynamic_sampling_context={}>".format(
             self._trace_id,
             self._span_id,
             self.parent_span_id,
             self.parent_sampled,
+            self.baggage,
             self.dynamic_sampling_context,
         )