Update sample rate in DSC (#4018)

- update `sample_rate` in DSC after the initial sampling decision is made
- fix some typos

Part of #3999

Author: sentrivana

Diff for: sentry_sdk/scope.py

@@ -1043,6 +1043,18 @@ def start_transaction(
         sampling_context.update(custom_sampling_context)
         transaction._set_initial_sampling_decision(sampling_context=sampling_context)
 
+        # update the sample rate in the dsc
+        if transaction.sample_rate is not None:
+            propagation_context = self.get_active_propagation_context()
+            if propagation_context:
+                dsc = propagation_context.dynamic_sampling_context
+                if dsc is not None:
+                    dsc["sample_rate"] = str(transaction.sample_rate)
+            if transaction._baggage:
+                transaction._baggage.sentry_items["sample_rate"] = str(
+                    transaction.sample_rate
+                )
+
         if transaction.sampled:
             profile = Profile(
                 transaction.sampled, transaction._start_timestamp_monotonic_ns

Diff for: sentry_sdk/tracing.py

@@ -1070,7 +1070,6 @@ def get_baggage(self):
 
         The first time a new baggage with Sentry items is made,
         it will be frozen."""
-
         if not self._baggage or self._baggage.mutable:
             self._baggage = Baggage.populate_from_transaction(self)
 

Diff for: sentry_sdk/tracing_utils.py

@@ -392,7 +392,7 @@ def __init__(
         self.parent_sampled = parent_sampled
         """Boolean indicator if the parent span was sampled.
         Important when the parent span originated in an upstream service,
-        because we watn to sample the whole trace, or nothing from the trace."""
+        because we want to sample the whole trace, or nothing from the trace."""
 
         self.dynamic_sampling_context = dynamic_sampling_context
         """Data that is used for dynamic sampling decisions."""

Diff for: tests/integrations/stdlib/test_httplib.py

@@ -185,12 +185,13 @@ def test_outgoing_trace_headers(sentry_init, monkeypatch):
 
     sentry_init(traces_sample_rate=1.0)
 
-    headers = {}
-    headers["baggage"] = (
-        "other-vendor-value-1=foo;bar;baz, sentry-trace_id=771a43a4192642f0b136d5159a501700, "
-        "sentry-public_key=49d0f7386ad645858ae85020e393bef3, sentry-sample_rate=0.01337, "
-        "sentry-user_id=Am%C3%A9lie, other-vendor-value-2=foo;bar;"
-    )
+    headers = {
+        "baggage": (
+            "other-vendor-value-1=foo;bar;baz, sentry-trace_id=771a43a4192642f0b136d5159a501700, "
+            "sentry-public_key=49d0f7386ad645858ae85020e393bef3, sentry-sample_rate=0.01337, "
+            "sentry-user_id=Am%C3%A9lie, other-vendor-value-2=foo;bar;"
+        ),
+    }
 
     transaction = Transaction.continue_from_headers(headers)
 
@@ -220,7 +221,7 @@ def test_outgoing_trace_headers(sentry_init, monkeypatch):
         expected_outgoing_baggage = (
             "sentry-trace_id=771a43a4192642f0b136d5159a501700,"
             "sentry-public_key=49d0f7386ad645858ae85020e393bef3,"
-            "sentry-sample_rate=0.01337,"
+            "sentry-sample_rate=1.0,"
             "sentry-user_id=Am%C3%A9lie"
         )
 

Diff for: tests/test_dsc.py

@@ -8,6 +8,9 @@
 This is not tested in this file.
 """
 
+import random
+from unittest import mock
+
 import pytest
 
 import sentry_sdk
@@ -115,7 +118,85 @@ def test_dsc_continuation_of_trace(sentry_init, capture_envelopes):
 
     assert "sample_rate" in envelope_trace_header
     assert type(envelope_trace_header["sample_rate"]) == str
-    assert envelope_trace_header["sample_rate"] == "0.01337"
+    assert envelope_trace_header["sample_rate"] == "1.0"
+
+    assert "sampled" in envelope_trace_header
+    assert type(envelope_trace_header["sampled"]) == str
+    assert envelope_trace_header["sampled"] == "true"
+
+    assert "release" in envelope_trace_header
+    assert type(envelope_trace_header["release"]) == str
+    assert envelope_trace_header["release"] == "[email protected]"
+
+    assert "environment" in envelope_trace_header
+    assert type(envelope_trace_header["environment"]) == str
+    assert envelope_trace_header["environment"] == "bird"
+
+    assert "transaction" in envelope_trace_header
+    assert type(envelope_trace_header["transaction"]) == str
+    assert envelope_trace_header["transaction"] == "bar"
+
+
+def test_dsc_continuation_of_trace_sample_rate_changed_in_traces_sampler(
+    sentry_init, capture_envelopes
+):
+    """
+    Another service calls our service and passes tracing information to us.
+    Our service is continuing the trace, but modifies the sample rate.
+    The DSC propagated further should contain the updated sample rate.
+    """
+
+    def my_traces_sampler(sampling_context):
+        return 0.25
+
+    sentry_init(
+        dsn="https://[email protected]/12312012",
+        release="[email protected]",
+        environment="canary",
+        traces_sampler=my_traces_sampler,
+    )
+    envelopes = capture_envelopes()
+
+    # This is what the upstream service sends us
+    sentry_trace = "771a43a4192642f0b136d5159a501700-1234567890abcdef-1"
+    baggage = (
+        "other-vendor-value-1=foo;bar;baz, "
+        "sentry-trace_id=771a43a4192642f0b136d5159a501700, "
+        "sentry-public_key=frontendpublickey, "
+        "sentry-sample_rate=1.0, "
+        "sentry-sampled=true, "
+        "[email protected], "
+        "sentry-environment=bird, "
+        "sentry-transaction=bar, "
+        "other-vendor-value-2=foo;bar;"
+    )
+    incoming_http_headers = {
+        "HTTP_SENTRY_TRACE": sentry_trace,
+        "HTTP_BAGGAGE": baggage,
+    }
+
+    # We continue the incoming trace and start a new transaction
+    with mock.patch.object(random, "random", return_value=0.2):
+        transaction = sentry_sdk.continue_trace(incoming_http_headers)
+        with sentry_sdk.start_transaction(transaction, name="foo"):
+            pass
+
+    assert len(envelopes) == 1
+
+    transaction_envelope = envelopes[0]
+    envelope_trace_header = transaction_envelope.headers["trace"]
+
+    assert "trace_id" in envelope_trace_header
+    assert type(envelope_trace_header["trace_id"]) == str
+    assert envelope_trace_header["trace_id"] == "771a43a4192642f0b136d5159a501700"
+
+    assert "public_key" in envelope_trace_header
+    assert type(envelope_trace_header["public_key"]) == str
+    assert envelope_trace_header["public_key"] == "frontendpublickey"
+
+    assert "sample_rate" in envelope_trace_header
+    assert type(envelope_trace_header["sample_rate"]) == str
+    assert envelope_trace_header["sample_rate"] == "0.25"
 
     assert "sampled" in envelope_trace_header
     assert type(envelope_trace_header["sampled"]) == str

Diff for: tests/tracing/test_integration_tests.py

@@ -53,9 +53,11 @@ def test_basic(sentry_init, capture_events, sample_rate):
         assert not events
 
 
-@pytest.mark.parametrize("sampled", [True, False, None])
+@pytest.mark.parametrize("parent_sampled", [True, False, None])
 @pytest.mark.parametrize("sample_rate", [0.0, 1.0])
-def test_continue_from_headers(sentry_init, capture_envelopes, sampled, sample_rate):
+def test_continue_from_headers(
+    sentry_init, capture_envelopes, parent_sampled, sample_rate
+):
     """
     Ensure data is actually passed along via headers, and that they are read
     correctly.
@@ -66,7 +68,7 @@ def test_continue_from_headers(sentry_init, capture_envelopes, sampled, sample_r
     # make a parent transaction (normally this would be in a different service)
     with start_transaction(name="hi", sampled=True if sample_rate == 0 else None):
         with start_span() as old_span:
-            old_span.sampled = sampled
+            old_span.sampled = parent_sampled
             headers = dict(
                 sentry_sdk.get_current_scope().iter_trace_propagation_headers(old_span)
             )
@@ -81,7 +83,7 @@ def test_continue_from_headers(sentry_init, capture_envelopes, sampled, sample_r
     # child transaction, to prove that we can read 'sentry-trace' header data correctly
     child_transaction = Transaction.continue_from_headers(headers, name="WRONG")
     assert child_transaction is not None
-    assert child_transaction.parent_sampled == sampled
+    assert child_transaction.parent_sampled == parent_sampled
     assert child_transaction.trace_id == old_span.trace_id
     assert child_transaction.same_process_as_parent is False
     assert child_transaction.parent_span_id == old_span.span_id
@@ -106,8 +108,8 @@ def test_continue_from_headers(sentry_init, capture_envelopes, sampled, sample_r
         sentry_sdk.get_current_scope().transaction = "ho"
         capture_message("hello")
 
-    # in this case the child transaction won't be captured
-    if sampled is False or (sample_rate == 0 and sampled is None):
+    if parent_sampled is False or (sample_rate == 0 and parent_sampled is None):
+        # in this case the child transaction won't be captured
         trace1, message = envelopes
         message_payload = message.get_event()
         trace1_payload = trace1.get_transaction_event()
@@ -129,12 +131,17 @@ def test_continue_from_headers(sentry_init, capture_envelopes, sampled, sample_r
             == message_payload["contexts"]["trace"]["trace_id"]
         )
 
+        if parent_sampled is not None:
+            expected_sample_rate = str(float(parent_sampled))
+        else:
+            expected_sample_rate = str(sample_rate)
+
         assert trace2.headers["trace"] == baggage.dynamic_sampling_context()
         assert trace2.headers["trace"] == {
             "public_key": "49d0f7386ad645858ae85020e393bef3",
             "trace_id": "771a43a4192642f0b136d5159a501700",
             "user_id": "Amelie",
-            "sample_rate": "0.01337",
+            "sample_rate": expected_sample_rate,
         }
 
     assert message_payload["message"] == "hello"

Diff for: tests/tracing/test_sampling.py

@@ -198,20 +198,19 @@ def test_passes_parent_sampling_decision_in_sampling_context(
     transaction = Transaction.continue_from_headers(
         headers={"sentry-trace": sentry_trace_header}, name="dogpark"
     )
-    spy = mock.Mock(wraps=transaction)
-    start_transaction(transaction=spy)
 
-    # there's only one call (so index at 0) and kwargs are always last in a call
-    # tuple (so index at -1)
-    sampling_context = spy._set_initial_sampling_decision.mock_calls[0][-1][
-        "sampling_context"
-    ]
-    assert "parent_sampled" in sampling_context
-    # because we passed in a spy, attribute access requires unwrapping
-    assert sampling_context["parent_sampled"]._mock_wraps is parent_sampling_decision
+    def mock_set_initial_sampling_decision(_, sampling_context):
+        assert "parent_sampled" in sampling_context
+        assert sampling_context["parent_sampled"] is parent_sampling_decision
 
+    with mock.patch(
+        "sentry_sdk.tracing.Transaction._set_initial_sampling_decision",
+        mock_set_initial_sampling_decision,
+    ):
+        start_transaction(transaction=transaction)
 
-def test_passes_custom_samling_context_from_start_transaction_to_traces_sampler(
+
+def test_passes_custom_sampling_context_from_start_transaction_to_traces_sampler(
     sentry_init, DictionaryContaining  # noqa: N803
 ):
     traces_sampler = mock.Mock()