Skip to content

ref(allocation-policies): Share sliding-window bookkeeping between the bytes policies #8640

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): Share sliding-window bookkeeping between the bytes policies #8640

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
48 changes: 48 additions & 0 deletions snuba/query/allocation_policies/bytes_scanned_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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()),
)
83 changes: 18 additions & 65 deletions snuba/query/allocation_policies/bytes_scanned_rejecting_policy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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")

Expand All @@ -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)
Expand All @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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")
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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"))},
Expand Down Expand Up @@ -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,
)
79 changes: 16 additions & 63 deletions snuba/query/allocation_policies/bytes_scanned_window_policy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,28 +1,19 @@
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,
MAX_THRESHOLD,
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")

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand All @@ -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,
)
Expand Down Expand Up @@ -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:
Expand Down
Loading