feat(sentry): Keep only a fraction of the traces that callers sampled

snuba/environment.py

@@ -1,5 +1,7 @@
 import logging
 import os
+import re
+from typing import Any
 
 import sentry_sdk
 import structlog
@@ -160,6 +162,48 @@ def before_send(event: Event, hint: Hint) -> Event | None:
     return event
 
 
+_HEALTH_CHECK_PATH = re.compile(r"^/health(_envoy)?/?$")
+
+
+def _inherited_sample_rate(sampling_context: dict[str, Any]) -> float:
+    """
+    Keep a fraction of the traces that the caller sampled, chosen by trace
+    id so that every snuba request of one trace gets the same decision.
+
+    The bucketing matches the one sentry applies in its own transport, so
+    both services keep the same traces when they run with the same rate.
+    """
+    rate = settings.SENTRY_INHERITED_TRACE_SAMPLE_RATE
+    trace_id = sampling_context.get("transaction_context", {}).get("trace_id")
+    if rate >= 1.0 or not trace_id:
+        return rate
+    return 1.0 if int(trace_id[:8], 16) % 100 < int(rate * 100) else 0.0
+
+
+def traces_sampler(sampling_context: dict[str, Any]) -> float:
+    """
+    Decide the sample rate for a root span. This replaces the server-side
+    dynamic sampling rules for the snuba project.
+
+    A request that the caller did not sample is never sampled. A request
+    that the caller sampled is kept at `SENTRY_INHERITED_TRACE_SAMPLE_RATE`.
+    Only traces that start in snuba get a fresh decision.
+    """
+    parent_sampled = sampling_context.get("parent_sampled")
+    if parent_sampled is not None:
+        return _inherited_sample_rate(sampling_context) if parent_sampled else 0.0
+
+    environment = sentry_sdk.get_client().options.get("environment") or ""
+    if any(marker in environment for marker in settings.SENTRY_ALWAYS_SAMPLED_ENVIRONMENTS):
+        return 1.0
+
+    path = sampling_context.get("wsgi_environ", {}).get("PATH_INFO", "")
+    if _HEALTH_CHECK_PATH.match(path):
+        return settings.SENTRY_HEALTH_CHECK_TRACE_SAMPLE_RATE
+
+    return settings.SENTRY_TRACE_SAMPLE_RATE
+
+
 def setup_sentry() -> None:
     sentry_sdk.init(
         dsn=settings.SENTRY_DSN,
@@ -179,7 +223,7 @@ def setup_sentry() -> None:
         # the value for release is also computed in rust-snuba, please keep the
         # logic in sync
         release=os.getenv("SNUBA_RELEASE"),
-        traces_sample_rate=settings.SENTRY_TRACE_SAMPLE_RATE,
+        traces_sampler=traces_sampler,
         profiles_sample_rate=settings.SNUBA_PROFILES_SAMPLE_RATE,
         # Stream spans as they finish. Disables the legacy tracing API
         # (start_span/start_transaction/update_current_span/scope.span).

snuba/settings/__init__.py

@@ -218,7 +218,18 @@ class RedisClusters(TypedDict):
 
 # Sentry Options
 SENTRY_DSN: str | None = None
-SENTRY_TRACE_SAMPLE_RATE = 0
+# Sample rate for traces that start in snuba (no incoming sampling decision).
+SENTRY_TRACE_SAMPLE_RATE = 0.0
+# Fraction of the traces sampled by the caller that snuba keeps. Nearly all
+# snuba traffic carries a caller decision, so this is the main lever on the
+# volume of the snuba project.
+SENTRY_INHERITED_TRACE_SAMPLE_RATE = float(
+    os.environ.get("SENTRY_INHERITED_TRACE_SAMPLE_RATE", 0.1)
+)
+# Sample rate for traces started by health check requests.
+SENTRY_HEALTH_CHECK_TRACE_SAMPLE_RATE = 0.0
+# Environments whose traces are always sampled, matched as substrings.
+SENTRY_ALWAYS_SAMPLED_ENVIRONMENTS = ("debug", "dev", "local", "qa", "test")
 
 # Snuba Admin Options
 SLACK_API_TOKEN = os.environ.get("SLACK_API_TOKEN")

tests/test_environment.py

@@ -1,3 +1,8 @@
+from collections.abc import Callable
+from typing import Any
+
+import pytest
+import sentry_sdk
 from arroyo.processing.strategies.run_task_with_multiprocessing import (
     ChildProcessTerminated,
 )
@@ -6,7 +11,8 @@
 from redis.exceptions import TimeoutError as RedisTimeoutError
 from sentry_sdk.types import Event, Hint
 
-from snuba.environment import before_send
+from snuba import settings
+from snuba.environment import before_send, traces_sampler
 from snuba.query.allocation_policies import AllocationPolicyViolations
 from snuba.web.rpc.common.exceptions import RPCAllocationPolicyException
 
@@ -144,3 +150,101 @@ def test_before_send_terminates_on_cyclic_cause_chain() -> None:
     except ValueError:
         err.__context__ = err
         assert before_send(event, _hint_for(err)) is event
+
+
+@pytest.fixture
+def sentry_environment(monkeypatch: pytest.MonkeyPatch) -> Callable[[str | None], None]:
+    def _set(environment: str | None) -> None:
+        monkeypatch.setattr(sentry_sdk.get_client(), "options", {"environment": environment})
+
+    _set(None)
+    return _set
+
+
+@pytest.fixture
+def sample_rates(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(settings, "SENTRY_TRACE_SAMPLE_RATE", 0.5)
+    monkeypatch.setattr(settings, "SENTRY_HEALTH_CHECK_TRACE_SAMPLE_RATE", 0.01)
+    monkeypatch.setattr(settings, "SENTRY_INHERITED_TRACE_SAMPLE_RATE", 1.0)
+
+
+def _request(
+    path: str, parent_sampled: bool | None = None, trace_id: str | None = None
+) -> dict[str, Any]:
+    transaction_context: dict[str, Any] = {"name": "generic WSGI request"}
+    if trace_id is not None:
+        transaction_context["trace_id"] = trace_id
+    return {
+        "transaction_context": transaction_context,
+        "parent_sampled": parent_sampled,
+        "wsgi_environ": {"PATH_INFO": path},
+    }
+
+
+@pytest.mark.usefixtures("sentry_environment", "sample_rates")
+def test_traces_sampler_inherits_parent_decision() -> None:
+    assert traces_sampler(_request("/health", parent_sampled=True)) == 1.0
+    assert traces_sampler(_request("/query", parent_sampled=False)) == 0.0
+
+
+@pytest.mark.usefixtures("sentry_environment", "sample_rates")
+def test_traces_sampler_keeps_a_fraction_of_sampled_traces_by_trace_id(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(settings, "SENTRY_INHERITED_TRACE_SAMPLE_RATE", 0.1)
+
+    # int("00000005", 16) % 100 == 5 < 10, int("00000063", 16) % 100 == 99
+    kept = "00000005" + "0" * 24
+    dropped = "00000063" + "0" * 24
+    assert traces_sampler(_request("/query", parent_sampled=True, trace_id=kept)) == 1.0
+    assert traces_sampler(_request("/query", parent_sampled=True, trace_id=dropped)) == 0.0
+    assert traces_sampler(_request("/query", parent_sampled=False, trace_id=kept)) == 0.0
+
+
+@pytest.mark.usefixtures("sentry_environment", "sample_rates")
+def test_traces_sampler_falls_back_to_the_inherited_rate_without_trace_id(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(settings, "SENTRY_INHERITED_TRACE_SAMPLE_RATE", 0.1)
+    assert traces_sampler(_request("/query", parent_sampled=True)) == 0.1
+
+
+@pytest.mark.usefixtures("sentry_environment", "sample_rates")
+def test_traces_sampler_uses_base_rate_for_root_requests() -> None:
+    assert traces_sampler(_request("/query")) == 0.5
+    assert traces_sampler(_request("/rpc/EndpointTraceItemTable/v1")) == 0.5
+
+
+@pytest.mark.usefixtures("sentry_environment", "sample_rates")
+def test_traces_sampler_uses_health_check_rate() -> None:
+    assert traces_sampler(_request("/health")) == 0.01
+    assert traces_sampler(_request("/health/")) == 0.01
+    assert traces_sampler(_request("/health_envoy")) == 0.01
+    assert traces_sampler(_request("/healthy")) == 0.5
+
+
+@pytest.mark.usefixtures("sample_rates")
+def test_traces_sampler_uses_base_rate_without_wsgi_environ(
+    sentry_environment: Callable[[str | None], None],
+) -> None:
+    sentry_environment("us")
+    assert traces_sampler({"transaction_context": {"name": "[cli init] api"}}) == 0.5
+
+
+@pytest.mark.usefixtures("sample_rates")
+@pytest.mark.parametrize("environment", ["dev", "local-simon", "test", "qa-eu", "debug"])
+def test_traces_sampler_samples_everything_in_development_environments(
+    sentry_environment: Callable[[str | None], None], environment: str
+) -> None:
+    sentry_environment(environment)
+    assert traces_sampler(_request("/health")) == 1.0
+    assert traces_sampler(_request("/query")) == 1.0
+
+
+@pytest.mark.usefixtures("sample_rates")
+@pytest.mark.parametrize("environment", ["us", "de", "s4s2", "production"])
+def test_traces_sampler_keeps_rates_in_production_environments(
+    sentry_environment: Callable[[str | None], None], environment: str
+) -> None:
+    sentry_environment(environment)
+    assert traces_sampler(_request("/query")) == 0.5