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 25, 2025
1 parent 2b3f4f7 commit 11cfacd
Show file tree
Hide file tree
Showing 4 changed files with 82 additions and 3 deletions.
13 changes: 13 additions & 0 deletions sentry_sdk/scope.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
logger,
)

import typing
from typing import TYPE_CHECKING

if TYPE_CHECKING:
Expand Down Expand Up @@ -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,
Expand Down
8 changes: 7 additions & 1 deletion sentry_sdk/tracing.py
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,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 @@ -485,6 +487,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:
Expand All @@ -495,7 +499,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
21 changes: 19 additions & 2 deletions sentry_sdk/tracing_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand All @@ -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
"""
Expand All @@ -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
Expand Down
43 changes: 43 additions & 0 deletions tests/tracing/test_sample_rand_propagation.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
"""
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.
"""

from unittest import mock
from unittest.mock import Mock

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"


def test_continue_trace_missing_sample_rand():
"""
Test that a missing sample_rand is filled in onto the transaction's baggage.
"""

headers = {
"sentry-trace": f"00000000000000000000000000000000-0000000000000000",
"baggage": f"sentry-placeholder=asdf",
}

mock_uniform = Mock(return_value=0.5)

with mock.patch("sentry_sdk.tracing_utils.Random.uniform", mock_uniform):
transaction = sentry_sdk.continue_trace(headers)

assert transaction.get_baggage().sentry_items["sample_rand"] == "0.500000"

0 comments on commit 11cfacd

Please sign in to comment.