diff --git a/snuba/query/allocation_policies/__init__.py b/snuba/query/allocation_policies/__init__.py index d5868154c2c..86970b05e6b 100644 --- a/snuba/query/allocation_policies/__init__.py +++ b/snuba/query/allocation_policies/__init__.py @@ -18,6 +18,7 @@ logger, ) from snuba.datasets.storages.storage_key import StorageKey +from snuba.state import get_config as get_runtime_config from snuba.utils.metrics.wrapper import MetricsWrapper from snuba.utils.registered_class import import_submodules_in_directory from snuba.utils.serializable_exception import JsonSerializable, SerializableException @@ -26,6 +27,13 @@ CAPMAN_PREFIX = "capman" CAPMAN_HASH = "capman" +# Global runtime config (visible in the runtime config admin UI, NOT Capacity +# Management) holding a comma-separated list of organization ids that are exempt +# from ALL allocation policies (rate limits), e.g. "1,42,1337". It is global +# rather than per-policy because it is meant to exempt an org from every policy +# at once. See AllocationPolicy.is_rate_limit_bypassed. +ORG_RATE_LIMIT_BYPASS_CONFIG = "organization_rate_limit_bypass_allowlist" + IS_ACTIVE = "is_active" IS_ENFORCED = "is_enforced" MAX_THREADS = "max_threads" @@ -424,6 +432,30 @@ def __eq__(self, other: Any) -> bool: def is_cross_org_query(self, tenant_ids: dict[str, str | int]) -> bool: return bool(tenant_ids.get("cross_org_query", False)) + def is_rate_limit_bypassed(self, tenant_ids: dict[str, str | int]) -> bool: + """Returns True if the query's organization is in the global rate limit + bypass allowlist. Orgs in this allowlist skip ALL allocation policies + entirely, in the same way an inactive policy is skipped. + + The allowlist is a single global runtime config shared across every + policy (see ORG_RATE_LIMIT_BYPASS_CONFIG), since it is meant to exempt an + org from all rate limits at once. It is stored as a comma-separated list + of organization ids, e.g. "1,42,1337". + """ + if not tenant_ids: + return False + org_id = tenant_ids.get("organization_id") + if org_id is None: + return False + # Use `is None` rather than a falsy check: runtime config coerces types, + # so a lone org id of `0` would be stored as the integer 0 and a `not + # allowlist` check would wrongly treat it as unset. + allowlist = get_runtime_config(ORG_RATE_LIMIT_BYPASS_CONFIG) + if allowlist is None: + return False + bypassed_orgs = {part.strip() for part in str(allowlist).split(",") if part.strip()} + return str(org_id) in bypassed_orgs + @classmethod def from_kwargs(cls, **kwargs: str) -> AllocationPolicy: required_tenant_types = kwargs.pop("required_tenant_types", None) @@ -457,7 +489,7 @@ def get_quota_allowance( for t, tid in tenant_ids.items(): span.set_data(f"tenant_ids.{t}", str(tid)) try: - if not self.is_active: + if not self.is_active or self.is_rate_limit_bypassed(tenant_ids): allowance = QuotaAllowance( can_run=True, max_threads=self.max_threads, @@ -547,7 +579,7 @@ def update_quota_balance( result_or_error: QueryResultOrError, ) -> None: try: - if not self.is_active: + if not self.is_active or self.is_rate_limit_bypassed(tenant_ids): return None return self._update_quota_balance(tenant_ids, query_id, result_or_error) except InvalidTenantsForAllocationPolicy: diff --git a/tests/query/allocation_policies/test_allocation_policy_base.py b/tests/query/allocation_policies/test_allocation_policy_base.py index a6c2273bbee..56defae7fa1 100644 --- a/tests/query/allocation_policies/test_allocation_policy_base.py +++ b/tests/query/allocation_policies/test_allocation_policy_base.py @@ -12,6 +12,7 @@ MAX_THRESHOLD, NO_SUGGESTION, NO_UNITS, + ORG_RATE_LIMIT_BYPASS_CONFIG, AllocationPolicy, AllocationPolicyConfig, InvalidTenantsForAllocationPolicy, @@ -524,6 +525,94 @@ def test_is_not_enforced() -> None: assert throttled_metrics[1].tags["is_enforced"] == "False" +@pytest.mark.redis_db +def test_org_rate_limit_bypass_allowlist() -> None: + MAX_THREADS = 100 + reject_policy = RejectingEverythingAllocationPolicy( + StorageKey("some_storage"), + [], + {"is_active": 1, "is_enforced": 1, "max_threads": MAX_THREADS}, + ) + allowlisted_tenant: dict[str, int | str] = { + "organization_id": 123, + "referrer": "some_referrer", + } + other_tenant: dict[str, int | str] = { + "organization_id": 999, + "referrer": "some_referrer", + } + + # By default the org is rejected by the (enforced) policy. + assert not reject_policy.get_quota_allowance(allowlisted_tenant, "deadbeef").can_run + + # Add the org to the global bypass allowlist. + set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "123,456") + + # Allowlisted org bypasses the policy entirely (passes through with full threads). + allowance = reject_policy.get_quota_allowance(allowlisted_tenant, "deadbeef") + assert allowance.can_run + assert allowance.max_threads == MAX_THREADS + + # An org not in the allowlist is still rejected. + assert not reject_policy.get_quota_allowance(other_tenant, "deadbeef").can_run + + +@pytest.mark.redis_db +def test_org_rate_limit_bypass_skips_update_quota_balance() -> None: + # BadlyWrittenAllocationPolicy raises in both _get_quota_allowance and + # _update_quota_balance, so a successful call proves those private methods + # were never invoked (i.e. the org bypassed the policy). + policy = BadlyWrittenAllocationPolicy( + StorageKey("some_storage"), + [], + {"is_active": 1, "is_enforced": 1}, + ) + tenant_ids: dict[str, int | str] = { + "organization_id": 123, + "referrer": "some_referrer", + } + result_or_error = QueryResultOrError( + query_result=QueryResult( + result={"profile": {"bytes": 420}}, + extra={"stats": {}, "sql": "", "experiments": {}}, + ), + error=None, + ) + + set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "123") + + # Neither the buggy _get_quota_allowance nor _update_quota_balance is reached. + assert policy.get_quota_allowance(tenant_ids, "deadbeef").can_run + policy.update_quota_balance(tenant_ids, "deadbeef", result_or_error) + + +@pytest.mark.redis_db +def test_org_rate_limit_bypass_no_org_id() -> None: + reject_policy = RejectingEverythingAllocationPolicy( + StorageKey("some_storage"), + [], + {"is_active": 1, "is_enforced": 1}, + ) + set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "123") + # tenant_ids without an organization_id should never be bypassed. + assert not reject_policy.get_quota_allowance({"referrer": "some_referrer"}, "deadbeef").can_run + + +@pytest.mark.redis_db +def test_org_rate_limit_bypass_zero_org_id() -> None: + # A lone org id of 0 is coerced to the integer 0 by runtime config; make + # sure it is still honored (i.e. not treated as an unset allowlist). + reject_policy = RejectingEverythingAllocationPolicy( + StorageKey("some_storage"), + [], + {"is_active": 1, "is_enforced": 1}, + ) + set_config(ORG_RATE_LIMIT_BYPASS_CONFIG, "0") + assert reject_policy.get_quota_allowance( + {"organization_id": 0, "referrer": "some_referrer"}, "deadbeef" + ).can_run + + @pytest.mark.redis_db def test_configs_with_delimiter_values() -> None: # test that configs with dots can be stored and read