Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions src/sentry/dynamic_sampling/per_org/calculations.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import logging
from collections.abc import Iterable
from dataclasses import replace
from typing import TYPE_CHECKING, cast

import orjson
Expand All @@ -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_override_for_project
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,
Expand Down Expand Up @@ -63,6 +65,25 @@ 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.
"""
return [
replace(item, new_sample_rate=override)
if (override := get_sample_rate_override_for_project(int(item.id))) is not None
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)
Expand Down
2 changes: 2 additions & 0 deletions src/sentry/dynamic_sampling/per_org/scheduler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down
7 changes: 7 additions & 0 deletions src/sentry/dynamic_sampling/rules/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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
Comment thread
shellmayr marked this conversation as resolved.
# 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
Comment thread
shellmayr marked this conversation as resolved.
sample_rate = organization.get_option("sentry:target_sample_rate")
else:
sample_rate = quotas.backend.get_blended_sample_rate(
Expand Down
28 changes: 28 additions & 0 deletions src/sentry/dynamic_sampling/sample_rate_override.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
from __future__ import annotations

from sentry import options


def get_sample_rate_override_for_project(project_id: int) -> float | None:
"""
Return a per-project sample rate override for custom dynamic sampling, if one is
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. Returns
``None`` when no override applies (no entry, or an out-of-range/invalid value, which
we ignore rather than emit an invalid rule).
"""
overrides = options.get("dynamic-sampling.sample-rate-override-per-project")
raw = overrides.get(str(project_id))
if raw is None:
return None

try:
rate = float(raw)
except (TypeError, ValueError):
return None

if 0.0 <= rate <= 1.0:
return rate
return None
10 changes: 10 additions & 0 deletions src/sentry/options/defaults.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions tests/sentry/dynamic_sampling/per_org/test_calculations.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down
Loading