From b400fb175698bf1d614c1b577905e879e9f4d6e5 Mon Sep 17 00:00:00 2001 From: Daniel Szoke Date: Tue, 11 Feb 2025 19:49:52 +0100 Subject: [PATCH] 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 --- sentry_sdk/scope.py | 13 ++++ sentry_sdk/tracing.py | 8 ++- sentry_sdk/tracing_utils.py | 21 +++++- tests/tracing/test_sample_rand_propagation.py | 71 +++++++++++++++++++ 4 files changed, 110 insertions(+), 3 deletions(-) create mode 100644 tests/tracing/test_sample_rand_propagation.py diff --git a/sentry_sdk/scope.py b/sentry_sdk/scope.py index 4e3bb87489..e425031766 100644 --- a/sentry_sdk/scope.py +++ b/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, diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py index 9d50d38963..b99469ba2b 100644 --- a/sentry_sdk/tracing.py +++ b/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( diff --git a/sentry_sdk/tracing_utils.py b/sentry_sdk/tracing_utils.py index 38934dd211..f5a0afff9e 100644 --- a/sentry_sdk/tracing_utils.py +++ b/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 diff --git a/tests/tracing/test_sample_rand_propagation.py b/tests/tracing/test_sample_rand_propagation.py new file mode 100644 index 0000000000..9732f1a9d4 --- /dev/null +++ b/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"