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

`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.

Depends on:
  - #4038 

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

@@ -1,5 +1,4 @@
 import uuid
-import random
 import warnings
 from datetime import datetime, timedelta, timezone
 from enum import Enum
@@ -477,6 +476,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
@@ -485,6 +486,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:
@@ -495,7 +498,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(
@@ -779,6 +784,7 @@ class Transaction(Span):
         "_profile",
         "_continuous_profile",
         "_baggage",
+        "_sample_rand",
     )
 
     def __init__(  # type: ignore[misc]
@@ -803,6 +809,14 @@ def __init__(  # type: ignore[misc]
         self._continuous_profile = None  # type: Optional[ContinuousProfile]
         self._baggage = baggage
 
+        baggage_sample_rand = (
+            None if self._baggage is None else self._baggage._sample_rand()
+        )
+        if baggage_sample_rand is not None:
+            self._sample_rand = baggage_sample_rand
+        else:
+            self._sample_rand = _generate_sample_rand(self.trace_id)
+
     def __repr__(self):
         # type: () -> str
         return (
@@ -1173,10 +1187,10 @@ def _set_initial_sampling_decision(self, sampling_context):
             self.sampled = False
             return
 
-        # Now we roll the dice. random.random is inclusive of 0, but not of 1,
+        # Now we roll the dice. self._sample_rand is inclusive of 0, but not of 1,
         # so strict < is safe here. In case sample_rate is a boolean, cast it
         # to a float (True becomes 1.0 and False becomes 0.0)
-        self.sampled = random.random() < self.sample_rate
+        self.sampled = self._sample_rand < self.sample_rate
 
         if self.sampled:
             logger.debug(
@@ -1333,6 +1347,7 @@ async def my_async_function():
     Baggage,
     EnvironHeaders,
     extract_sentrytrace_data,
+    _generate_sample_rand,
     has_tracing_enabled,
     maybe_create_breadcrumbs_from_span,
 )

sentry_sdk/tracing_utils.py

@@ -531,6 +531,14 @@ def _fill_sample_rand(self):
             f"{sample_rand:.6f}"  # noqa: E231
         )
 
+    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:
     """
@@ -553,8 +561,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
         """
@@ -577,6 +590,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
@@ -628,6 +645,7 @@ def populate_from_transaction(cls, transaction):
         options = client.options or {}
 
         sentry_items["trace_id"] = transaction.trace_id
+        sentry_items["sample_rand"] = str(transaction._sample_rand)
 
         if options.get("environment"):
             sentry_items["environment"] = options["environment"]
@@ -700,6 +718,20 @@ def strip_sentry_baggage(header):
             )
         )
 
+    def _sample_rand(self):
+        # type: () -> Optional[Decimal]
+        """Convenience method to get the sample_rand value from the sentry_items.
+
+        We validate the value and parse it as a Decimal before returning it. The value is considered
+        valid if it is a Decimal in the range [0, 1).
+        """
+        sample_rand = try_convert(Decimal, self.sentry_items.get("sample_rand"))
+
+        if sample_rand is not None and Decimal(0) <= sample_rand < Decimal(1):
+            return sample_rand
+
+        return None
+
     def __repr__(self):
         # type: () -> str
         return f'<Baggage "{self.serialize(include_third_party=True)}", mutable={self.mutable}>'

tests/integrations/aiohttp/test_aiohttp.py

@@ -626,18 +626,19 @@ async def handler(request):
 
     raw_server = await aiohttp_raw_server(handler)
 
-    with start_transaction(
-        name="/interactions/other-dogs/new-dog",
-        op="greeting.sniff",
-        trace_id="0123456789012345678901234567890",
-    ):
-        client = await aiohttp_client(raw_server)
-        resp = await client.get("/", headers={"bagGage": "custom=value"})
-
-        assert (
-            resp.request_info.headers["baggage"]
-            == "custom=value,sentry-trace_id=0123456789012345678901234567890,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0,sentry-sampled=true"
-        )
+    with mock.patch("sentry_sdk.tracing_utils.Random.uniform", return_value=0.5):
+        with start_transaction(
+            name="/interactions/other-dogs/new-dog",
+            op="greeting.sniff",
+            trace_id="0123456789012345678901234567890",
+        ):
+            client = await aiohttp_client(raw_server)
+            resp = await client.get("/", headers={"bagGage": "custom=value"})
+
+            assert (
+                resp.request_info.headers["baggage"]
+                == "custom=value,sentry-trace_id=0123456789012345678901234567890,sentry-sample_rand=0.500000,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0,sentry-sampled=true"
+            )
 
 
 @pytest.mark.asyncio

tests/integrations/celery/test_celery.py

@@ -509,22 +509,25 @@ def test_baggage_propagation(init_celery):
     def dummy_task(self, x, y):
         return _get_headers(self)
 
-    with start_transaction() as transaction:
-        result = dummy_task.apply_async(
-            args=(1, 0),
-            headers={"baggage": "custom=value"},
-        ).get()
-
-        assert sorted(result["baggage"].split(",")) == sorted(
-            [
-                "sentry-release=abcdef",
-                "sentry-trace_id={}".format(transaction.trace_id),
-                "sentry-environment=production",
-                "sentry-sample_rate=1.0",
-                "sentry-sampled=true",
-                "custom=value",
-            ]
-        )
+    # patch random.uniform to return a predictable sample_rand value
+    with mock.patch("sentry_sdk.tracing_utils.Random.uniform", return_value=0.5):
+        with start_transaction() as transaction:
+            result = dummy_task.apply_async(
+                args=(1, 0),
+                headers={"baggage": "custom=value"},
+            ).get()
+
+            assert sorted(result["baggage"].split(",")) == sorted(
+                [
+                    "sentry-release=abcdef",
+                    "sentry-trace_id={}".format(transaction.trace_id),
+                    "sentry-environment=production",
+                    "sentry-sample_rand=0.500000",
+                    "sentry-sample_rate=1.0",
+                    "sentry-sampled=true",
+                    "custom=value",
+                ]
+            )
 
 
 def test_sentry_propagate_traces_override(init_celery):

tests/integrations/httpx/test_httpx.py

@@ -170,30 +170,32 @@ def test_outgoing_trace_headers_append_to_baggage(
 
     url = "http://example.com/"
 
-    with start_transaction(
-        name="/interactions/other-dogs/new-dog",
-        op="greeting.sniff",
-        trace_id="01234567890123456789012345678901",
-    ) as transaction:
-        if asyncio.iscoroutinefunction(httpx_client.get):
-            response = asyncio.get_event_loop().run_until_complete(
-                httpx_client.get(url, headers={"baGGage": "custom=data"})
+    # patch random.uniform to return a predictable sample_rand value
+    with mock.patch("sentry_sdk.tracing_utils.Random.uniform", return_value=0.5):
+        with start_transaction(
+            name="/interactions/other-dogs/new-dog",
+            op="greeting.sniff",
+            trace_id="01234567890123456789012345678901",
+        ) as transaction:
+            if asyncio.iscoroutinefunction(httpx_client.get):
+                response = asyncio.get_event_loop().run_until_complete(
+                    httpx_client.get(url, headers={"baGGage": "custom=data"})
+                )
+            else:
+                response = httpx_client.get(url, headers={"baGGage": "custom=data"})
+
+            request_span = transaction._span_recorder.spans[-1]
+            assert response.request.headers[
+                "sentry-trace"
+            ] == "{trace_id}-{parent_span_id}-{sampled}".format(
+                trace_id=transaction.trace_id,
+                parent_span_id=request_span.span_id,
+                sampled=1,
+            )
+            assert (
+                response.request.headers["baggage"]
+                == "custom=data,sentry-trace_id=01234567890123456789012345678901,sentry-sample_rand=0.500000,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0,sentry-sampled=true"
             )
-        else:
-            response = httpx_client.get(url, headers={"baGGage": "custom=data"})
-
-        request_span = transaction._span_recorder.spans[-1]
-        assert response.request.headers[
-            "sentry-trace"
-        ] == "{trace_id}-{parent_span_id}-{sampled}".format(
-            trace_id=transaction.trace_id,
-            parent_span_id=request_span.span_id,
-            sampled=1,
-        )
-        assert (
-            response.request.headers["baggage"]
-            == "custom=data,sentry-trace_id=01234567890123456789012345678901,sentry-environment=production,sentry-release=d08ebdb9309e1b004c6f52202de58a09c2268e42,sentry-transaction=/interactions/other-dogs/new-dog,sentry-sample_rate=1.0,sentry-sampled=true"
-        )
 
 
 @pytest.mark.parametrize(

tests/integrations/stdlib/test_httplib.py

@@ -1,4 +1,3 @@
-import random
 from http.client import HTTPConnection, HTTPSConnection
 from socket import SocketIO
 from urllib.error import HTTPError
@@ -189,7 +188,7 @@ def test_outgoing_trace_headers(sentry_init, monkeypatch):
         "baggage": (
             "other-vendor-value-1=foo;bar;baz, sentry-trace_id=771a43a4192642f0b136d5159a501700, "
             "sentry-public_key=49d0f7386ad645858ae85020e393bef3, sentry-sample_rate=0.01337, "
-            "sentry-user_id=Am%C3%A9lie, other-vendor-value-2=foo;bar;"
+            "sentry-user_id=Am%C3%A9lie, sentry-sample_rand=0.132521102938283, other-vendor-value-2=foo;bar;"
         ),
     }
 
@@ -222,7 +221,8 @@ def test_outgoing_trace_headers(sentry_init, monkeypatch):
             "sentry-trace_id=771a43a4192642f0b136d5159a501700,"
             "sentry-public_key=49d0f7386ad645858ae85020e393bef3,"
             "sentry-sample_rate=1.0,"
-            "sentry-user_id=Am%C3%A9lie"
+            "sentry-user_id=Am%C3%A9lie,"
+            "sentry-sample_rand=0.132521102938283"
         )
 
         assert request_headers["baggage"] == expected_outgoing_baggage
@@ -235,11 +235,9 @@ def test_outgoing_trace_headers_head_sdk(sentry_init, monkeypatch):
     mock_send = mock.Mock()
     monkeypatch.setattr(HTTPSConnection, "send", mock_send)
 
-    # make sure transaction is always sampled
-    monkeypatch.setattr(random, "random", lambda: 0.1)
-
     sentry_init(traces_sample_rate=0.5, release="foo")
-    transaction = Transaction.continue_from_headers({})
+    with mock.patch("sentry_sdk.tracing_utils.Random.uniform", return_value=0.25):
+        transaction = Transaction.continue_from_headers({})
 
     with start_transaction(transaction=transaction, name="Head SDK tx") as transaction:
         HTTPSConnection("www.squirrelchasers.com").request("GET", "/top-chasers")
@@ -261,6 +259,7 @@ def test_outgoing_trace_headers_head_sdk(sentry_init, monkeypatch):
 
         expected_outgoing_baggage = (
             "sentry-trace_id=%s,"
+            "sentry-sample_rand=0.250000,"
             "sentry-environment=production,"
             "sentry-release=foo,"
             "sentry-sample_rate=0.5,"

tests/test_api.py

@@ -1,4 +1,6 @@
 import pytest
+
+import re
 from unittest import mock
 
 import sentry_sdk
@@ -95,10 +97,10 @@ def test_baggage_with_tracing_disabled(sentry_init):
 def test_baggage_with_tracing_enabled(sentry_init):
     sentry_init(traces_sample_rate=1.0, release="1.0.0", environment="dev")
     with start_transaction() as transaction:
-        expected_baggage = "sentry-trace_id={},sentry-environment=dev,sentry-release=1.0.0,sentry-sample_rate=1.0,sentry-sampled={}".format(
+        expected_baggage_re = r"^sentry-trace_id={},sentry-sample_rand=0\.\d{{6}},sentry-environment=dev,sentry-release=1\.0\.0,sentry-sample_rate=1\.0,sentry-sampled={}$".format(
             transaction.trace_id, "true" if transaction.sampled else "false"
         )
-        assert get_baggage() == expected_baggage
+        assert re.match(expected_baggage_re, get_baggage())
 
 
 @pytest.mark.forked

tests/test_dsc.py

@@ -8,7 +8,6 @@
 This is not tested in this file.
 """
 
-import random
 from unittest import mock
 
 import pytest
@@ -176,7 +175,7 @@ def my_traces_sampler(sampling_context):
     }
 
     # We continue the incoming trace and start a new transaction
-    with mock.patch.object(random, "random", return_value=0.2):
+    with mock.patch("sentry_sdk.tracing_utils.Random.uniform", return_value=0.125):
         transaction = sentry_sdk.continue_trace(incoming_http_headers)
         with sentry_sdk.start_transaction(transaction, name="foo"):
             pass