Port sample_rate update to potel-base

sentry_sdk/integrations/opentelemetry/sampler.py

@@ -50,18 +50,39 @@ def get_parent_sampled(parent_context, trace_id):
     return None
 
 
+def get_parent_sample_rate(parent_context, trace_id):
+    # type: (Optional[SpanContext], int) -> Optional[float]
+    if parent_context is None:
+        return None
+
+    is_span_context_valid = parent_context is not None and parent_context.is_valid
+
+    if is_span_context_valid and parent_context.trace_id == trace_id:
+        parent_sample_rate = parent_context.trace_state.get(TRACESTATE_SAMPLE_RATE_KEY)
+        if parent_sample_rate is None:
+            return None
+
+        try:
+            return float(parent_sample_rate)
+        except Exception:
+            return None
+
+    return None
+
+
 def dropped_result(parent_span_context, attributes, sample_rate=None):
     # type: (SpanContext, Attributes, Optional[float]) -> SamplingResult
     # these will only be added the first time in a root span sampling decision
+    # if sample_rate is provided, it'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 and TRACESTATE_SAMPLE_RATE_KEY not in trace_state:
-        trace_state = trace_state.add(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
+    if sample_rate is not None:
+        trace_state = trace_state.update(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
 
     is_root_span = not (
         parent_span_context.is_valid and not parent_span_context.is_remote
@@ -88,17 +109,18 @@ def dropped_result(parent_span_context, attributes, sample_rate=None):
 
 
 def sampled_result(span_context, attributes, sample_rate):
-    # type: (SpanContext, Attributes, float) -> SamplingResult
+    # type: (SpanContext, Attributes, Optional[float]) -> SamplingResult
     # these will only be added the first time in a root span sampling decision
+    # if sample_rate is provided, it'll be updated in trace state
     trace_state = span_context.trace_state
 
     if TRACESTATE_SAMPLED_KEY not in trace_state:
         trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "true")
     elif trace_state.get(TRACESTATE_SAMPLED_KEY) == "deferred":
         trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "true")
 
-    if TRACESTATE_SAMPLE_RATE_KEY not in trace_state:
-        trace_state = trace_state.add(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
+    if sample_rate is not None:
+        trace_state = trace_state.update(TRACESTATE_SAMPLE_RATE_KEY, str(sample_rate))
 
     return SamplingResult(
         Decision.RECORD_AND_SAMPLE,
@@ -142,9 +164,13 @@ def should_sample(
             if is_root_span:
                 sample_rate = float(custom_sampled)
                 if sample_rate > 0:
-                    return sampled_result(parent_span_context, attributes, sample_rate)
+                    return sampled_result(
+                        parent_span_context, attributes, sample_rate=sample_rate
+                    )
                 else:
-                    return dropped_result(parent_span_context, attributes)
+                    return dropped_result(
+                        parent_span_context, attributes, sample_rate=sample_rate
+                    )
             else:
                 logger.debug(
                     f"[Tracing] Ignoring sampled param for non-root span {name}"
@@ -154,19 +180,27 @@ def should_sample(
         # Traces_sampler is responsible to check parent sampled to have full transactions.
         has_traces_sampler = callable(client.options.get("traces_sampler"))
 
+        sample_rate_to_propagate = None
+
         if is_root_span and has_traces_sampler:
             sampling_context = create_sampling_context(
                 name, attributes, parent_span_context, trace_id
             )
             sample_rate = client.options["traces_sampler"](sampling_context)
+            sample_rate_to_propagate = sample_rate
         else:
             # Check if there is a parent with a sampling decision
             parent_sampled = get_parent_sampled(parent_span_context, trace_id)
+            parent_sample_rate = get_parent_sample_rate(parent_span_context, trace_id)
             if parent_sampled is not None:
-                sample_rate = parent_sampled
+                sample_rate = bool(parent_sampled)
+                sample_rate_to_propagate = (
+                    parent_sample_rate if parent_sample_rate else sample_rate
+                )
             else:
                 # Check if there is a traces_sample_rate
                 sample_rate = client.options.get("traces_sample_rate")
+                sample_rate_to_propagate = sample_rate
 
         # If the sample rate is invalid, drop the span
         if not is_valid_sample_rate(sample_rate, source=self.__class__.__name__):
@@ -178,15 +212,21 @@ def should_sample(
         # 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
         sample_rate = float(cast("Union[bool, float, int]", sample_rate))
         sampled = random.random() < sample_rate
 
         if sampled:
-            return sampled_result(parent_span_context, attributes, sample_rate)
+            return sampled_result(
+                parent_span_context, attributes, sample_rate=sample_rate_to_propagate
+            )
         else:
-            return dropped_result(parent_span_context, attributes, sample_rate)
+            return dropped_result(
+                parent_span_context, attributes, sample_rate=sample_rate_to_propagate
+            )
 
     def get_description(self) -> str:
         return self.__class__.__name__

sentry_sdk/integrations/opentelemetry/scope.py

@@ -123,7 +123,7 @@ def _incoming_otel_span_context(self):
         # for twp to work, we also need to consider deferred sampling when the sampling
         # flag is not present, so the above TraceFlags are not sufficient
         if self._propagation_context.parent_sampled is None:
-            trace_state = trace_state.add(TRACESTATE_SAMPLED_KEY, "deferred")
+            trace_state = trace_state.update(TRACESTATE_SAMPLED_KEY, "deferred")
 
         span_context = SpanContext(
             trace_id=int(self._propagation_context.trace_id, 16),

tests/test_dsc.py

@@ -8,6 +8,7 @@
 This is not tested in this file.
 """
 
+import random
 from unittest import mock
 
 import pytest
@@ -117,7 +118,7 @@ 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"] == "1.0"
+    assert envelope_trace_header["sample_rate"] == "0.01337"
 
     assert "sampled" in envelope_trace_header
     assert type(envelope_trace_header["sampled"]) == str
@@ -137,7 +138,7 @@ def test_dsc_continuation_of_trace(sentry_init, capture_envelopes):
 
 
 def test_dsc_continuation_of_trace_sample_rate_changed_in_traces_sampler(
-    sentry_init, capture_envelopes
+    sentry_init, capture_envelopes, monkeypatch
 ):
     """
     Another service calls our service and passes tracing information to us.
@@ -175,10 +176,10 @@ def my_traces_sampler(sampling_context):
     }
 
     # We continue the incoming trace and start a new transaction
-    with mock.patch("sentry_sdk.tracing_utils.Random.uniform", return_value=0.125):
-        with sentry_sdk.continue_trace(incoming_http_headers):
-            with sentry_sdk.start_span(name="foo"):
-                pass
+    monkeypatch.setattr(random, "random", lambda: 0.125)
+    with sentry_sdk.continue_trace(incoming_http_headers):
+        with sentry_sdk.start_span(name="foo"):
+            pass
 
     assert len(envelopes) == 1
 
@@ -214,6 +215,212 @@ def my_traces_sampler(sampling_context):
     assert envelope_trace_header["transaction"] == "bar"
 
 
+@pytest.mark.parametrize(
+    "test_data, expected_sample_rate, expected_sampled",
+    [
+        # Test data:
+        # "incoming_sample_rate":
+        #       The "sentry-sample_rate" in the incoming `baggage` header.
+        # "incoming_sampled":
+        #       The "sentry-sampled" in the incoming `baggage` header.
+        # "sentry_trace_header_parent_sampled":
+        #       The number at the end in the `sentry-trace` header, called "parent_sampled".
+        # "use_local_traces_sampler":
+        #       Whether the local traces sampler is used.
+        # "local_traces_sampler_result":
+        #       The result of the local traces sampler.
+        # "local_traces_sample_rate":
+        #       The `traces_sample_rate` setting in the local `sentry_init` call.
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "true",
+                "sentry_trace_header_parent_sampled": 1,
+                "use_local_traces_sampler": False,
+                "local_traces_sampler_result": None,
+                "local_traces_sample_rate": 0.7,
+            },
+            1.0,  # expected_sample_rate
+            "true",  # expected_sampled
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "true",
+                "sentry_trace_header_parent_sampled": 1,
+                "use_local_traces_sampler": True,
+                "local_traces_sampler_result": 0.5,
+                "local_traces_sample_rate": 0.7,
+            },
+            0.5,  # expected_sample_rate
+            "true",  # expected_sampled
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "false",
+                "sentry_trace_header_parent_sampled": 0,
+                "use_local_traces_sampler": False,
+                "local_traces_sampler_result": None,
+                "local_traces_sample_rate": 0.7,
+            },
+            None,  # expected_sample_rate
+            "tracing-disabled-no-transactions-should-be-sent",  # expected_sampled (because the parent sampled is 0)
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "false",
+                "sentry_trace_header_parent_sampled": 0,
+                "use_local_traces_sampler": True,
+                "local_traces_sampler_result": 0.5,
+                "local_traces_sample_rate": 0.7,
+            },
+            0.5,  # expected_sample_rate
+            "false",  # expected_sampled (traces sampler can override parent sampled)
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "true",
+                "sentry_trace_header_parent_sampled": 1,
+                "use_local_traces_sampler": False,
+                "local_traces_sampler_result": None,
+                "local_traces_sample_rate": None,
+            },
+            None,  # expected_sample_rate
+            "tracing-disabled-no-transactions-should-be-sent",  # expected_sampled (traces_sample_rate=None disables all transaction creation)
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "true",
+                "sentry_trace_header_parent_sampled": 1,
+                "use_local_traces_sampler": True,
+                "local_traces_sampler_result": 0.5,
+                "local_traces_sample_rate": None,
+            },
+            0.5,  # expected_sample_rate
+            "true",  # expected_sampled (traces sampler overrides the traces_sample_rate setting, so transactions are created)
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "false",
+                "sentry_trace_header_parent_sampled": 0,
+                "use_local_traces_sampler": False,
+                "local_traces_sampler_result": None,
+                "local_traces_sample_rate": None,
+            },
+            None,  # expected_sample_rate
+            "tracing-disabled-no-transactions-should-be-sent",  # expected_sampled (traces_sample_rate=None disables all transaction creation)
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": "false",
+                "sentry_trace_header_parent_sampled": 0,
+                "use_local_traces_sampler": True,
+                "local_traces_sampler_result": 0.5,
+                "local_traces_sample_rate": None,
+            },
+            0.5,  # expected_sample_rate
+            "false",  # expected_sampled
+        ),
+        (
+            {
+                "incoming_sample_rate": 1.0,
+                "incoming_sampled": None,
+                "sentry_trace_header_parent_sampled": None,
+                "use_local_traces_sampler": False,
+                "local_traces_sampler_result": 0.5,
+                "local_traces_sample_rate": 0.7,
+            },
+            0.7,  # expected_sample_rate
+            "true",  # expected_sampled
+        ),
+    ],
+    ids=(
+        "1 traces_sample_rate does not override incoming",
+        "2 traces_sampler overrides incoming",
+        "3 traces_sample_rate does not overrides incoming sample rate or parent (incoming not sampled)",
+        "4 traces_sampler overrides incoming (incoming not sampled)",
+        "5 forwarding incoming (traces_sample_rate not set)",
+        "6 traces_sampler overrides incoming  (traces_sample_rate not set)",
+        "7 forwarding incoming (traces_sample_rate not set) (incoming not sampled)",
+        "8 traces_sampler overrides incoming  (traces_sample_rate not set) (incoming not sampled)",
+        "9 traces_sample_rate overrides incoming (upstream deferred sampling decision)",
+    ),
+)
+def test_dsc_sample_rate_change(
+    sentry_init,
+    capture_envelopes,
+    test_data,
+    expected_sample_rate,
+    expected_sampled,
+):
+    """
+    Another service calls our service and passes tracing information to us.
+    Our service is continuing the trace, but modifies the sample rate.
+    The DSC in transaction envelopes should contain the updated sample rate.
+    """
+
+    def my_traces_sampler(sampling_context):
+        return test_data["local_traces_sampler_result"]
+
+    init_kwargs = {
+        "dsn": "https://[email protected]/12312012",
+        "release": "[email protected]",
+        "environment": "canary",
+    }
+
+    if test_data["local_traces_sample_rate"]:
+        init_kwargs["traces_sample_rate"] = test_data["local_traces_sample_rate"]
+
+    if test_data["use_local_traces_sampler"]:
+        init_kwargs["traces_sampler"] = my_traces_sampler
+
+    sentry_init(**init_kwargs)
+    envelopes = capture_envelopes()
+
+    # This is what the upstream service sends us
+    incoming_trace_id = "771a43a4192642f0b136d5159a501700"
+    if test_data["sentry_trace_header_parent_sampled"] is None:
+        sentry_trace = f"{incoming_trace_id}-1234567890abcdef"
+    else:
+        sentry_trace = f"{incoming_trace_id}-1234567890abcdef-{test_data['sentry_trace_header_parent_sampled']}"
+
+    baggage = (
+        f"sentry-trace_id={incoming_trace_id}, "
+        f"sentry-sample_rate={str(test_data['incoming_sample_rate'])}, "
+        f"sentry-sampled={test_data['incoming_sampled']}, "
+        "sentry-public_key=frontendpublickey, "
+        "[email protected], "
+        "sentry-environment=prod, "
+        "sentry-transaction=foo, "
+    )
+    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):
+        with sentry_sdk.continue_trace(incoming_http_headers):
+            with sentry_sdk.start_span(name="foo"):
+                pass
+
+    if expected_sampled == "tracing-disabled-no-transactions-should-be-sent":
+        assert len(envelopes) == 0
+    else:
+        transaction_envelope = envelopes[0]
+        dsc_in_envelope_header = transaction_envelope.headers["trace"]
+
+        assert dsc_in_envelope_header["sample_rate"] == str(expected_sample_rate)
+        assert dsc_in_envelope_header["sampled"] == str(expected_sampled).lower()
+        assert dsc_in_envelope_header["trace_id"] == incoming_trace_id
+
+
 def test_dsc_issue(sentry_init, capture_envelopes):
     """
     Our service is a standalone service that does not have tracing enabled. Just uses Sentry for error reporting.