diff --git a/sentry-options/schemas/snuba/schema.json b/sentry-options/schemas/snuba/schema.json index c6bbb6cbf..57cd46ed2 100644 --- a/sentry-options/schemas/snuba/schema.json +++ b/sentry-options/schemas/snuba/schema.json @@ -387,19 +387,8 @@ "type": "integer", "optional": true }, - "cross_org_referrer_limits": { - "type": "object", - "additionalProperties": { - "type": "object", - "properties": { - "max_threads": { - "type": "integer" - }, - "concurrent_limit": { - "type": "integer" - } - } - }, + "key_by": { + "type": "string", "optional": true } }, diff --git a/snuba/datasets/configuration/outcomes/entities/outcomes_cross_org_hourly.yaml b/snuba/datasets/configuration/outcomes/entities/outcomes_cross_org_hourly.yaml index 49f613495..d3800446b 100644 --- a/snuba/datasets/configuration/outcomes/entities/outcomes_cross_org_hourly.yaml +++ b/snuba/datasets/configuration/outcomes/entities/outcomes_cross_org_hourly.yaml @@ -31,7 +31,6 @@ query_processors: # Dedicated org-level outcomes entity for cross-org active-org discovery. Unlike the # shared `outcomes` entity it intentionally has an EMPTY validators list (no # EntityRequiredColumnValidator for org_id), so callers can group by org_id without a -# per-org filter. Access is governed at the storage layer by -# CrossOrgQueryAllocationPolicy. +# per-org filter. validators: [] validate_data_model: error diff --git a/snuba/datasets/configuration/outcomes/storages/cross_org_hourly.yaml b/snuba/datasets/configuration/outcomes/storages/cross_org_hourly.yaml index f29fcae04..44ed2bc2b 100644 --- a/snuba/datasets/configuration/outcomes/storages/cross_org_hourly.yaml +++ b/snuba/datasets/configuration/outcomes/storages/cross_org_hourly.yaml @@ -35,5 +35,4 @@ query_processors: # shared outcomes_hourly storage over the same ClickHouse tables but intentionally # OMITS the OrgIdEnforcer mandatory_condition_checker so a query can scan across # organizations. The shared `outcomes` storage keeps OrgIdEnforcer; this storage is -# reachable only via the dedicated `outcomes_cross_org_hourly` entity and is gated by -# CrossOrgQueryAllocationPolicy (registered referrers only). +# reachable only via the dedicated `outcomes_cross_org_hourly` entity. diff --git a/snuba/query/allocation_policies/concurrent_rate_limit.py b/snuba/query/allocation_policies/concurrent_rate_limit.py index bf4cce405..329bc1125 100644 --- a/snuba/query/allocation_policies/concurrent_rate_limit.py +++ b/snuba/query/allocation_policies/concurrent_rate_limit.py @@ -2,8 +2,6 @@ import logging import typing -from collections.abc import Callable -from typing import cast from snuba import state from snuba.configs.configuration import Configuration @@ -16,7 +14,6 @@ AllocationPolicyViolations, InvalidTenantsForAllocationPolicy, QueryResultOrError, - QueryType, QuotaAllowance, ) from snuba.state.rate_limit import ( @@ -43,22 +40,18 @@ } QUOTA_UNIT = "concurrent_queries" +KEY_BY_TENANT = "tenant" +KEY_BY_ORGANIZATION = "organization" +KEY_BY_ORGANIZATION_REFERRER = "organization_referrer" +_KEY_BY_VALUES = {KEY_BY_TENANT, KEY_BY_ORGANIZATION, KEY_BY_ORGANIZATION_REFERRER} +# Never configured in any region. Changing it would move every count to new Redis keys. +_RATE_LIMIT_SHARD_FACTOR = 1 SUGGESTION = "A customer is sending too many queries to snuba. The customer may be abusing an API or the queries may be innefficient" class BaseConcurrentRateLimitAllocationPolicy(AllocationPolicy): def _additional_config_definitions(self) -> list[Configuration]: return [ - Configuration( - name="rate_limit_shard_factor", - description="""number of shards that each redis set is supposed to have. - increasing this value multiplies the number of redis keys by that - factor, and (on average) reduces the size of each redis set. You probably don't need to change this - unless you're scaling out redis for some reason - """, - value_type=int, - default=1, - ), Configuration( name="max_query_duration_s", description="""maximum duration of a query in seconds. Queries that exceed this duration are considered finished by the rate limiter. This reduces memory usage. If you turn this down lower than the actual timeout period, the system can start undercounting concurrent queries""", @@ -71,6 +64,50 @@ def _additional_config_definitions(self) -> list[Configuration]: def rate_limit_name(self) -> str: raise NotImplementedError + # Template method: every concurrency policy counts in-flight queries the same way. + # Subclasses only choose the bucket and limit (_rate_limit_params), which queries + # they ignore (_pass_through_allowance, _skips_balance_update) and how a decision + # becomes an allowance (_allowance). + def _get_quota_allowance( + self, tenant_ids: dict[str, str | int], query_id: str + ) -> QuotaAllowance: + pass_through = self._pass_through_allowance(tenant_ids) + if pass_through is not None: + return pass_through + rate_limit_params = self._rate_limit_params(tenant_ids) + rate_limit_stats, can_run, why = self._is_within_rate_limit(query_id, rate_limit_params) + return self._allowance(tenant_ids, rate_limit_params, rate_limit_stats, can_run, why) + + def _update_quota_balance( + self, + tenant_ids: dict[str, str | int], + query_id: str, + result_or_error: QueryResultOrError, + ) -> None: + if self._skips_balance_update(tenant_ids): + return + self._end_query(query_id, self._rate_limit_params(tenant_ids), result_or_error) + + def _pass_through_allowance(self, tenant_ids: dict[str, str | int]) -> QuotaAllowance | None: + """An allowance for queries this policy doesn't count, or None to count it.""" + return None + + def _skips_balance_update(self, tenant_ids: dict[str, str | int]) -> bool: + return False + + def _rate_limit_params(self, tenant_ids: dict[str, str | int]) -> RateLimitParameters: + raise NotImplementedError + + def _allowance( + self, + tenant_ids: dict[str, str | int], + rate_limit_params: RateLimitParameters, + rate_limit_stats: RateLimitStats, + can_run: bool, + why: str, + ) -> QuotaAllowance: + raise NotImplementedError + def _is_within_rate_limit( self, query_id: str, rate_limit_params: RateLimitParameters ) -> tuple[RateLimitStats, bool, str]: @@ -79,18 +116,14 @@ def _is_within_rate_limit( # configuration parameter. It's used for the per-second caclulation but that calculation # is fundamentally flawed rate_history_s = 1 - rate_limit_shard_factor = self.get_config_value("rate_limit_shard_factor") assert isinstance(rate_history_s, (int, float)) - assert isinstance(rate_limit_shard_factor, int) assert rate_limit_params.concurrent_limit is not None, "concurrent_limit must be set" - assert rate_limit_shard_factor > 0 - rate_limit_stats = rate_limit_start_request( rate_limit_params, query_id, rate_history_s, - rate_limit_shard_factor, + _RATE_LIMIT_SHARD_FACTOR, rate_limit_prefix, self.get_config_value("max_query_duration_s"), ) @@ -113,8 +146,6 @@ def _end_query( # removes the current query from the rate limit bookkeeping so it is no longer counted # in rate limits rate_limit_prefix = f"{self.component_name()}.rate_limit" - rate_limit_shard_factor = self.get_config_value("rate_limit_shard_factor") - was_rate_limited = result_or_error.error is not None and isinstance( result_or_error.error.__cause__, AllocationPolicyViolations, @@ -122,7 +153,7 @@ def _end_query( rate_limit_finish_request( rate_limit_params, query_id, - rate_limit_shard_factor, + _RATE_LIMIT_SHARD_FACTOR, was_rate_limited, rate_limit_prefix, self.get_config_value("max_query_duration_s"), @@ -139,55 +170,13 @@ def _additional_config_definitions(self) -> list[Configuration]: default=DEFAULT_CONCURRENT_QUERIES_LIMIT, ), Configuration( - name="referrer_project_override", - description="override concurrent limit for a specific project, referrer combo", - value_type=int, - default=-1, - param_types={"referrer": str, "project_id": int}, - ), - Configuration( - name="referrer_organization_override", - description="override concurrent limit for a specific organization_id, referrer combo", - value_type=int, - default=-1, - param_types={"referrer": str, "organization_id": int}, - ), - Configuration( - name="project_override", - description="override concurrent limit for a specific project_id", - value_type=int, - default=-1, - param_types={"project_id": int}, - ), - Configuration( - name="organization_override", - description="override concurrent limit for a specific organization_id", - value_type=int, - default=-1, - param_types={"organization_id": int}, + name="key_by", + description=f"Which tenant the concurrency count is kept for: {KEY_BY_TENANT} (project_id if present, else organization_id), {KEY_BY_ORGANIZATION} or {KEY_BY_ORGANIZATION_REFERRER}. Per-tenant limits go in match blocks. Queries without an organization_id use {KEY_BY_TENANT}", + value_type=str, + default=KEY_BY_TENANT, ), ] - def _get_overrides(self, tenant_ids: dict[str, str | int]) -> dict[str, int]: - overrides = {} - available_tenant_ids = set(tenant_ids.keys()) - # get all overrides that can be retrieved with the tenant_ids - # e.g. if organization_id and referrer are passed in, retrieve - # ('organization_override, 'referrer_organization_override') - for config_definition in self._additional_config_definitions(): - if config_definition.name.endswith("_override"): - param_types = config_definition.param_types - if set(param_types.keys()).issubset(available_tenant_ids): - params = {param: tenant_ids[param] for param in param_types} - config_value = self.get_config_value(config_definition.name, params) - if config_value != config_definition.default: - key = "|".join( - [f"{param}__{tenant_id}" for param, tenant_id in sorted(params.items())] - ) - - overrides[key] = config_value - return overrides - def _get_tenant_key_and_value(self, tenant_ids: dict[str, str | int]) -> tuple[str, str | int]: if "project_id" in tenant_ids: return "project_id", tenant_ids["project_id"] @@ -203,29 +192,33 @@ def _get_tenant_key_and_value(self, tenant_ids: dict[str, str | int]) -> tuple[s def rate_limit_name(self) -> str: return "concurrent_rate_limit_policy" - def _get_rate_limit_params( - self, tenant_ids: dict[str, str | int] - ) -> tuple[RateLimitParameters, dict[str, int]]: - tenant_key, tenant_value = self._get_tenant_key_and_value(tenant_ids) - overrides = self._get_overrides(tenant_ids) - concurrent_limit = self.get_config_value("concurrent_limit") - if overrides: - concurrent_limit = min(overrides.values()) - tenant_value = min(overrides, key=cast(Callable[[str], int], overrides.get)) - - return ( - RateLimitParameters( - self.rate_limit_name, - bucket=str(tenant_value), - per_second_limit=None, - concurrent_limit=concurrent_limit, - ), - overrides, + def _get_rate_limit_params(self, tenant_ids: dict[str, str | int]) -> RateLimitParameters: + return RateLimitParameters( + self.rate_limit_name, + bucket=self._get_bucket(tenant_ids), + per_second_limit=None, + concurrent_limit=self.get_config_value("concurrent_limit"), ) - def _get_quota_allowance( - self, tenant_ids: dict[str, str | int], query_id: str - ) -> QuotaAllowance: + def _get_bucket(self, tenant_ids: dict[str, str | int]) -> str: + key_by = self.get_config_value("key_by") + if key_by not in _KEY_BY_VALUES: + logger.warning("Unknown key_by %s on %s", key_by, self.component_name()) + key_by = KEY_BY_TENANT + if key_by != KEY_BY_TENANT and "organization_id" in tenant_ids: + # Same bucket names the removed organization_override and + # referrer_organization_override params used. + bucket = f"organization_id__{tenant_ids['organization_id']}" + if key_by == KEY_BY_ORGANIZATION_REFERRER: + bucket += f"|referrer__{tenant_ids.get('referrer', 'no_referrer')}" + return bucket + _, tenant_value = self._get_tenant_key_and_value(tenant_ids) + return str(tenant_value) + + def _rate_limit_params(self, tenant_ids: dict[str, str | int]) -> RateLimitParameters: + return self._get_rate_limit_params(tenant_ids) + + def _pass_through_allowance(self, tenant_ids: dict[str, str | int]) -> QuotaAllowance | None: if tenant_ids.get("referrer", "no_referrer") in _PASS_THROUGH_REFERRERS: return QuotaAllowance( can_run=True, @@ -250,42 +243,28 @@ def _get_quota_allowance( quota_unit=QUOTA_UNIT, suggestion=CROSS_ORG_SUGGESTION, ) + return None - rate_limit_params, overrides = self._get_rate_limit_params(tenant_ids) - rate_limit_stats, within_rate_limit, why = self._is_within_rate_limit( - query_id, - rate_limit_params, - ) + def _skips_balance_update(self, tenant_ids: dict[str, str | int]) -> bool: + # Pass-through referrers still end their query; only cross-org queries skip it. + return self.is_cross_org_query(tenant_ids) + def _allowance( + self, + tenant_ids: dict[str, str | int], + rate_limit_params: RateLimitParameters, + rate_limit_stats: RateLimitStats, + can_run: bool, + why: str, + ) -> QuotaAllowance: return QuotaAllowance( - can_run=within_rate_limit, - max_threads=self.max_threads if within_rate_limit else 0, - explanation={"reason": why, "overrides": overrides}, + can_run=can_run, + max_threads=self.max_threads if can_run else 0, + explanation={"reason": why}, is_throttled=False, throttle_threshold=typing.cast(int, rate_limit_params.concurrent_limit), rejection_threshold=typing.cast(int, rate_limit_params.concurrent_limit), quota_used=rate_limit_stats.concurrent, quota_unit=QUOTA_UNIT, - suggestion=NO_SUGGESTION if within_rate_limit else SUGGESTION, + suggestion=NO_SUGGESTION if can_run else SUGGESTION, ) - - def _update_quota_balance( - self, - tenant_ids: dict[str, str | int], - query_id: str, - result_or_error: QueryResultOrError, - ) -> None: - if self.is_cross_org_query(tenant_ids): - return - rate_limit_params, _ = self._get_rate_limit_params(tenant_ids) - self._end_query(query_id, rate_limit_params, result_or_error) - - -class DeleteConcurrentRateLimitAllocationPolicy(ConcurrentRateLimitAllocationPolicy): - @property - def rate_limit_name(self) -> str: - return "delete_concurrent_rate_limit_policy" - - @property - def query_type(self) -> QueryType: - return QueryType.DELETE diff --git a/snuba/query/allocation_policies/cross_org.py b/snuba/query/allocation_policies/cross_org.py deleted file mode 100644 index 2f582960c..000000000 --- a/snuba/query/allocation_policies/cross_org.py +++ /dev/null @@ -1,176 +0,0 @@ -from __future__ import annotations - -import logging -import typing -from typing import Any, cast - -from snuba.configs.configuration import Configuration, ResourceIdentifier -from snuba.query.allocation_policies import ( - MAX_THRESHOLD, - NO_SUGGESTION, - NO_UNITS, - QueryResultOrError, - QuotaAllowance, -) -from snuba.query.allocation_policies.concurrent_rate_limit import ( - BaseConcurrentRateLimitAllocationPolicy, -) -from snuba.redis import RedisClientKey, get_redis_client -from snuba.state.rate_limit import RateLimitParameters -from snuba.utils.serializable_exception import JsonSerializable - -DEFAULT_CONCURRENT_QUERIES_LIMIT = 22 -DEFAULT_PER_SECOND_QUERIES_LIMIT = 50 - -rds = get_redis_client(RedisClientKey.RATE_LIMITER) - -logger = logging.getLogger("snuba.query.allocation_policy_cross_org") - -_RATE_LIMIT_NAME = "concurrent_limit_policy" -_UNREGISTERED_REFERRER_MAX_THREADS = 1 -_UNREGISTERED_REFERRER_CONCURRENT_QUERIES = 1 - -QUOTA_UNIT = "concurrent_queries" -SUGGESTION = "scan less concurrent queries" - - -class CrossOrgQueryAllocationPolicy(BaseConcurrentRateLimitAllocationPolicy): - """A case-by-case allocation policy for cross-org queries. All referrers affected by this policy have to be registered - in this class's configuration through the `cross_org_referrer_limits` parameter. Example: - - ```yaml - - name: CrossOrgQueryAllocationPolicy - is_enforced: 0 - cross_org_referrer_limits: - dynamic_sampling.counters.get_org_transaction_volumes: - max_threads: 4 - concurrent_limit: 10 - ``` - - Each referrer gets a concurrent limit (applied per referrer) and a max_threads limit (applied to every query made by that referrer). - Both limits are static but changeable at runtime. - - unregistered referrers are assigned a default concurrent limit and max_threads limit of 1. - """ - - @property - def rate_limit_name(self) -> str: - return "cross_org_query_policy" - - def _additional_config_definitions(self) -> list[Configuration]: - return super()._additional_config_definitions() + [ - Configuration( - name="referrer_concurrent_override", - description="""override the concurrent limit for a referrer""", - value_type=int, - param_types={"referrer": str}, - default=-1, - ), - ] - - def _validate_cross_org_referrer_limits( - self, cross_org_referrer_limits: dict[str, dict[str, int]] - ) -> None: - for referrer, limits in cross_org_referrer_limits.items(): - concurrent_limit = limits.get("concurrent_limit", None) - max_threads = limits.get("max_threads", None) - if not isinstance(max_threads, int): - raise ValueError(f"max_threads is required for {referrer} and must be an int") - if not isinstance(concurrent_limit, int): - raise ValueError(f"concurrent_limit is required for {referrer} and must be an int") - - def __init__( - self, - storage_key: ResourceIdentifier, - **kwargs: Any, - ) -> None: - super().__init__( - storage_key, - **kwargs, - ) - self._registered_cross_org_referrers = cast( - "dict[str, dict[str, int]]", kwargs.get("cross_org_referrer_limits", {}) - ) - self._validate_cross_org_referrer_limits(self._registered_cross_org_referrers) - - def _get_max_threads(self, referrer: str) -> int: - if not self._referrer_is_registered(referrer): - return _UNREGISTERED_REFERRER_MAX_THREADS - return int(self._registered_cross_org_referrers[referrer]["max_threads"]) - - def _get_concurrent_limit(self, referrer: str) -> int: - if not self._referrer_is_registered(referrer): - return _UNREGISTERED_REFERRER_CONCURRENT_QUERIES - concurrent_override = int( - self.get_config_value("referrer_concurrent_override", {"referrer": referrer}) - ) - return ( - concurrent_override - if concurrent_override != -1 - else int(self._registered_cross_org_referrers[referrer]["concurrent_limit"]) - ) - - def _referrer_is_registered(self, referrer: str) -> bool: - return referrer in self._registered_cross_org_referrers - - def _get_quota_allowance( - self, tenant_ids: dict[str, str | int], query_id: str - ) -> QuotaAllowance: - referrer = str(tenant_ids.get("referrer", "no_referrer")) - if not self._referrer_is_registered(referrer) and not self.is_cross_org_query(tenant_ids): - # This is not a cross org query and the referrer is not registered. This is outside the responsibility of this policy - return QuotaAllowance( - can_run=True, - max_threads=self.max_threads, - explanation={"reason": "pass_through"}, - is_throttled=False, - throttle_threshold=MAX_THRESHOLD, - rejection_threshold=MAX_THRESHOLD, - quota_used=0, - quota_unit=NO_UNITS, - suggestion=NO_SUGGESTION, - ) - - concurrent_limit = self._get_concurrent_limit(referrer) - rate_limit_params = RateLimitParameters( - self.rate_limit_name, referrer, None, concurrent_limit - ) - rate_limit_stats, can_run, explanation = self._is_within_rate_limit( - query_id, - rate_limit_params, - ) - decision_explanation: dict[str, JsonSerializable] = {"reason": explanation} - if not self._referrer_is_registered(referrer): - decision_explanation["cross_org_query"] = ( - f"This referrer is not registered for the current storage {self._resource_identifier.value}, if you want to increase its limits, register it in the yaml of the CrossOrgQueryAllocationPolicy" - ) - - return QuotaAllowance( - can_run=can_run, - max_threads=self._get_max_threads(referrer) if can_run else 0, - explanation=decision_explanation, - is_throttled=False, - throttle_threshold=typing.cast(int, rate_limit_params.concurrent_limit), - rejection_threshold=typing.cast(int, rate_limit_params.concurrent_limit), - quota_used=rate_limit_stats.concurrent, - quota_unit=QUOTA_UNIT, - suggestion=NO_SUGGESTION if can_run else SUGGESTION, - ) - - def _update_quota_balance( - self, - tenant_ids: dict[str, str | int], - query_id: str, - result_or_error: QueryResultOrError, - ) -> None: - referrer = str(tenant_ids.get("referrer", "no_referrer")) - if not self._referrer_is_registered(referrer) and not self.is_cross_org_query(tenant_ids): - return - rate_limit_params = RateLimitParameters( - self.rate_limit_name, - referrer, - None, - # limit number does not matter for ending a query so I just picked 22 - 22, - ) - self._end_query(query_id, rate_limit_params, result_or_error) diff --git a/snuba/query/allocation_policies/per_referrer.py b/snuba/query/allocation_policies/per_referrer.py index 9eeddeeb1..40e6c8fa0 100644 --- a/snuba/query/allocation_policies/per_referrer.py +++ b/snuba/query/allocation_policies/per_referrer.py @@ -5,14 +5,13 @@ from snuba.configs.configuration import Configuration from snuba.query.allocation_policies import ( NO_SUGGESTION, - QueryResultOrError, QuotaAllowance, ) from snuba.query.allocation_policies.concurrent_rate_limit import ( BaseConcurrentRateLimitAllocationPolicy, ) from snuba.redis import RedisClientKey, get_redis_client -from snuba.state.rate_limit import RateLimitParameters +from snuba.state.rate_limit import RateLimitParameters, RateLimitStats from snuba.utils.serializable_exception import JsonSerializable rds = get_redis_client(RedisClientKey.RATE_LIMITER) @@ -20,7 +19,7 @@ logger = logging.getLogger("snuba.query.allocation_policy_per_referrer") _DEFAULT_CONCURRENT_REQUEST_PER_REFERRER = 100 -_REFERRER_CONCURRENT_OVERRIDE = -1 +_UNSET_CONCURRENT_LIMIT = -1 _REQUESTS_THROTTLE_DIVIDER = 1.5 _THREADS_THROTTLE_DIVIDER = 2 @@ -59,15 +58,14 @@ def _additional_config_definitions(self) -> list[Configuration]: default=_DEFAULT_CONCURRENT_REQUEST_PER_REFERRER, ), Configuration( - name="referrer_concurrent_override", - description="""override the concurrent limit for a referrer""", + name="concurrent_limit", + description="""maximum concurrent queries per referrer, the same setting name as ConcurrentRateLimitAllocationPolicy. Per-referrer values go in match blocks. Unset (-1) falls back to default_concurrent_request_per_referrer""", value_type=int, - param_types={"referrer": str}, - default=_REFERRER_CONCURRENT_OVERRIDE, + default=_UNSET_CONCURRENT_LIMIT, ), Configuration( name="requests_throttle_divider", - description="default_concurrent_request_per_referrer divided by this value will be the threshold at which we will decrease the number of threads (THROTTLED_THREADS) used to execute queries", + description="the concurrent limit divided by this value will be the threshold at which we will decrease the number of threads (THROTTLED_THREADS) used to execute queries", value_type=float, default=_REQUESTS_THROTTLE_DIVIDER, ), @@ -79,33 +77,33 @@ def _additional_config_definitions(self) -> list[Configuration]: ), ] - def _get_concurrent_limit(self, referrer: str) -> int: - concurrent_override = int( - self.get_config_value("referrer_concurrent_override", {"referrer": referrer}) - ) - default_concurrent_value = int( - self.get_config_value("default_concurrent_request_per_referrer") + def _get_concurrent_limit(self) -> int: + concurrent_limit = int(self.get_config_value("concurrent_limit")) + if concurrent_limit != _UNSET_CONCURRENT_LIMIT: + return concurrent_limit + return int(self.get_config_value("default_concurrent_request_per_referrer")) + + def _rate_limit_params(self, tenant_ids: dict[str, str | int]) -> RateLimitParameters: + referrer = str(tenant_ids.get("referrer", "no_referrer")) + return RateLimitParameters( + self.rate_limit_name, referrer, None, self._get_concurrent_limit() ) - return concurrent_override if concurrent_override != -1 else default_concurrent_value - def _get_quota_allowance( - self, tenant_ids: dict[str, str | int], query_id: str + def _allowance( + self, + tenant_ids: dict[str, str | int], + rate_limit_params: RateLimitParameters, + rate_limit_stats: RateLimitStats, + can_run: bool, + why: str, ) -> QuotaAllowance: referrer = str(tenant_ids.get("referrer", "no_referrer")) - concurrent_limit = self._get_concurrent_limit(referrer) - rate_limit_params = RateLimitParameters( - self.rate_limit_name, referrer, None, concurrent_limit - ) - rate_limit_stats, can_run, explanation = self._is_within_rate_limit( - query_id, - rate_limit_params, - ) assert rate_limit_params.concurrent_limit is not None, "concurrent_limit must be set" num_threads = self.max_threads requests_throttle_threshold = max( 1, int( - self.get_config_value("default_concurrent_request_per_referrer") + rate_limit_params.concurrent_limit // self.get_config_value("requests_throttle_divider") ), ) @@ -122,7 +120,7 @@ def _get_quota_allowance( tags={"referrer": referrer}, ) decision_explanation: dict[str, JsonSerializable] = { - "reason": explanation, + "reason": why, "policy": self.rate_limit_name, "referrer": referrer, } @@ -138,19 +136,3 @@ def _get_quota_allowance( quota_unit=QUOTA_UNIT, suggestion=NO_SUGGESTION if can_run else SUGGESTION, ) - - def _update_quota_balance( - self, - tenant_ids: dict[str, str | int], - query_id: str, - result_or_error: QueryResultOrError, - ) -> None: - referrer = str(tenant_ids.get("referrer", "no_referrer")) - rate_limit_params = RateLimitParameters( - self.rate_limit_name, - referrer, - None, - # limit number does not matter for ending a query so I just picked 22 - 22, - ) - self._end_query(query_id, rate_limit_params, result_or_error) diff --git a/tests/datasets/configuration/test_storage_loader.py b/tests/datasets/configuration/test_storage_loader.py index 580af38a9..17f0307b5 100644 --- a/tests/datasets/configuration/test_storage_loader.py +++ b/tests/datasets/configuration/test_storage_loader.py @@ -79,7 +79,7 @@ def test_processor_with_constructor(self) -> None: delete_allocation_policies: - - name: DeleteConcurrentRateLimitAllocationPolicy + name: ConcurrentRateLimitAllocationPolicy """ with tempfile.TemporaryDirectory() as tmpdirname: filename = os.path.join(tmpdirname, "file.yaml") @@ -127,7 +127,7 @@ def test_processor_with_constructor(self) -> None: assert len(delete_policies := storage.get_delete_allocation_policies()) == 1 assert {p.class_name() for p in delete_policies} == { - "DeleteConcurrentRateLimitAllocationPolicy", + "ConcurrentRateLimitAllocationPolicy", } def test_column_parser(self) -> None: diff --git a/tests/query/allocation_policies/test_allocation_policy_base.py b/tests/query/allocation_policies/test_allocation_policy_base.py index ff3edbf94..63b1ca768 100644 --- a/tests/query/allocation_policies/test_allocation_policy_base.py +++ b/tests/query/allocation_policies/test_allocation_policy_base.py @@ -28,7 +28,6 @@ from snuba.query.allocation_policies.concurrent_rate_limit import ( ConcurrentRateLimitAllocationPolicy, ) -from snuba.query.allocation_policies.cross_org import CrossOrgQueryAllocationPolicy from snuba.query.allocation_policies.per_referrer import ReferrerGuardRailPolicy from snuba.utils.metrics.backends.testing import get_recorded_metric_calls from tests.configs.component_config import ( @@ -508,7 +507,6 @@ def get_all_allocation_policy_classes(self) -> list[type[AllocationPolicy]]: ReferrerGuardRailPolicy, ConcurrentRateLimitAllocationPolicy, BytesScannedRejectingPolicy, - CrossOrgQueryAllocationPolicy, BytesScannedWindowAllocationPolicy, ] diff --git a/tests/query/allocation_policies/test_attachment.py b/tests/query/allocation_policies/test_attachment.py index 7c671059d..c52db21d2 100644 --- a/tests/query/allocation_policies/test_attachment.py +++ b/tests/query/allocation_policies/test_attachment.py @@ -16,12 +16,6 @@ {"name": "BytesScannedRejectingPolicy", "is_enforced": 0}, {"name": "BytesScannedWindowAllocationPolicy", "is_enforced": 0}, {"name": "ConcurrentRateLimitAllocationPolicy", "is_enforced": 0}, - { - "name": "CrossOrgQueryAllocationPolicy", - "is_enforced": 0, - "cross_org_referrer_limits": {"some.referrer": {"max_threads": 2, "concurrent_limit": 4}}, - }, - {"name": "DeleteConcurrentRateLimitAllocationPolicy", "is_enforced": 0}, {"name": "PassthroughPolicy", "is_enforced": 0}, {"name": "ReferrerGuardRailPolicy", "is_enforced": 0}, ] diff --git a/tests/query/allocation_policies/test_concurrent_rate_limit_policy.py b/tests/query/allocation_policies/test_concurrent_rate_limit_policy.py index bd2d667fc..e7d4fea89 100644 --- a/tests/query/allocation_policies/test_concurrent_rate_limit_policy.py +++ b/tests/query/allocation_policies/test_concurrent_rate_limit_policy.py @@ -5,6 +5,7 @@ import pytest +from snuba.configs.configuration import ResourceIdentifier from snuba.datasets.storage import StorageKey from snuba.query.allocation_policies import ( AllocationPolicyViolations, @@ -161,166 +162,43 @@ def test_tenant_selection(policy: ConcurrentRateLimitAllocationPolicy): policy._get_tenant_key_and_value({}) -OVERRIDE_TEST_CASES = [ - pytest.param( - [("organization_override", 1, {"organization_id": 123})], - {"organization_id": 123}, - {"organization_id__123": 1}, - 1, - id="organization_override", - ), - pytest.param( - [("organization_override", 1, {"organization_id": 123})], - {"organization_id": 456}, - {}, - MAX_CONCURRENT_QUERIES, - id="non-matching tenant_id", - ), - pytest.param( - [ - ( - "referrer_organization_override", - 1, - {"referrer": "abcd", "organization_id": 456}, - ) - ], - {"organization_id": 456, "referrer": "abcd"}, - {"organization_id__456|referrer__abcd": 1}, - 1, - ), - pytest.param( - [ - ("referrer_project_override", 1, {"referrer": "abcd", "project_id": 134}), - ("project_override", 4, {"project_id": 134}), - ], - {"organization_id": 456, "referrer": "abcd", "project_id": 134}, - {"project_id__134|referrer__abcd": 1, "project_id__134": 4}, - 1, - ), - pytest.param( - [ - ( - "referrer_organization_override", - 1, - {"referrer": "abcd", "organization_id": 123}, - ), - ], - {"organization_id": 123, "referrer": "abcd", "project_id": 134}, - {"organization_id__123|referrer__abcd": 1}, - 1, - id="referrer_organization_override", - ), - pytest.param( - [ - ( - "referrer_project_override", - 1, - {"referrer": "abcd", "project_id": 456}, - ), - ], - {"organization_id": 123, "referrer": "abcd", "project_id": 456}, - {"project_id__456|referrer__abcd": 1}, - 1, - id="referrer_organization_override", - ), - pytest.param( - [ - ( - "referrer_project_override", - MAX_CONCURRENT_QUERIES * 2, - {"referrer": "abcd", "project_id": 456}, - ), - ], - {"organization_id": 123, "referrer": "abcd", "project_id": 456}, - {"project_id__456|referrer__abcd": MAX_CONCURRENT_QUERIES * 2}, - MAX_CONCURRENT_QUERIES * 2, - id="override to a greater number", - ), -] - - -@pytest.mark.redis_db @pytest.mark.parametrize( - "overrides,tenant_ids,expected_overrides,expected_concurrent_limit", - OVERRIDE_TEST_CASES, + ("key_by", "tenant_ids", "bucket"), + [ + pytest.param("tenant", {"organization_id": 1, "project_id": 2}, "2", id="default, project"), + pytest.param("tenant", {"organization_id": 1}, "1", id="default, org"), + pytest.param( + "organization", + {"organization_id": 1, "project_id": 2, "referrer": "r"}, + "organization_id__1", + id="organization", + ), + pytest.param( + "organization_referrer", + {"organization_id": 1, "project_id": 2, "referrer": "r"}, + "organization_id__1|referrer__r", + id="organization_referrer", + ), + pytest.param("organization", {"project_id": 2}, "2", id="no org falls back"), + pytest.param("typo", {"organization_id": 1, "project_id": 2}, "2", id="unknown value"), + ], ) -def test_apply_overrides( - policy: ConcurrentRateLimitAllocationPolicy, - overrides, - tenant_ids, - expected_overrides, - expected_concurrent_limit, -) -> None: - for override in overrides: - set_component_config(policy, *override) - for i in range(expected_concurrent_limit): - policy.get_quota_allowance(tenant_ids=tenant_ids, query_id=f"{i}") - allowance = policy.get_quota_allowance( - tenant_ids=tenant_ids, query_id=f"{expected_concurrent_limit + 1}" +def test_key_by_bucket(key_by: str, tenant_ids: dict[str, str | int], bucket: str) -> None: + policy = ConcurrentRateLimitAllocationPolicy( + ResourceIdentifier(StorageKey("errors")), key_by=key_by ) - assert not allowance.can_run and allowance.max_threads == 0 - assert allowance.explanation["overrides"] == expected_overrides + # The default must keep today's bucket, or every running count resets. + assert policy._get_bucket(tenant_ids) == bucket @pytest.mark.redis_db -def test_override_isolation( - policy: ConcurrentRateLimitAllocationPolicy, -) -> None: - override_concurrent_limit = 1 - project_id = 1234 - overridden_referrer = "overridden_referrer" - set_component_config( - policy, - "referrer_project_override", - override_concurrent_limit, - {"project_id": project_id, "referrer": overridden_referrer}, - ) - for i in range(MAX_CONCURRENT_QUERIES): - policy.get_quota_allowance( - tenant_ids={"project_id": project_id, "referrer": "a_different_referrer"}, - query_id=str(i), - ) - # query the override referrer, it should not reject because overrides are counted on their own - assert policy.get_quota_allowance( - tenant_ids={"project_id": project_id, "referrer": overridden_referrer}, - query_id="uniq_string_1", - ).can_run - - # finish the overridden referrer query, make sure another one can run - policy.update_quota_balance( - tenant_ids={"project_id": project_id, "referrer": overridden_referrer}, - query_id="uniq_string_1", - result_or_error=_RESULT_SUCCESS, - ) - try: - policy.get_quota_allowance( - tenant_ids={"project_id": project_id, "referrer": overridden_referrer}, - query_id="uniq_string_2", - ) - except Exception: - pytest.fail("overridden query was finished, another one should have been able to run") - - # finish a non-overidden query - policy.update_quota_balance( - tenant_ids={"project_id": project_id, "referrer": "a_different_referrer"}, - query_id="1", - result_or_error=_RESULT_SUCCESS, +def test_key_by_organization_shares_one_count() -> None: + policy = ConcurrentRateLimitAllocationPolicy( + ResourceIdentifier(StorageKey("errors")), key_by="organization", concurrent_limit=1 ) - - # our overridden referrer should still error because an update to the non-overridden limits shouldn't affect the overridden ones - assert not policy.get_quota_allowance( - tenant_ids={"project_id": project_id, "referrer": overridden_referrer}, - query_id="uniq_string_3", - ).can_run - - -def test_pass_through(policy: ConcurrentRateLimitAllocationPolicy) -> None: - ## should not be blocked because the subscriptions_executor referrer is not rate limited - for i in range(MAX_CONCURRENT_QUERIES * 2): - assert policy.get_quota_allowance( - tenant_ids={"referrer": "subscriptions_executor", "project_id": 1234}, - query_id=f"abc{i}", - ).can_run + assert policy.get_quota_allowance({"organization_id": 1, "project_id": 2}, "q1").can_run + assert not policy.get_quota_allowance({"organization_id": 1, "project_id": 3}, "q2").can_run + assert policy.get_quota_allowance({"organization_id": 9, "project_id": 3}, "q3").can_run @pytest.mark.redis_db diff --git a/tests/query/allocation_policies/test_cross_org_policy.py b/tests/query/allocation_policies/test_cross_org_policy.py deleted file mode 100644 index b311e9aea..000000000 --- a/tests/query/allocation_policies/test_cross_org_policy.py +++ /dev/null @@ -1,170 +0,0 @@ -import pytest - -from snuba.query.allocation_policies import QueryResultOrError -from snuba.query.allocation_policies.cross_org import CrossOrgQueryAllocationPolicy -from snuba.web import QueryResult -from tests.configs.component_config import set_component_config - -_RESULT_SUCCESS = QueryResultOrError( - QueryResult( - result={"profile": {"bytes": 42069}}, - extra={"stats": {}, "sql": "", "experiments": {}}, - ), - error=None, -) - - -class TestCrossOrgQueryAllocationPolicy: - @pytest.mark.redis_db - def test_policy_pass_basic(self): - policy = CrossOrgQueryAllocationPolicy.from_kwargs( - **{ - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": { - "concurrent_limit": 1, - "max_threads": 1, - }, - }, - } - ) - unimportant_allowance = policy.get_quota_allowance( - tenant_ids={"referrer": "some_referrer"}, query_id="1" - ) - assert unimportant_allowance.can_run is True - assert unimportant_allowance.max_threads == 10 - assert unimportant_allowance.explanation == { - "reason": "pass_through", - "storage_key": "metrics_counters", - } - cross_org_allowance = policy.get_quota_allowance( - tenant_ids={"referrer": "statistical_detectors"}, query_id="2" - ) - assert cross_org_allowance.can_run is True - assert cross_org_allowance.max_threads == 1 - - quota_allowance = policy.get_quota_allowance( - tenant_ids={"referrer": "statistical_detectors"}, query_id="3" - ) - assert not quota_allowance.can_run and quota_allowance.max_threads == 0 - policy.update_quota_balance( - tenant_ids={"referrer": "statistical_detectors"}, - query_id="2", - result_or_error=_RESULT_SUCCESS, - ) - assert policy.get_quota_allowance( - tenant_ids={"referrer": "statistical_detectors"}, query_id="3" - ).can_run - - @pytest.mark.parametrize( - "config", - [ - { - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": { - "max_threads": 1, - }, - }, - }, - { - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": { - "concurrent_limit": 2, - }, - }, - }, - { - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": {"max_threads": {}, "concurrent_limit": 2}, - }, - }, - ], - ) - def test_bad_config(self, config) -> None: - with pytest.raises(ValueError): - CrossOrgQueryAllocationPolicy.from_kwargs(**config) - - @pytest.mark.redis_db - def test_override(self): - policy = CrossOrgQueryAllocationPolicy.from_kwargs( - **{ - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": { - "concurrent_limit": 1, - "max_threads": 1, - }, - }, - } - ) - set_component_config( - policy, - "referrer_concurrent_override", - 0, - {"referrer": "statistical_detectors"}, - ) - - quota_allowance = policy.get_quota_allowance( - tenant_ids={"referrer": "statistical_detectors"}, query_id="2" - ) - assert not quota_allowance.can_run and quota_allowance.max_threads == 0 - - @pytest.mark.redis_db - def test_respects_max_threads_config(self) -> None: - policy = CrossOrgQueryAllocationPolicy.from_kwargs( - **{ - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": { - "concurrent_limit": 10, - "max_threads": 4, - }, - }, - } - ) - set_component_config(policy, "max_threads", 7) - assert ( - policy.get_quota_allowance( - tenant_ids={"referrer": "some_referrer"}, query_id="1" - ).max_threads - == 7 - ) - assert ( - policy.get_quota_allowance( - tenant_ids={"referrer": "statistical_detectors"}, query_id="2" - ).max_threads - == 4 - ) - - @pytest.mark.redis_db - def test_throttle_cross_org_query_with_unregistered_referrer(self): - policy = CrossOrgQueryAllocationPolicy.from_kwargs( - **{ - "storage_key": "metrics_counters", - "cross_org_referrer_limits": { - "statistical_detectors": { - "concurrent_limit": 1, - "max_threads": 1, - }, - }, - } - ) - allowance = policy.get_quota_allowance( - tenant_ids={"referrer": "unregistered", "cross_org_query": 1}, - query_id="1", - ) - assert allowance.can_run is True - assert allowance.max_threads == 1 - - allowance = policy.get_quota_allowance( - tenant_ids={"referrer": "unregistered", "cross_org_query": 1}, - query_id="2", - ) - assert not allowance.can_run - assert ( - allowance.explanation["cross_org_query"] - == "This referrer is not registered for the current storage metrics_counters, if you want to increase its limits, register it in the yaml of the CrossOrgQueryAllocationPolicy" - ) diff --git a/tests/query/allocation_policies/test_per_referrer.py b/tests/query/allocation_policies/test_per_referrer.py index 95234c496..916ec1f6d 100644 --- a/tests/query/allocation_policies/test_per_referrer.py +++ b/tests/query/allocation_policies/test_per_referrer.py @@ -78,15 +78,7 @@ def test_throttle(self) -> None: @pytest.mark.redis_db def test_override(self): policy = ReferrerGuardRailPolicy.from_kwargs( - **{ - "storage_key": "metrics_counters", - } - ) - set_component_config( - policy, - "referrer_concurrent_override", - 0, - {"referrer": "statistical_detectors"}, + storage_key="metrics_counters", concurrent_limit=0 ) quota_allowance = policy.get_quota_allowance( diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py index 889b07f98..9fd1f61c0 100644 --- a/tests/test_snql_api.py +++ b/tests/test_snql_api.py @@ -28,7 +28,6 @@ from snuba.querylog.query_metadata import QueryStatus from snuba.utils.metrics.backends.testing import get_recorded_metric_calls from tests.base import BaseApiTest -from tests.configs.component_config import set_component_config from tests.fixtures import get_raw_event, get_raw_transaction from tests.helpers import override_entity_column_validator, write_unprocessed_events @@ -285,9 +284,8 @@ def test_project_rate_limiting(self) -> None: ) policy = ConcurrentRateLimitAllocationPolicy( - ResourceIdentifier(StorageKey("errors")), + ResourceIdentifier(StorageKey("errors")), concurrent_limit=0 ) - set_component_config(policy, "project_override", 0, {"project_id": self.project_id}) with patch( "snuba.web.db_query._get_allocation_policies", diff --git a/tests/web/rpc/v1/test_storage_routing.py b/tests/web/rpc/v1/test_storage_routing.py index 2d4fa9fb5..1ee320aac 100644 --- a/tests/web/rpc/v1/test_storage_routing.py +++ b/tests/web/rpc/v1/test_storage_routing.py @@ -362,7 +362,6 @@ def _output_metrics(self, routing_context: RoutingContext) -> None: "can_run": True, "max_threads": 10, "explanation": { - "overrides": {}, "reason": "within limit", "storage_key": "EAP", },