feat: Ensure PropagationContext has sample_rand

Author: szokeasaurusrex

sentry_sdk/tracing_utils.py

@@ -6,6 +6,7 @@
 from collections.abc import Mapping
 from datetime import timedelta
 from functools import wraps
+from random import Random
 from urllib.parse import quote, unquote
 import uuid
 
@@ -397,6 +398,8 @@ def __init__(
         self.dynamic_sampling_context = dynamic_sampling_context
         """Data that is used for dynamic sampling decisions."""
 
+        self._fill_sample_rand()
+
     @classmethod
     def from_incoming_data(cls, incoming_data):
         # type: (Dict[str, Any]) -> Optional[PropagationContext]
@@ -418,6 +421,8 @@ def from_incoming_data(cls, incoming_data):
                     propagation_context = PropagationContext()
                 propagation_context.update(sentrytrace_data)
 
+        propagation_context._fill_sample_rand()
+
         return propagation_context
 
     @property
@@ -426,6 +431,7 @@ def trace_id(self):
         """The trace id of the Sentry trace."""
         if not self._trace_id:
             self._trace_id = uuid.uuid4().hex
+            self._fill_sample_rand()
 
         return self._trace_id
 
@@ -469,6 +475,45 @@ def __repr__(self):
             self.dynamic_sampling_context,
         )
 
+    def _fill_sample_rand(self):
+        """
+        If the sample_rand is missing from the Dynamic Sampling Context (or invalid),
+        we generate it here.
+
+        We only generate a sample_rand if the trace_id is set.
+
+        If we have a parent_sampled value and a sample_rate in the DSC, we compute
+        a sample_rand value randomly in the range [0, sample_rate) if parent_sampled is True,
+        or in the range [sample_rate, 1) if parent_sampled is False. If either parent_sampled
+        or sample_rate is missing, we generate a random value in the range [0, 1).
+
+        The sample_rand is deterministically generated from the trace_id.
+        """
+        if self._trace_id is None:
+            # We only want to generate a sample_rand if the _trace_id is set.
+            return
+
+        # Ensure that the dynamic_sampling_context is a dict
+        self.dynamic_sampling_context = self.dynamic_sampling_context or {}
+
+        sample_rand = _try_float(self.dynamic_sampling_context.get("sample_rand"))
+        if sample_rand is not None and 0 <= sample_rand < 1:
+            # sample_rand is present and valid, so don't overwrite it
+            return
+
+        # Get a random value in [0, 1)
+        random_value = Random(self.trace_id).random()
+
+        # Get the sample rate and compute the transformation that will map the random value
+        # to the desired range: [0, 1), [0, sample_rate), or [sample_rate, 1).
+        sample_rate = _try_float(self.dynamic_sampling_context.get("sample_rate"))
+        factor, offset = _sample_rand_transormation(self.parent_sampled, sample_rate)
+
+        # Transform the random value to the desired range
+        self.dynamic_sampling_context["sample_rand"] = str(
+            random_value * factor + offset
+        )
+
 
 class Baggage:
     """
@@ -744,6 +789,35 @@ def get_current_span(scope=None):
     return current_span
 
 
+def _try_float(value):
+    # type: (object) -> Optional[float]
+    """Small utility to convert a value to a float, if possible."""
+    try:
+        return float(value)
+    except (ValueError, TypeError):
+        return None
+
+
+def _sample_rand_transormation(parent_sampled, sample_rate):
+    # type: (Optional[bool], Optional[float]) -> tuple[float, float]
+    """
+    Compute the factor and offset to scale and translate a random number in [0, 1) to
+    a range consistent with the parent_sampled and sample_rate values.
+
+    The return value is a tuple (factor, offset) such that, given random_value in [0, 1),
+    and new_value = random_value * factor + offset:
+      - new_value will be unchanged if either parent_sampled or sample_rate is None
+      - if parent_sampled and sample_rate are both set, new_value will be in [0, sample_rate)
+        if parent_sampled is True, or in [sample_rate, 1) if parent_sampled is False
+    """
+    if parent_sampled is None or sample_rate is None:
+        return 1.0, 0.0
+    elif parent_sampled is True:
+        return sample_rate, 0.0
+    else:  # parent_sampled is False
+        return 1.0 - sample_rate, sample_rate
+
+
 # Circular imports
 from sentry_sdk.tracing import (
     BAGGAGE_HEADER_NAME,

tests/test_propagationcontext.py

@@ -1,18 +1,27 @@
+import pytest
+
 from sentry_sdk.tracing_utils import PropagationContext
 
 
 def test_empty_context():
     ctx = PropagationContext()
 
+    # DSC is None before calling ctx.trace_id
+    assert ctx.dynamic_sampling_context is None
+
     assert ctx.trace_id is not None
     assert len(ctx.trace_id) == 32
 
+    # ctx.trace_id lazily computes the trace_id and therefor also the sample_rand
+    assert ctx.dynamic_sampling_context is not None
+    sample_rand = float(ctx.dynamic_sampling_context["sample_rand"])
+    assert 0 <= sample_rand < 1
+
     assert ctx.span_id is not None
     assert len(ctx.span_id) == 16
 
     assert ctx.parent_span_id is None
     assert ctx.parent_sampled is None
-    assert ctx.dynamic_sampling_context is None
 
 
 def test_context_with_values():
@@ -32,6 +41,8 @@ def test_context_with_values():
     assert ctx.parent_sampled
     assert ctx.dynamic_sampling_context == {
         "foo": "bar",
+        # sample_rand deterministically generated from trace_id
+        "sample_rand": "0.20286121767364262",
     }
 
 
@@ -81,3 +92,45 @@ def test_update():
     assert ctx.dynamic_sampling_context is None
 
     assert not hasattr(ctx, "foo")
+
+
+def test_existing_sample_rand_kept():
+    ctx = PropagationContext(
+        trace_id="00000000000000000000000000000000",
+        dynamic_sampling_context={"sample_rand": "0.5"},
+    )
+
+    # If sample_rand was regenerated, the value would be 0.8766381713144122 based on the trace_id
+    assert ctx.dynamic_sampling_context["sample_rand"] == "0.5"
+
+
+@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_sample_rand_filled(parent_sampled, sample_rate, expected_sample_rand):
+    """When continuing a trace, we want to fill in the sample_rand value if it's missing."""
+    dsc = {}
+    if sample_rate is not None:
+        dsc["sample_rate"] = sample_rate
+
+    ctx = PropagationContext(
+        trace_id="00000000000000000000000000000000",
+        parent_sampled=parent_sampled,
+        dynamic_sampling_context=dsc,
+    )
+
+    assert ctx.dynamic_sampling_context["sample_rand"] == expected_sample_rand