From fece2c88428c061758f150ff532a93605a7ac477 Mon Sep 17 00:00:00 2001 From: Dani Ezzeddine Date: Thu, 29 Oct 2026 14:01:42 -0700 Subject: [PATCH] ref(allocation-policies): Template method for the concurrency policies ConcurrentRateLimitAllocationPolicy, ReferrerGuardRailPolicy and CrossOrgQueryAllocationPolicy each re-implemented the same count-check-end flow. BaseConcurrentRateLimitAllocationPolicy now runs it; subclasses supply the bucket and limit, which queries they skip, and how a decision becomes an allowance. Buckets, limits, explanations and decisions are unchanged. --- .../concurrent_rate_limit.py | 88 ++++++++++++++----- snuba/query/allocation_policies/cross_org.py | 59 ++++++------- .../query/allocation_policies/per_referrer.py | 44 ++++------ 3 files changed, 107 insertions(+), 84 deletions(-) diff --git a/snuba/query/allocation_policies/concurrent_rate_limit.py b/snuba/query/allocation_policies/concurrent_rate_limit.py index bf4cce405..c4d56f52a 100644 --- a/snuba/query/allocation_policies/concurrent_rate_limit.py +++ b/snuba/query/allocation_policies/concurrent_rate_limit.py @@ -71,6 +71,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]: @@ -223,9 +267,11 @@ def _get_rate_limit_params( overrides, ) - def _get_quota_allowance( - self, tenant_ids: dict[str, str | int], query_id: str - ) -> QuotaAllowance: + def _rate_limit_params(self, tenant_ids: dict[str, str | int]) -> RateLimitParameters: + rate_limit_params, _ = self._get_rate_limit_params(tenant_ids) + return rate_limit_params + + 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,36 +296,32 @@ 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, "overrides": self._get_overrides(tenant_ids)}, 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 diff --git a/snuba/query/allocation_policies/cross_org.py b/snuba/query/allocation_policies/cross_org.py index 2f582960c..6d7a7cff9 100644 --- a/snuba/query/allocation_policies/cross_org.py +++ b/snuba/query/allocation_policies/cross_org.py @@ -9,14 +9,13 @@ 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.state.rate_limit import RateLimitParameters, RateLimitStats from snuba.utils.serializable_exception import JsonSerializable DEFAULT_CONCURRENT_QUERIES_LIMIT = 22 @@ -113,11 +112,14 @@ def _get_concurrent_limit(self, referrer: str) -> int: 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: + def _is_out_of_scope(self, tenant_ids: dict[str, str | int]) -> bool: 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 not self._referrer_is_registered(referrer) and not self.is_cross_org_query( + tenant_ids + ) + + def _pass_through_allowance(self, tenant_ids: dict[str, str | int]) -> QuotaAllowance | None: + if self._is_out_of_scope(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, @@ -130,16 +132,27 @@ def _get_quota_allowance( quota_unit=NO_UNITS, suggestion=NO_SUGGESTION, ) + return None - 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, + def _skips_balance_update(self, tenant_ids: dict[str, str | int]) -> bool: + return self._is_out_of_scope(tenant_ids) + + 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(referrer) ) - decision_explanation: dict[str, JsonSerializable] = {"reason": explanation} + + 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")) + decision_explanation: dict[str, JsonSerializable] = {"reason": why} 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" @@ -156,21 +169,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")) - 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..5fdfbdab4 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) @@ -88,18 +87,21 @@ def _get_concurrent_limit(self, referrer: str) -> int: ) 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 - ) -> QuotaAllowance: + def _rate_limit_params(self, tenant_ids: dict[str, str | int]) -> RateLimitParameters: 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, + return RateLimitParameters( + self.rate_limit_name, referrer, None, self._get_concurrent_limit(referrer) ) + + 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")) assert rate_limit_params.concurrent_limit is not None, "concurrent_limit must be set" num_threads = self.max_threads requests_throttle_threshold = max( @@ -122,7 +124,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 +140,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)