diff --git a/src/sentry/dynamic_sampling/per_org/calculations.py b/src/sentry/dynamic_sampling/per_org/calculations.py index 4584bf76d397a2..7bb8e63b87fd44 100644 --- a/src/sentry/dynamic_sampling/per_org/calculations.py +++ b/src/sentry/dynamic_sampling/per_org/calculations.py @@ -2,6 +2,7 @@ import logging from collections.abc import Iterable +from dataclasses import replace from typing import TYPE_CHECKING, cast import orjson @@ -23,6 +24,7 @@ from sentry.dynamic_sampling.per_org.gate import project_balancing_debug_project_ids from sentry.dynamic_sampling.per_org.queries import ProjectTransactionCounts, ProjectVolume from sentry.dynamic_sampling.rules.utils import get_redis_client_for_ds +from sentry.dynamic_sampling.sample_rate_override import get_sample_rate_overrides from sentry.dynamic_sampling.tasks.common import sample_rate_to_float from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( generate_boost_low_volume_projects_cache_key, @@ -63,6 +65,29 @@ def run_project_balancing( ) +def apply_project_sample_rate_overrides( + rebalanced_projects: list[RebalancedItem], +) -> list[RebalancedItem]: + """ + Hard-replace the balanced sample rate of any project that has a per-project override + configured via the ``dynamic-sampling.sample-rate-override-per-project`` option. + + Applied as an explicit step in the scheduler (rather than inside the balancing model) + so the override is surfaced in the pipeline. The result feeds the cached project + sample rates and the downstream transaction balancing. + """ + overrides = get_sample_rate_overrides() + if not overrides: + return rebalanced_projects + + return [ + replace(item, new_sample_rate=overrides[int(item.id)]) + if int(item.id) in overrides + else item + for item in rebalanced_projects + ] + + def get_cached_rebalanced_project_sample_rates(org_id: int) -> dict[int, float | None]: redis_client = get_redis_client_for_ds() cache_key = generate_boost_low_volume_projects_cache_key(org_id=org_id) diff --git a/src/sentry/dynamic_sampling/per_org/scheduler.py b/src/sentry/dynamic_sampling/per_org/scheduler.py index 3a6774a9a6e771..05dd06446e0791 100644 --- a/src/sentry/dynamic_sampling/per_org/scheduler.py +++ b/src/sentry/dynamic_sampling/per_org/scheduler.py @@ -9,6 +9,7 @@ from taskbroker_client.retry import Retry from sentry.dynamic_sampling.per_org.calculations import ( + apply_project_sample_rate_overrides, compare_rebalanced_projects_with_cache, compare_rebalanced_transactions_with_cache, get_cached_rebalanced_project_sample_rates, @@ -126,6 +127,7 @@ def run_calculations_per_org_task(org_id: OrganizationId) -> DynamicSamplingStat if config.should_balance_projects: rebalanced_projects = run_project_balancing(config, project_volumes) + rebalanced_projects = apply_project_sample_rate_overrides(rebalanced_projects) config.set_rebalanced_project_sample_rates(rebalanced_projects) cached_sample_rates = get_cached_rebalanced_project_sample_rates(config.organization.id) compare_rebalanced_projects_with_cache( diff --git a/src/sentry/dynamic_sampling/rules/base.py b/src/sentry/dynamic_sampling/rules/base.py index bf76a98d4dfe47..078c8610259b52 100644 --- a/src/sentry/dynamic_sampling/rules/base.py +++ b/src/sentry/dynamic_sampling/rules/base.py @@ -9,6 +9,7 @@ from sentry.dynamic_sampling.rules.biases.base import Bias from sentry.dynamic_sampling.rules.combine import get_relay_biases from sentry.dynamic_sampling.rules.utils import PolymorphicRule, RuleType, get_enabled_user_biases +from sentry.dynamic_sampling.sample_rate_override import get_sample_rate_override_for_project from sentry.dynamic_sampling.tasks.helpers.boost_low_volume_projects import ( get_boost_low_volume_projects_sample_rate, ) @@ -58,6 +59,12 @@ def get_guarded_project_sample_rate(organization: Organization, project: Project return float(project.get_option("sentry:target_sample_rate", TARGET_SAMPLE_RATE_DEFAULT)) if has_custom_dynamic_sampling(organization): + # A per-project override (configured via options) hard-replaces the rate the custom + # dynamic sampling path would otherwise compute, winning over project/org targets and + # the boosted/rebalanced rate. + override = get_sample_rate_override_for_project(project.id) + if override is not None: + return override sample_rate = organization.get_option("sentry:target_sample_rate") else: sample_rate = quotas.backend.get_blended_sample_rate( diff --git a/src/sentry/dynamic_sampling/sample_rate_override.py b/src/sentry/dynamic_sampling/sample_rate_override.py new file mode 100644 index 00000000000000..b1dfe6d0c6fc67 --- /dev/null +++ b/src/sentry/dynamic_sampling/sample_rate_override.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from sentry import options + + +def get_sample_rate_overrides() -> dict[int, float]: + """ + Return the validated per-project sample rate overrides for custom dynamic sampling, + as configured via the ``dynamic-sampling.sample-rate-override-per-project`` option. + + The option maps a stringified project id to a fixed sample rate that hard-replaces + whatever rate the custom dynamic sampling path would otherwise compute. Entries with + an invalid id or an out-of-range/invalid rate are skipped (rather than emitting an + invalid rule). Reads the option once so callers iterating over many projects don't + re-read it per project. + """ + overrides: dict[int, float] = {} + for raw_id, raw_rate in options.get( + "dynamic-sampling.sample-rate-override-per-project" + ).items(): + try: + project_id = int(raw_id) + rate = float(raw_rate) + except (TypeError, ValueError): + continue + if 0.0 <= rate <= 1.0: + overrides[project_id] = rate + return overrides + + +def get_sample_rate_override_for_project(project_id: int) -> float | None: + """Return the sample rate override for a single project, or ``None`` if none applies.""" + return get_sample_rate_overrides().get(project_id) diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py index 0202306afe1a78..2a36ee87619ec9 100644 --- a/src/sentry/options/defaults.py +++ b/src/sentry/options/defaults.py @@ -2148,6 +2148,16 @@ flags=FLAG_AUTOMATOR_MODIFIABLE, ) +# Per-project sample rate overrides for custom dynamic sampling. Maps a stringified +# project id to a fixed sample rate (0.0-1.0) that hard-replaces the rate the custom +# dynamic sampling path would otherwise compute for that project. Example: +# {"12345": 0.5}. An empty mapping disables the override. +register( + "dynamic-sampling.sample-rate-override-per-project", + default={}, + flags=FLAG_AUTOMATOR_MODIFIABLE, +) + # Controls the intensity of dynamic sampling transaction rebalancing. 0.0 = explict rebalancing # not performed, 1.0= full rebalancing (tries to bring everything to mean). Note that even at 0.0 # there will still be some rebalancing between the explicit and implicit transactions ( so setting rebalancing diff --git a/tests/sentry/dynamic_sampling/per_org/test_calculations.py b/tests/sentry/dynamic_sampling/per_org/test_calculations.py index 8db67c1dfb2116..c2043718371a29 100644 --- a/tests/sentry/dynamic_sampling/per_org/test_calculations.py +++ b/tests/sentry/dynamic_sampling/per_org/test_calculations.py @@ -8,6 +8,7 @@ from sentry.dynamic_sampling.models.common import RebalancedItem from sentry.dynamic_sampling.models.projects_rebalancing import ProjectsRebalancingInput from sentry.dynamic_sampling.per_org.calculations import ( + apply_project_sample_rate_overrides, compare_rebalanced_projects_with_cache, compare_rebalanced_transactions_with_cache, get_cached_rebalanced_project_sample_rates, @@ -25,6 +26,7 @@ generate_boost_low_volume_transactions_cache_key, ) from sentry.testutils.cases import TestCase +from sentry.testutils.helpers.options import override_options def _project_volume(project_id: int, total: int = 100, keep: int = 25) -> ProjectVolume: @@ -71,6 +73,32 @@ def test_run_project_balancing_returns_rebalanced_projects(self) -> None: ] assert result == rebalanced_projects + def test_apply_project_sample_rate_overrides(self) -> None: + overridden_id = 1001 + normal_id = 1002 + rebalanced_projects = [ + RebalancedItem(id=overridden_id, count=100, new_sample_rate=0.25), + RebalancedItem(id=normal_id, count=100, new_sample_rate=0.25), + ] + + with override_options( + {"dynamic-sampling.sample-rate-override-per-project": {str(overridden_id): 0.9}} + ): + result = apply_project_sample_rate_overrides(rebalanced_projects) + + result_by_id = {item.id: item.new_sample_rate for item in result} + # Overridden project gets the option value; the other keeps its balanced rate. + assert result_by_id[overridden_id] == 0.9 + assert result_by_id[normal_id] == 0.25 + + def test_apply_project_sample_rate_overrides_noop_without_option(self) -> None: + rebalanced_projects = [ + RebalancedItem(id=2001, count=100, new_sample_rate=0.25), + ] + # No overrides configured -> the balanced rates are returned untouched. + result = apply_project_sample_rate_overrides(rebalanced_projects) + assert result == rebalanced_projects + def test_compare_rebalanced_projects_with_cache_logs_per_project(self) -> None: org = self.create_organization() project_with_volume = self.create_project(organization=org) diff --git a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py index cd383b31a2682f..a66f03086adb35 100644 --- a/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py +++ b/tests/sentry/dynamic_sampling/tasks/test_boost_low_volume_projects.py @@ -27,6 +27,7 @@ from sentry.testutils.cases import BaseMetricsLayerTestCase, SnubaTestCase, TestCase from sentry.testutils.helpers.datetime import freeze_time from sentry.testutils.helpers.features import with_feature +from sentry.testutils.helpers.options import override_options MOCK_DATETIME = (timezone.now() - timedelta(days=1)).replace( hour=0, minute=0, second=0, microsecond=0 @@ -167,6 +168,43 @@ def test_simple_one_org_one_project_task_target_sample_rate(self) -> None: ) assert (sample_rate, got_value) == (0.5, True) + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_per_project_sample_rate_override(self) -> None: + # A per-project override configured via options hard-replaces the rate the + # custom dynamic sampling path would otherwise resolve for that project -- + # winning even over the recently-added 100% boost -- and leaves other projects + # untouched. + org1 = self.create_organization("am3-override-org") + org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION) + org1.update_option("sentry:target_sample_rate", 0.5) + overridden = self.create_project(organization=org1) + normal = self.create_project(organization=org1) + + # Baseline: freshly-created projects are boosted to 1.0 by the recently-added + # rule, so neither resolves to the org target yet. + assert get_guarded_project_sample_rate(org1, overridden) == 1.0 + + with override_options( + {"dynamic-sampling.sample-rate-override-per-project": {str(overridden.id): 0.9}} + ): + assert get_guarded_project_sample_rate(org1, overridden) == 0.9 + # Not in the override map -> unaffected by the override. + assert get_guarded_project_sample_rate(org1, normal) == 1.0 + + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) + def test_per_project_sample_rate_override_ignores_out_of_range(self) -> None: + org1 = self.create_organization("am3-override-org-bad") + org1.update_option("sentry:sampling_mode", DynamicSamplingMode.ORGANIZATION) + org1.update_option("sentry:target_sample_rate", 0.5) + project = self.create_project(organization=org1) + + baseline = get_guarded_project_sample_rate(org1, project) + with override_options( + {"dynamic-sampling.sample-rate-override-per-project": {str(project.id): 2.0}} + ): + # Out-of-range override is ignored; the resolved rate is unchanged. + assert get_guarded_project_sample_rate(org1, project) == baseline + @with_feature(["organizations:dynamic-sampling", "organizations:dynamic-sampling-custom"]) def test_project_mode_sampling_with_query(self) -> None: org1 = self.create_organization("test-org")