feat(tracing): Propagate sample_rand to transaction's baggage

`continue_trace` now propagates incoming `sample_rand` values to the transaction's baggage. Also, in the case where `sample_rand` is missing from the incoming trace and needs to be backfilled, this change introduces a mechanism for the backfilled value from the scope's propagation context to be propagated to the transaction's baggage.

The transaction still does not use the `sample_rand` for making sampling decisions; this PR only enables propagation. A future PR will add support for reading the incoming/backfilled `sample_rand` and for using this value to make sampling decisions.

Ref #3998

Author: szokeasaurusrex

sentry_sdk/scope.py

@@ -43,6 +43,7 @@
     logger,
 )
 
+import typing
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -1146,8 +1147,20 @@ def continue_trace(
         """
         self.generate_propagation_context(environ_or_headers)
 
+        # When we generate the propagation context, the sample_rand value is set
+        # if missing or invalid (we use the original value if it's valid).
+        # We want the transaction to use the same sample_rand value. Due to duplicated
+        # propagation logic in the transaction, we pass it in to avoid recomputing it
+        # in the transaction.
+        # TYPE SAFETY: self.generate_propagation_context() ensures that self._propagation_context
+        # is not None.
+        sample_rand = typing.cast(
+            PropagationContext, self._propagation_context
+        )._sample_rand()
+
         transaction = Transaction.continue_from_headers(
             normalize_incoming_data(environ_or_headers),
+            _sample_rand=sample_rand,
             op=op,
             origin=origin,
             name=name,

sentry_sdk/tracing.py

@@ -468,6 +468,8 @@ def continue_from_environ(
     def continue_from_headers(
         cls,
         headers,  # type: Mapping[str, str]
+        *,
+        _sample_rand=None,  # type: Optional[str]
         **kwargs,  # type: Any
     ):
         # type: (...) -> Transaction
@@ -476,6 +478,8 @@ def continue_from_headers(
         the ``sentry-trace`` and ``baggage`` headers).
 
         :param headers: The dictionary with the HTTP headers to pull information from.
+        :param _sample_rand: If provided, we override the sample_rand value from the
+            incoming headers with this value. (internal use only)
         """
         # TODO move this to the Transaction class
         if cls is Span:
@@ -486,7 +490,9 @@ def continue_from_headers(
 
         # TODO-neel move away from this kwargs stuff, it's confusing and opaque
         # make more explicit
-        baggage = Baggage.from_incoming_header(headers.get(BAGGAGE_HEADER_NAME))
+        baggage = Baggage.from_incoming_header(
+            headers.get(BAGGAGE_HEADER_NAME), _sample_rand=_sample_rand
+        )
         kwargs.update({BAGGAGE_HEADER_NAME: baggage})
 
         sentrytrace_kwargs = extract_sentrytrace_data(

sentry_sdk/tracing_utils.py

@@ -516,6 +516,14 @@ def _fill_sample_rand(self):
             # sample_rand value in this case.
             pass
 
+    def _sample_rand(self):
+        # type: () -> Optional[str]
+        """Convenience method to get the sample_rand value from the dynamic_sampling_context."""
+        if self.dynamic_sampling_context is None:
+            return None
+
+        return self.dynamic_sampling_context.get("sample_rand")
+
 
 class Baggage:
     """
@@ -538,8 +546,13 @@ def __init__(
         self.mutable = mutable
 
     @classmethod
-    def from_incoming_header(cls, header):
-        # type: (Optional[str]) -> Baggage
+    def from_incoming_header(
+        cls,
+        header,  # type: Optional[str]
+        *,
+        _sample_rand=None,  # type: Optional[str]
+    ):
+        # type: (...) -> Baggage
         """
         freeze if incoming header already has sentry baggage
         """
@@ -562,6 +575,10 @@ def from_incoming_header(cls, header):
                     else:
                         third_party_items += ("," if third_party_items else "") + item
 
+        if _sample_rand is not None:
+            sentry_items["sample_rand"] = str(_sample_rand)
+            mutable = False
+
         return Baggage(sentry_items, third_party_items, mutable)
 
     @classmethod

tests/tracing/test_sample_rand_propagation.py

@@ -0,0 +1,71 @@
+"""
+These tests exist to verify that Scope.continue_trace() correctly propagates the
+sample_rand value onto the transaction's baggage.
+
+We check both the case where there is an incoming sample_rand, as well as the case
+where we need to compute it because it is missing.
+"""
+
+import pytest
+
+import sentry_sdk
+
+
+def test_continue_trace_with_sample_rand():
+    """
+    Test that an incoming sample_rand is propagated onto the transaction's baggage.
+    """
+    headers = {
+        "sentry-trace": "00000000000000000000000000000000-0000000000000000-0",
+        "baggage": "sentry-sample_rand=0.1,sentry-sample_rate=0.5",
+    }
+
+    transaction = sentry_sdk.continue_trace(headers)
+    assert transaction.get_baggage().sentry_items["sample_rand"] == "0.1"
+
+
+@pytest.mark.parametrize(
+    ("parent_sampled", "sample_rate", "expected_sample_rand"),
+    (
+        (None, None, "0.8766381713144122"),
+        (None, "0.5", "0.8766381713144122"),
+        (False, None, "0.8766381713144122"),
+        (True, None, "0.8766381713144122"),
+        (False, "0.0", "0.8766381713144122"),
+        (False, "0.01", "0.8778717896012681"),
+        (True, "0.01", "0.008766381713144122"),
+        (False, "0.1", "0.888974354182971"),
+        (True, "0.1", "0.08766381713144122"),
+        (False, "0.5", "0.9383190856572061"),
+        (True, "0.5", "0.4383190856572061"),
+        (True, "1.0", "0.8766381713144122"),
+    ),
+)
+def test_continue_trace_missing_sample_rand(
+    parent_sampled, sample_rate, expected_sample_rand
+):
+    """
+    Test that a missing sample_rand is filled in onto the transaction's baggage. The sample_rand
+    is pseudorandomly generated based on the trace_id, so we assert the exact values that should
+    be generated.
+    """
+    headers = {
+        "sentry-trace": f"00000000000000000000000000000000-0000000000000000{sampled_flag(parent_sampled)}",
+        "baggage": f"sentry-sample_rate={sample_rate}",
+    }
+
+    transaction = sentry_sdk.continue_trace(headers)
+    assert transaction.get_baggage().sentry_items["sample_rand"] == expected_sample_rand
+
+
+def sampled_flag(sampled):
+    """
+    convenience function to get the sampled flag on the sentry-trace header, given a parent
+    sampling decision.
+    """
+    if sampled is None:
+        return ""
+    elif sampled is True:
+        return "-1"
+    else:
+        return "-0"