feat(tracemetrics): Add sample_rate to metrics api calls

This allows for a sample_rate (0, 1.0] to be sent on a per metric basis.

Author: k-fish

sentry_sdk/metrics.py

@@ -1,6 +1,6 @@
 """
-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
@@ -19,6 +19,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 +38,19 @@ 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
+
+        if sample_rate != 1.0:
+            attrs["sentry.client_sample_rate"] = sample_rate
+
     metric = {
         "timestamp": time.time(),
         "trace_id": None,
@@ -56,26 +70,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

@@ -267,3 +267,75 @@ 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
+
+    # Should not have sample_rate attribute when not provided
+    assert "sentry.client_sample_rate" not in metrics[0]["attributes"]