Skip to content

Commit

Permalink
feat(tracing): Propagate sample_rand to transaction's baggage
Browse files Browse the repository at this point in the history
`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
  • Loading branch information
szokeasaurusrex committed Feb 11, 2025
1 parent 0a0de06 commit bea45b9
Show file tree
Hide file tree
Showing 4 changed files with 103 additions and 3 deletions.
8 changes: 8 additions & 0 deletions sentry_sdk/scope.py
Original file line number Diff line number Diff line change
Expand Up @@ -1122,8 +1122,16 @@ 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.
sample_rand = 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,
Expand Down
8 changes: 7 additions & 1 deletion sentry_sdk/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -467,6 +467,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
Expand All @@ -475,6 +477,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.
"""
# TODO move this to the Transaction class
if cls is Span:
Expand All @@ -485,7 +489,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(
Expand Down
19 changes: 17 additions & 2 deletions sentry_sdk/tracing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -517,6 +517,14 @@ def _fill_sample_rand(self):
random_value * factor + offset
)

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:
"""
Expand All @@ -539,8 +547,8 @@ def __init__(
self.mutable = mutable

@classmethod
def from_incoming_header(cls, header):
# type: (Optional[str]) -> Baggage
def from_incoming_header(cls, header, *, sample_rand=None):
# type: (Optional[str], ..., Optional[str]) -> Baggage
"""
freeze if incoming header already has sentry baggage
"""
Expand All @@ -563,6 +571,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
Expand Down Expand Up @@ -686,6 +698,9 @@ def strip_sentry_baggage(header):
)
)

def __repr__(self):
return f"<Baggage {self.serialize()}>"


def should_propagate_trace(client, url):
# type: (sentry_sdk.client.BaseClient, str) -> bool
Expand Down
71 changes: 71 additions & 0 deletions tests/tracing/test_sample_rand_propagation.py
Original file line number Diff line number Diff line change
@@ -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"

0 comments on commit bea45b9

Please sign in to comment.