feat(tracemetrics): Add sample_rate to metrics api calls

sentry_sdk/metrics.py

@@ -1,15 +1,17 @@
 """
-NOTE: This file contains experimental code that may be changed or removed at any
-time without prior notice.
+NOTE: This file contains experimental code that may be changed or removed at
+any time without prior notice.
 """
 
 import time
-from typing import Any, Optional, TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 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
     from sentry_sdk._types import Metric, MetricType
 
 
@@ -19,6 +21,7 @@ def _capture_metric(
     value,  # type: float
     unit=None,  # type: Optional[str]
     attributes=None,  # type: Optional[dict[str, Any]]
+    sample_rate=None,  # type: Optional[float]
 ):
     # type: (...) -> None
     client = sentry_sdk.get_client()
@@ -37,6 +40,37 @@ def _capture_metric(
                 else safe_repr(v)
             )
 
+    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:
+            attrs["sentry.client_sample_rate"] = sample_rate
+
     metric = {
         "timestamp": time.time(),
         "trace_id": None,
@@ -56,26 +90,29 @@ def count(
     value,  # type: float
     unit=None,  # type: Optional[str]
     attributes=None,  # type: Optional[dict[str, Any]]
+    sample_rate=None,  # type: Optional[float]
 ):
     # type: (...) -> None
-    _capture_metric(name, "counter", value, unit, attributes)
+    _capture_metric(name, "counter", value, unit, attributes, sample_rate)
 
 
 def gauge(
     name,  # type: str
     value,  # type: float
     unit=None,  # type: Optional[str]
     attributes=None,  # type: Optional[dict[str, Any]]
+    sample_rate=None,  # type: Optional[float]
 ):
     # type: (...) -> None
-    _capture_metric(name, "gauge", value, unit, attributes)
+    _capture_metric(name, "gauge", value, unit, attributes, sample_rate)
 
 
 def distribution(
     name,  # type: str
     value,  # type: float
     unit=None,  # type: Optional[str]
     attributes=None,  # type: Optional[dict[str, Any]]
+    sample_rate=None,  # type: Optional[float]
 ):
     # type: (...) -> None
-    _capture_metric(name, "distribution", value, unit, attributes)
+    _capture_metric(name, "distribution", value, unit, attributes, sample_rate)

tests/test_metrics.py

@@ -1,6 +1,7 @@
 import json
 import sys
 from typing import List, Any, Mapping
+from unittest import mock
 import pytest
 
 import sentry_sdk
@@ -267,3 +268,113 @@ def record_lost_event(reason, data_category, quantity):
     assert len(lost_event_calls) == 5
     for lost_event_call in lost_event_calls:
         assert lost_event_call == ("queue_overflow", "trace_metric", 1)
+
+
+def test_metrics_sample_rate_basic(sentry_init, capture_envelopes):
+    sentry_init()
+    envelopes = capture_envelopes()
+
+    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
+
+    assert metrics[0]["name"] == "test.counter"
+    assert metrics[0]["attributes"]["sentry.client_sample_rate"] == 0.5
+
+    assert metrics[1]["name"] == "test.gauge"
+    assert metrics[1]["attributes"]["sentry.client_sample_rate"] == 0.8
+
+    assert metrics[2]["name"] == "test.distribution"
+    assert "sentry.client_sample_rate" not in metrics[2]["attributes"]
+
+
+def test_metrics_sample_rate_normalization(sentry_init, capture_envelopes, monkeypatch):
+    sentry_init()
+    envelopes = capture_envelopes()
+    client = sentry_sdk.get_client()
+
+    lost_event_calls = []
+
+    def record_lost_event(reason, data_category, quantity):
+        lost_event_calls.append((reason, data_category, quantity))
+
+    monkeypatch.setattr(client.transport, "record_lost_event", record_lost_event)
+
+    sentry_sdk.metrics.count("test.counter1", 1, sample_rate=0.0)  # <= 0
+    sentry_sdk.metrics.count("test.counter2", 1, sample_rate=-0.5)  # < 0
+    sentry_sdk.metrics.count("test.counter3", 1, sample_rate=0.5)  # > 0 but < 1.0
+    sentry_sdk.metrics.count("test.counter4", 1, sample_rate=1.0)  # = 1.0
+    sentry_sdk.metrics.count("test.counter4", 1, sample_rate=1.5)  # > 1.0
+
+    client.flush()
+    metrics = envelopes_to_metrics(envelopes)
+
+    assert len(metrics) == 2
+
+    assert metrics[0]["attributes"]["sentry.client_sample_rate"] == 0.5
+    assert (
+        "sentry.client_sample_rate" not in metrics[1]["attributes"]
+    )  # 1.0 does not need a sample rate, it's implied to be 1.0
+
+    assert len(lost_event_calls) == 3
+    assert lost_event_calls[0] == ("invalid_sample_rate", "trace_metric", 1)
+    assert lost_event_calls[1] == ("invalid_sample_rate", "trace_metric", 1)
+    assert lost_event_calls[2] == ("invalid_sample_rate", "trace_metric", 1)
+
+
+def test_metrics_no_sample_rate(sentry_init, capture_envelopes):
+    sentry_init()
+    envelopes = capture_envelopes()
+
+    sentry_sdk.metrics.count("test.counter", 1)
+
+    get_client().flush()
+    metrics = envelopes_to_metrics(envelopes)
+
+    assert len(metrics) == 1
+
+    assert "sentry.client_sample_rate" not in metrics[0]["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(
+    sentry_init, capture_envelopes, sample_rate, sample_rand, monkeypatch
+):
+    sentry_init(traces_sample_rate=1.0)
+    envelopes = capture_envelopes()
+    client = sentry_sdk.get_client()
+
+    lost_event_calls = []
+
+    def record_lost_event(reason, data_category, quantity):
+        lost_event_calls.append((reason, data_category, quantity))
+
+    monkeypatch.setattr(client.transport, "record_lost_event", record_lost_event)
+
+    with mock.patch(
+        "sentry_sdk.tracing_utils.Random.randrange",
+        return_value=int(sample_rand * 1000000),
+    ):
+        with sentry_sdk.start_transaction() as transaction:
+            sentry_sdk.metrics.count("test.counter", 1, sample_rate=sample_rate)
+
+    get_client().flush()
+    metrics = envelopes_to_metrics(envelopes)
+
+    should_be_sampled = sample_rand < sample_rate and sample_rate > 0.0
+    assert len(metrics) == int(should_be_sampled)
+
+    if sample_rate <= 0.0:
+        assert len(lost_event_calls) == 1
+        assert lost_event_calls[0] == ("invalid_sample_rate", "trace_metric", 1)
+    elif not should_be_sampled:
+        assert len(lost_event_calls) == 1
+        assert lost_event_calls[0] == ("sample_rate", "trace_metric", 1)
+    else:
+        assert len(lost_event_calls) == 0