Skip to content

ref(allocation-policies): Template method for the concurrency policies #8641

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dezzeddi wants to merge 1 commit into from
Closed

ref(allocation-policies): Template method for the concurrency policies #8641

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
88 changes: 65 additions & 23 deletions snuba/query/allocation_policies/concurrent_rate_limit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down
59 changes: 27 additions & 32 deletions snuba/query/allocation_policies/cross_org.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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"
Expand All @@ -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)
44 changes: 15 additions & 29 deletions snuba/query/allocation_policies/per_referrer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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(
Expand All @@ -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,
}
Expand All @@ -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)
Loading