dont round incoming sample_rand

Author: sentrivana

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -8,7 +8,6 @@
 import sentry_sdk
 from sentry_sdk.tracing_utils import (
     _generate_sample_rand,
-    _round_sample_rand,
     has_tracing_enabled,
 )
 from sentry_sdk.utils import is_valid_sample_rate, logger
@@ -87,12 +86,12 @@ def get_parent_sample_rand(parent_context, trace_id):
         if parent_sample_rand is None:
             return None
 
-        return _round_sample_rand(parent_sample_rand)
+        return Decimal(parent_sample_rand)
 
     return None
 
 
-def dropped_result(parent_span_context, attributes, sample_rate=None, sample_rand=None):
+def dropped_result(span_context, attributes, sample_rate=None, sample_rand=None):
     # type: (SpanContext, Attributes, Optional[float], Optional[Decimal]) -> SamplingResult
     """
     React to a span getting unsampled and return a DROP SamplingResult.
@@ -104,26 +103,11 @@ def dropped_result(parent_span_context, attributes, sample_rate=None, sample_ran
     See for more info about OTel sampling:
     https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.sampling.html
     """
-    # these will only be added the first time in a root span sampling decision
-    # if sample_rate or sample_rand is provided, they'll be updated in trace state
-    trace_state = parent_span_context.trace_state
-
-    if TRACESTATE_SAMPLED_KEY not in trace_state:
-        trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "false")
-    elif trace_state.get(TRACESTATE_SAMPLED_KEY) == "deferred":
-        trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "false")
-
-    if sample_rate is not None:
-        trace_state = trace_state.update(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
-
-    if sample_rand is not None:
-        trace_state = trace_state.update(
-            TRACESTATE_SAMPLE_RAND_KEY, f"{sample_rand:.6f}"  # noqa: E231
-        )
-
-    is_root_span = not (
-        parent_span_context.is_valid and not parent_span_context.is_remote
+    trace_state = _update_trace_state(
+        span_context, sampled=False, sample_rate=sample_rate, sample_rand=sample_rand
     )
+
+    is_root_span = not (span_context.is_valid and not span_context.is_remote)
     if is_root_span:
         # Tell Sentry why we dropped the transaction/root-span
         client = sentry_sdk.get_client()
@@ -156,28 +140,38 @@ def sampled_result(span_context, attributes, sample_rate=None, sample_rand=None)
     See for more info about OTel sampling:
     https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.sampling.html
     """
-    # these will only be added the first time in a root span sampling decision
-    # if sample_rate or sample_rand is provided, they'll be updated in trace state
+    trace_state = _update_trace_state(
+        span_context, sampled=True, sample_rate=sample_rate, sample_rand=sample_rand
+    )
+
+    return SamplingResult(
+        Decision.RECORD_AND_SAMPLE,
+        attributes=attributes,
+        trace_state=trace_state,
+    )
+
+
+def _update_trace_state(span_context, sampled, sample_rate=None, sample_rand=None):
+    # type: (..., bool, Optional[float], Optional[Decimal]) -> TraceState
     trace_state = span_context.trace_state
 
+    sampled = "true" if sampled else "false"
     if TRACESTATE_SAMPLED_KEY not in trace_state:
-        trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "true")
+        trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, sampled)
     elif trace_state.get(TRACESTATE_SAMPLED_KEY) == "deferred":
-        trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "true")
+        trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, sampled)
 
+    # sample_rate should always be updated to ensure complete traces
     if sample_rate is not None:
         trace_state = trace_state.update(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
 
+    # only add sample_rand if there isn't one
     if sample_rand is not None:
-        trace_state = trace_state.update(
+        trace_state = trace_state.add(
             TRACESTATE_SAMPLE_RAND_KEY, f"{sample_rand:.6f}"  # noqa: E231
         )
 
-    return SamplingResult(
-        Decision.RECORD_AND_SAMPLE,
-        attributes=attributes,
-        trace_state=trace_state,
-    )
+    return trace_state
 
 
 class SentrySampler(Sampler):
@@ -274,17 +268,15 @@ def should_sample(
             logger.warning(
                 f"[Tracing] Discarding {name} because of invalid sample rate."
             )
-            return dropped_result(
-                parent_span_context, attributes, sample_rand=sample_rand
-            )
+            return dropped_result(parent_span_context, attributes)
 
         # Down-sample in case of back pressure monitor says so
         if is_root_span and client.monitor:
             sample_rate /= 2**client.monitor.downsample_factor
             if client.monitor.downsample_factor > 0:
                 sample_rate_to_propagate = sample_rate
 
-        # Roll the dice on sample rate
+        # Compare sample_rand to sample_rate to make the final sampling decision
         sample_rate = float(cast("Union[bool, float, int]", sample_rate))
         sampled = sample_rand < sample_rate
 
@@ -293,14 +285,14 @@ def should_sample(
                 parent_span_context,
                 attributes,
                 sample_rate=sample_rate_to_propagate,
-                sample_rand=sample_rand,
+                sample_rand=None if sample_rand == parent_sample_rand else sample_rand,
             )
         else:
             return dropped_result(
                 parent_span_context,
                 attributes,
                 sample_rate=sample_rate_to_propagate,
-                sample_rand=sample_rand,
+                sample_rand=None if sample_rand == parent_sample_rand else sample_rand,
             )
 
     def get_description(self) -> str:

sentry_sdk/tracing_utils.py

@@ -783,7 +783,9 @@ def _generate_sample_rand(
     while sample_rand >= upper:
         sample_rand = rng.uniform(lower, upper)
 
-    return _round_sample_rand(sample_rand)
+    return Decimal(sample_rand).quantize(
+        Decimal("0.000001"), rounding=ROUND_DOWN, context=Context(prec=6)
+    )
 
 
 def _sample_rand_range(parent_sampled, sample_rate):
@@ -801,21 +803,6 @@ def _sample_rand_range(parent_sampled, sample_rate):
         return sample_rate, 1.0
 
 
-def _round_sample_rand(sample_rand):
-    # type: (Union[Decimal, float, str]) -> Optional[Decimal]
-    # Round down to exactly six decimal-digit precision.
-    # Setting the context is needed to avoid an InvalidOperation exception
-    # in case the user has changed the default precision.
-    try:
-        return Decimal(sample_rand).quantize(
-            Decimal("0.000001"), rounding=ROUND_DOWN, context=Context(prec=6)
-        )
-    except Exception as ex:
-        logger.debug(f"Failed to round sample_rand {sample_rand}: {ex}")
-
-    return None
-
-
 # Circular imports
 from sentry_sdk.tracing import (
     BAGGAGE_HEADER_NAME,

tests/tracing/test_integration_tests.py

@@ -227,7 +227,7 @@ def test_trace_propagation_meta_head_sdk(sentry_init):
 
     assert 'meta name="baggage"' in baggage
     baggage_content = re.findall('content="([^"]*)"', baggage)[0]
-    assert baggage_content == root_span.get_baggage().serialize()
+    assert SortedBaggage(baggage_content) == root_span.get_baggage().serialize()
 
 
 @pytest.mark.parametrize(

tests/tracing/test_sample_rand.py

@@ -86,19 +86,18 @@ def test_decimal_context(sentry_init, capture_events):
 @pytest.mark.parametrize(
     "incoming_sample_rand,expected_sample_rand",
     (
-        ("0.0100015", "0.010001"),
-        ("0.1", "0.100000"),
+        ("0.0100015", "0.0100015"),
+        ("0.1", "0.1"),
     ),
 )
-def test_sample_rand_precision(
+def test_unexpected_incoming_sample_rand_precision(
     sentry_init, capture_events, incoming_sample_rand, expected_sample_rand
 ):
     """
-    Test that incoming sample_rand is correctly interpreted and optionally rounded.
+    Test that incoming sample_rand is correctly interpreted even if it looks unexpected.
 
     We shouldn't be getting arbitrary precision sample_rand in incoming headers,
-    but if we do for some reason, check that we have this covered and we round
-    it correctly.
+    but if we do for some reason, check that we don't tamper with it.
     """
     sentry_init(traces_sample_rate=1.0)
     events = capture_events()
@@ -118,8 +117,30 @@ def test_sample_rand_precision(
     assert len(events) == 1
 
 
+@pytest.mark.parametrize(
+    "incoming_sample_rand",
+    ("abc", "null", "47"),
+)
+def test_invalid_incoming_sample_rand(sentry_init, incoming_sample_rand):
+    """Test that we handle malformed incoming sample_rand."""
+    sentry_init(traces_sample_rate=1.0)
+
+    baggage = f"sentry-sample_rand={incoming_sample_rand},sentry-trace_id=771a43a4192642f0b136d5159a501700"  # noqa: E231
+    sentry_trace = "771a43a4192642f0b136d5159a501700-1234567890abcdef"
+
+    with sentry_sdk.continue_trace(
+        {BAGGAGE_HEADER_NAME: baggage, SENTRY_TRACE_HEADER_NAME: sentry_trace}
+    ):
+        with sentry_sdk.start_span():
+            pass
+
+    # The behavior here is undefined since we got a broken incoming trace,
+    # so as long as the SDK doesn't produce an error we consider this
+    # testcase a success.
+
+
 @pytest.mark.parametrize("incoming", ((0.0, "true"), (1.0, "false")))
-def test_invalid_sampled_and_sample_rate(sentry_init, incoming):
+def test_invalid_incoming_sampled_and_sample_rate(sentry_init, incoming):
     """
     Test that we don't error out in case we can't generate a sample_rand that
     would respect the incoming sampled and sample_rate.