diff --git a/snuba/query/allocation_policies/bytes_scanned_base.py b/snuba/query/allocation_policies/bytes_scanned_base.py new file mode 100644 index 000000000..f3f34676f --- /dev/null +++ b/snuba/query/allocation_policies/bytes_scanned_base.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import time + +from sentry_redis_tools.sliding_windows_rate_limiter import ( + GrantedQuota, + Quota, + RedisSlidingWindowRateLimiter, + RequestedQuota, +) + +from snuba.query.allocation_policies import AllocationPolicy +from snuba.redis import RedisClientKey, get_redis_client + +_RATE_LIMITER = RedisSlidingWindowRateLimiter(get_redis_client(RedisClientKey.RATE_LIMITER)) + + +class BaseBytesScannedPolicy(AllocationPolicy): + """Sliding-window bookkeeping shared by the bytes-scanned policies. Each subclass + decides which window a query counts against, its limit and what to do over it.""" + + WINDOW_SECONDS = 10 * 60 + WINDOW_GRANULARITY_SECONDS = 60 + + def _quota(self, window_prefix: str, limit: int) -> Quota: + return Quota( + window_seconds=self.WINDOW_SECONDS, + granularity_seconds=self.WINDOW_GRANULARITY_SECONDS, + limit=limit, + prefix_override=window_prefix, + ) + + def _granted_bytes(self, window_prefix: str, limit: int, requested: int) -> int: + """How many of `requested` bytes the window still has. Uses up nothing.""" + _, granted_quotas = _RATE_LIMITER.check_within_quotas( + [RequestedQuota(self.component_name(), requested, [self._quota(window_prefix, limit)])] + ) + return int(granted_quotas[0].granted) + + def _record_bytes( + self, request_name: str, window_prefix: str, limit: int, bytes_scanned: int + ) -> None: + # The query was already allowed, so record what it used as granted. + _RATE_LIMITER.use_quotas( + [RequestedQuota(request_name, bytes_scanned, [self._quota(window_prefix, limit)])], + grants=[GrantedQuota(request_name, granted=bytes_scanned, reached_quotas=[])], + timestamp=int(time.time()), + ) diff --git a/snuba/query/allocation_policies/bytes_scanned_rejecting_policy.py b/snuba/query/allocation_policies/bytes_scanned_rejecting_policy.py index 8f590383c..689724f65 100644 --- a/snuba/query/allocation_policies/bytes_scanned_rejecting_policy.py +++ b/snuba/query/allocation_policies/bytes_scanned_rejecting_policy.py @@ -1,16 +1,8 @@ from __future__ import annotations import logging -import time from typing import Any, cast -from sentry_redis_tools.sliding_windows_rate_limiter import ( - GrantedQuota, - Quota, - RedisSlidingWindowRateLimiter, - RequestedQuota, -) - from snuba.clickhouse.error_codes import ErrorCodes from snuba.clickhouse.errors import ClickhouseError from snuba.configs.configuration import Configuration @@ -19,12 +11,11 @@ MAX_THRESHOLD, NO_SUGGESTION, PASS_THROUGH_REFERRERS_SUGGESTION, - AllocationPolicy, InvalidTenantsForAllocationPolicy, QueryResultOrError, QuotaAllowance, ) -from snuba.redis import RedisClientKey, get_redis_client +from snuba.query.allocation_policies.bytes_scanned_base import BaseBytesScannedPolicy logger = logging.getLogger("snuba.query.bytes_scanned_window_policy") @@ -36,7 +27,6 @@ UNREASONABLY_LARGE_NUMBER_OF_BYTES_SCANNED_PER_QUERY = int(1e12) -_RATE_LIMITER = RedisSlidingWindowRateLimiter(get_redis_client(RedisClientKey.RATE_LIMITER)) DEFAULT_OVERRIDE_LIMIT = -1 TERABYTE = 10**12 DEFAULT_BYTES_SCANNED_LIMIT = int(1.28 * TERABYTE) @@ -47,7 +37,7 @@ SUGGESTION = "The feature, organization/project is scanning too many bytes, this usually means they are abusing that API" -class BytesScannedRejectingPolicy(AllocationPolicy): +class BytesScannedRejectingPolicy(BaseBytesScannedPolicy): """For every query that comes in, keep track of the amount of bytes scanned for every (project_id|organization_id, referrer) combination in the last 10 minutes (sliding window). If a specific combination scans too @@ -56,9 +46,6 @@ class BytesScannedRejectingPolicy(AllocationPolicy): cross-project queries use the organization_id, single project queries use the project_id """ - WINDOW_SECONDS = 10 * 60 - WINDOW_GRANULARITY_SECONDS = 60 - def _additional_config_definitions(self) -> list[Configuration]: # Overrides are checked in order of specificity; the first one set wins. # For organization_id queries: @@ -291,30 +278,16 @@ def _get_quota_allowance( throttle_threshold = max( 1, int(scan_limit // self.get_config_value("bytes_throttle_divider")) ) - timestamp, granted_quotas = _RATE_LIMITER.check_within_quotas( - [ - RequestedQuota( - self.component_name(), - # request a big number because we don't know how much we actually - # will use in this query. this doesn't use up any quota, we just want to know how much is left - UNREASONABLY_LARGE_NUMBER_OF_BYTES_SCANNED_PER_QUERY, - [ - Quota( - # TODO: Make window configurable but I don't know exactly how the rate limiter - # reacts to such changes - window_seconds=self.WINDOW_SECONDS, - granularity_seconds=self.WINDOW_GRANULARITY_SECONDS, - limit=scan_limit, - prefix_override=f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}-{referrer}", - ) - ], - ), - ] + # request a big number because we don't know how much we actually + # will use in this query. this doesn't use up any quota, we just want to know how much is left + granted = self._granted_bytes( + f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}-{referrer}", + scan_limit, + UNREASONABLY_LARGE_NUMBER_OF_BYTES_SCANNED_PER_QUERY, ) explanation: dict[str, Any] = {} - granted_quota = granted_quotas[0] - used_quota = scan_limit - granted_quota.granted - if granted_quota.granted <= 0: + used_quota = scan_limit - granted + if granted <= 0: if self.get_config_value("limit_bytes_instead_of_rejecting"): max_bytes_to_read = int( scan_limit / self.get_config_value("max_bytes_to_read_scan_limit_divider") @@ -324,7 +297,7 @@ def _get_quota_allowance( ] = f"""{customer_tenant_key} {customer_tenant_value} is over the bytes scanned limit of {scan_limit} for referrer {referrer}. The query will be limited to {max_bytes_to_read} bytes """ - explanation["granted_quota"] = granted_quota.granted + explanation["granted_quota"] = granted explanation["limit"] = scan_limit # This is technically a high cardinality tag value however these rejections # should not happen often therefore it should be safe to output these rejections as metris @@ -354,7 +327,7 @@ def _get_quota_allowance( ] = f"""{customer_tenant_key} {customer_tenant_value} is over the bytes scanned limit of {scan_limit} for referrer {referrer}. This policy is exceeded when a customer is abusing a specific feature in a way that puts load on clickhouse. If this is happening to "many customers, that may mean the feature is written in an inefficient way""" - explanation["granted_quota"] = granted_quota.granted + explanation["granted_quota"] = granted explanation["limit"] = scan_limit # This is technically a high cardinality tag value however these rejections # should not happen often therefore it should be safe to output these rejections as metris @@ -376,7 +349,7 @@ def _get_quota_allowance( ) # this checks to see if you reached the throttle threshold - if granted_quota.granted < scan_limit - throttle_threshold: + if granted < scan_limit - throttle_threshold: self.metrics.increment( "bytes_scanned_queries_throttled", tags={"referrer": str(tenant_ids.get("referrer", "no_referrer"))}, @@ -462,29 +435,9 @@ def _update_quota_balance( ): return scan_limit = self.__get_scan_limit(customer_tenant_key, customer_tenant_value, referrer) - # we can assume that the requested quota was granted (because it was) - # we just need to update the quota with however many bytes were consumed - _RATE_LIMITER.use_quotas( - [ - RequestedQuota( - f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}", - bytes_scanned, - [ - Quota( - window_seconds=self.WINDOW_SECONDS, - granularity_seconds=self.WINDOW_GRANULARITY_SECONDS, - limit=scan_limit, - prefix_override=f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}-{referrer}", - ) - ], - ) - ], - grants=[ - GrantedQuota( - f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}", - granted=bytes_scanned, - reached_quotas=[], - ) - ], - timestamp=int(time.time()), + self._record_bytes( + f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}", + f"{self.component_name()}-{customer_tenant_key}-{customer_tenant_value}-{referrer}", + scan_limit, + bytes_scanned, ) diff --git a/snuba/query/allocation_policies/bytes_scanned_window_policy.py b/snuba/query/allocation_policies/bytes_scanned_window_policy.py index 9876fe4e3..7adc0e0ad 100644 --- a/snuba/query/allocation_policies/bytes_scanned_window_policy.py +++ b/snuba/query/allocation_policies/bytes_scanned_window_policy.py @@ -1,16 +1,8 @@ from __future__ import annotations import logging -import time from typing import Any, cast -from sentry_redis_tools.sliding_windows_rate_limiter import ( - GrantedQuota, - Quota, - RedisSlidingWindowRateLimiter, - RequestedQuota, -) - from snuba.configs.configuration import Configuration from snuba.query.allocation_policies import ( CROSS_ORG_SUGGESTION, @@ -18,11 +10,10 @@ NO_SUGGESTION, NO_UNITS, PASS_THROUGH_REFERRERS_SUGGESTION, - AllocationPolicy, QueryResultOrError, QuotaAllowance, ) -from snuba.redis import RedisClientKey, get_redis_client +from snuba.query.allocation_policies.bytes_scanned_base import BaseBytesScannedPolicy logger = logging.getLogger("snuba.query.bytes_scanned_window_policy") @@ -75,17 +66,13 @@ UNREASONABLY_LARGE_NUMBER_OF_BYTES_SCANNED_PER_QUERY = int(1e10) -_RATE_LIMITER = RedisSlidingWindowRateLimiter(get_redis_client(RedisClientKey.RATE_LIMITER)) DEFAULT_OVERRIDE_LIMIT = -1 DEFAULT_BYTES_SCANNED_LIMIT = 10000000 QUOTA_UNIT = "bytes" SUGGESTION = "The feature, organization/project is scanning too many bytes, this usually means they are abusing that API" -class BytesScannedWindowAllocationPolicy(AllocationPolicy): - WINDOW_SECONDS = 10 * 60 - WINDOW_GRANULARITY_SECONDS = 60 - +class BytesScannedWindowAllocationPolicy(BaseBytesScannedPolicy): def _additional_config_definitions(self) -> list[Configuration]: return [ Configuration( @@ -180,37 +167,23 @@ def _get_quota_allowance( if org_id is not None: org_limit_bytes_scanned = self.__get_org_limit_bytes_scanned(org_id) - timestamp, granted_quotas = _RATE_LIMITER.check_within_quotas( - [ - RequestedQuota( - self.component_name(), - # request a big number because we don't know how much we actually - # will use in this query. this doesn't use up any quota, we just want to know how much is left - UNREASONABLY_LARGE_NUMBER_OF_BYTES_SCANNED_PER_QUERY, - [ - Quota( - # TODO: Make window configurable but I don't know exactly how the rate limiter - # reacts to such changes - window_seconds=self.WINDOW_SECONDS, - granularity_seconds=self.WINDOW_GRANULARITY_SECONDS, - limit=org_limit_bytes_scanned, - prefix_override=f"{self.component_name()}-organization_id-{org_id}", - ) - ], - ), - ] + # request a big number because we don't know how much we actually + # will use in this query. this doesn't use up any quota, we just want to know how much is left + granted = self._granted_bytes( + f"{self.component_name()}-organization_id-{org_id}", + org_limit_bytes_scanned, + UNREASONABLY_LARGE_NUMBER_OF_BYTES_SCANNED_PER_QUERY, ) num_threads = self.max_threads explanation: dict[str, Any] = {} - granted_quota = granted_quotas[0] is_throttled = False - if granted_quota.granted <= 0: + if granted <= 0: is_throttled = True explanation["reason"] = ( f"organization {org_id} is over the bytes scanned limit of {org_limit_bytes_scanned}" ) explanation["is_enforced"] = self.is_enforced - explanation["granted_quota"] = granted_quota.granted + explanation["granted_quota"] = granted explanation["limit"] = org_limit_bytes_scanned if self.is_enforced: @@ -223,7 +196,7 @@ def _get_quota_allowance( is_throttled=is_throttled, throttle_threshold=org_limit_bytes_scanned, rejection_threshold=MAX_THRESHOLD, - quota_used=org_limit_bytes_scanned - granted_quota.granted, + quota_used=org_limit_bytes_scanned - granted, quota_unit=QUOTA_UNIT, suggestion=SUGGESTION, ) @@ -279,31 +252,11 @@ def _update_quota_balance( org_limit_bytes_scanned = self.__get_org_limit_bytes_scanned( tenant_ids.get("organization_id") ) - # we can assume that the requested quota was granted (because it was) - # we just need to update the quota with however many bytes were consumed - _RATE_LIMITER.use_quotas( - [ - RequestedQuota( - f"{self.component_name()}-organization_id-{tenant_ids['organization_id']}", - bytes_scanned, - [ - Quota( - window_seconds=self.WINDOW_SECONDS, - granularity_seconds=self.WINDOW_GRANULARITY_SECONDS, - limit=org_limit_bytes_scanned, - prefix_override=f"{self.component_name()}-organization_id-{tenant_ids['organization_id']}", - ) - ], - ) - ], - grants=[ - GrantedQuota( - f"{self.component_name()}-organization_id-{tenant_ids['organization_id']}", - granted=bytes_scanned, - reached_quotas=[], - ) - ], - timestamp=int(time.time()), + self._record_bytes( + f"{self.component_name()}-organization_id-{tenant_ids['organization_id']}", + f"{self.component_name()}-organization_id-{tenant_ids['organization_id']}", + org_limit_bytes_scanned, + bytes_scanned, ) def __get_org_limit_bytes_scanned(self, org_id: Any) -> int: