diff --git a/snuba/environment.py b/snuba/environment.py index f689c3dcc5..71618ccca6 100644 --- a/snuba/environment.py +++ b/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). diff --git a/snuba/settings/__init__.py b/snuba/settings/__init__.py index f9225d3ef4..12189148b9 100644 --- a/snuba/settings/__init__.py +++ b/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") diff --git a/tests/test_environment.py b/tests/test_environment.py index e8359fc9d6..6b74396aa6 100644 --- a/tests/test_environment.py +++ b/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