Shuffle around where sampling happens

Author: k-fish

sentry_sdk/client.py

@@ -31,6 +31,7 @@
 )
 from sentry_sdk.serializer import serialize
 from sentry_sdk.tracing import trace
+from sentry_sdk.tracing_utils import _generate_sample_rand
 from sentry_sdk.transport import BaseHttpTransport, make_transport
 from sentry_sdk.consts import (
     SPANDATA,
@@ -181,7 +182,9 @@ class BaseClient:
 
     def __init__(self, options=None):
         # type: (Optional[Dict[str, Any]]) -> None
-        self.options = options if options is not None else DEFAULT_OPTIONS  # type: Dict[str, Any]
+        self.options = (
+            options if options is not None else DEFAULT_OPTIONS
+        )  # type: Dict[str, Any]
 
         self.transport = None  # type: Optional[Transport]
         self.monitor = None  # type: Optional[Monitor]
@@ -614,7 +617,9 @@ def _prepare_event(
                 event_scrubber.scrub_event(event)
 
         if scope is not None and scope._gen_ai_original_message_count:
-            spans = event.get("spans", [])  # type: List[Dict[str, Any]] | AnnotatedValue
+            spans = event.get(
+                "spans", []
+            )  # type: List[Dict[str, Any]] | AnnotatedValue
             if isinstance(spans, list):
                 for span in spans:
                     span_id = span.get("span_id", None)
@@ -1000,6 +1005,50 @@ def _capture_metric(self, metric):
         current_scope = sentry_sdk.get_current_scope()
         isolation_scope = sentry_sdk.get_isolation_scope()
 
+        # Determine trace_id and span_id using the same logic as the original metrics.py
+        trace_id = None
+        span_id = None
+        if current_scope.span is not None:
+            trace_id = current_scope.span.trace_id
+            span_id = current_scope.span.span_id
+        elif current_scope._propagation_context is not None:
+            trace_id = current_scope._propagation_context.trace_id
+            span_id = current_scope._propagation_context.span_id
+
+        sample_rate = metric["attributes"].get("sentry.client_sample_rate")
+        if sample_rate is not None:
+            sample_rate = float(sample_rate)
+            
+            # Always validate sample_rate range, regardless of trace context
+            if sample_rate <= 0.0 or sample_rate > 1.0:
+                if self.transport is not None:
+                    self.transport.record_lost_event(
+                        "invalid_sample_rate",
+                        data_category="trace_metric",
+                        quantity=1,
+                    )
+                return
+            
+            # If there's no trace context, remove the sample_rate attribute and continue
+            if trace_id is None:
+                del metric["attributes"]["sentry.client_sample_rate"]
+            else:
+                # There is a trace context, apply sampling logic
+                if sample_rate < 1.0:
+                    sample_rand = _generate_sample_rand(trace_id)
+                    if sample_rand >= sample_rate:
+                        if self.transport is not None:
+                            self.transport.record_lost_event(
+                                "sample_rate",
+                                data_category="trace_metric",
+                                quantity=1,
+                            )
+                        return
+                
+                # If sample_rate is 1.0, remove the attribute as it's implied
+                if sample_rate == 1.0:
+                    del metric["attributes"]["sentry.client_sample_rate"]
+
         metric["attributes"]["sentry.sdk.name"] = SDK_INFO["name"]
         metric["attributes"]["sentry.sdk.version"] = SDK_INFO["version"]
 
@@ -1011,10 +1060,6 @@ def _capture_metric(self, metric):
         if release is not None and "sentry.release" not in metric["attributes"]:
             metric["attributes"]["sentry.release"] = release
 
-        trace_context = current_scope.get_trace_context()
-        trace_id = trace_context.get("trace_id")
-        span_id = trace_context.get("span_id")
-
         metric["trace_id"] = trace_id or "00000000-0000-0000-0000-000000000000"
         if span_id is not None:
             metric["span_id"] = span_id

sentry_sdk/metrics.py

@@ -8,7 +8,6 @@
 
 import sentry_sdk
 from sentry_sdk.utils import safe_repr
-from sentry_sdk.tracing_utils import _generate_sample_rand
 
 if TYPE_CHECKING:
     from typing import Any, Optional, Union
@@ -41,35 +40,7 @@ def _capture_metric(
             )
 
     if sample_rate is not None:
-        if sample_rate <= 0.0 or sample_rate > 1.0:
-            if client.transport is not None:
-                client.transport.record_lost_event(
-                    "invalid_sample_rate",
-                    data_category="trace_metric",
-                    quantity=1,
-                )
-            return
-
-        trace_id = None
-        scope = sentry_sdk.get_current_scope()
-        if scope.span is not None:
-            trace_id = scope.span.trace_id
-        elif scope._propagation_context is not None:
-            trace_id = scope._propagation_context.trace_id
-
-        if trace_id is not None and sample_rate < 1.0:
-            sample_rand = _generate_sample_rand(trace_id)
-            if sample_rand >= sample_rate:
-                if client.transport is not None:
-                    client.transport.record_lost_event(
-                        "sample_rate",
-                        data_category="trace_metric",
-                        quantity=1,
-                    )
-                return
-
-        if sample_rate != 1.0 and trace_id is not None:
-            attrs["sentry.client_sample_rate"] = sample_rate
+        attrs["sentry.client_sample_rate"] = sample_rate
 
     metric = {
         "timestamp": time.time(),

tests/test_metrics.py

@@ -344,56 +344,6 @@ def test_metrics_no_sample_rate(sentry_init, capture_envelopes):
     assert "sentry.client_sample_rate" not in metrics[0]["attributes"]
 
 
-def test_metrics_sample_rate_no_trace_context(sentry_init, capture_envelopes):
-    """Test sentry.client_sample_rate not set when there's no trace context."""
-    sentry_init()
-    envelopes = capture_envelopes()
-
-    # Send metrics with sample_rate but without any active transaction/span
-    sentry_sdk.metrics.count("test.counter", 1, sample_rate=0.5)
-    sentry_sdk.metrics.gauge("test.gauge", 42, sample_rate=0.8)
-
-    get_client().flush()
-    metrics = envelopes_to_metrics(envelopes)
-
-    assert len(metrics) == 2
-
-    # When there's no trace context, no sampling is performed,
-    # so sentry.client_sample_rate should not be set
-    assert "sentry.client_sample_rate" not in metrics[0]["attributes"]
-    assert "sentry.client_sample_rate" not in metrics[1]["attributes"]
-
-
-def test_metrics_sample_rate_with_trace_context(
-    sentry_init, capture_envelopes, monkeypatch
-):
-    """Test sentry.client_sample_rate is set when there's a trace context."""
-    sentry_init(traces_sample_rate=1.0)
-    envelopes = capture_envelopes()
-
-    # Mock the random sampling to ensure all metrics are included
-    with mock.patch(
-        "sentry_sdk.tracing_utils.Random.randrange",
-        return_value=0,  # Always sample (0 < any positive sample_rate)
-    ):
-        # Send metrics with sample_rate within an active transaction
-        with sentry_sdk.start_transaction() as _:
-            sentry_sdk.metrics.count("test.counter", 1, sample_rate=0.5)
-            sentry_sdk.metrics.gauge("test.gauge", 42, sample_rate=0.8)
-            sentry_sdk.metrics.distribution("test.distribution", 200, sample_rate=1.0)
-
-    get_client().flush()
-    metrics = envelopes_to_metrics(envelopes)
-
-    assert len(metrics) == 3
-
-    # When there's a trace context and sample_rate < 1.0, set attribute
-    assert metrics[0]["attributes"]["sentry.client_sample_rate"] == 0.5
-    assert metrics[1]["attributes"]["sentry.client_sample_rate"] == 0.8
-    # sample_rate=1.0 doesn't need the attribute
-    assert "sentry.client_sample_rate" not in metrics[2]["attributes"]
-
-
 @pytest.mark.parametrize("sample_rand", (0.0, 0.25, 0.5, 0.75))
 @pytest.mark.parametrize("sample_rate", (0.0, 0.25, 0.5, 0.75, 1.0))
 def test_metrics_sampling_decision(