feat(allocation_policies): pardon rejections when dynamic allocation policy is on

sentry-options/schemas/snuba/schema.json

@@ -274,6 +274,11 @@
       "default": false,
       "description": "When true, the storage routing strategy fetches ClickHouse cluster load info to inform routing decisions."
     },
+    "storage_routing.enable_dynamic_allocation_policy": {
+      "type": "boolean",
+      "default": false,
+      "description": "When true, allocation-policy rejections are pardoned. Load thresholds are not applied yet."
+    },
     "retention_days": {
       "type": "object",
       "properties": {

snuba/web/rpc/storage_routing/load_retriever.py → snuba/clusters/load_info.py

@@ -15,11 +15,12 @@
 from snuba.clusters.cluster import ClickhouseClientSettings, get_cluster
 from snuba.clusters.storage_sets import StorageSetKey
 from snuba.redis import RedisClientKey, get_redis_client
+from snuba.state.sentry_options import get_option
 from snuba.utils.metrics.wrapper import MetricsWrapper
 
 metrics = MetricsWrapper(
     environment.metrics,
-    "snuba.web.rpc.storage_routing.load_retriever",
+    "snuba.clusters.load_info",
 )
 
 
@@ -47,6 +48,9 @@ def from_dict(cls, load_info_dict: dict[str, float | int | None]) -> LoadInfo:
             }
         )
 
+    def should_pardon(self) -> bool:
+        return get_option("storage_routing.enable_dynamic_allocation_policy", False)
+
 
 def cache(
     ttl_secs: int = 60,
@@ -86,9 +90,17 @@ def wrapper(*args: Any, **kwargs: Any) -> LoadInfo:
     return decorator
 
 
-@cache(ttl_secs=60)
 def get_cluster_loadinfo(
     storage_set_key: StorageSetKey = StorageSetKey.EVENTS_ANALYTICS_PLATFORM,
+) -> LoadInfo | None:
+    if not get_option("storage_routing.enable_get_cluster_loadinfo", False):
+        return None
+    return _get_cluster_loadinfo(storage_set_key)
+
+
+@cache(ttl_secs=60)
+def _get_cluster_loadinfo(
+    storage_set_key: StorageSetKey = StorageSetKey.EVENTS_ANALYTICS_PLATFORM,
 ) -> LoadInfo:
     cluster_name = None
     try:

snuba/query/allocation_policies/__init__.py

@@ -4,13 +4,14 @@
 import os
 from abc import ABC, abstractmethod
 from dataclasses import asdict, dataclass, field
-from enum import Enum
-from typing import Any, cast
+from enum import Enum, IntEnum
+from typing import Any, assert_never, cast
 
 from redis.exceptions import TimeoutError as RedisTimeoutError
 from sentry_sdk import traces
 
 from snuba import environment, settings
+from snuba.clusters.load_info import LoadInfo
 from snuba.configs.configuration import (
     ConfigurableComponent,
     ConfigurableComponentData,
@@ -26,6 +27,7 @@
 from snuba.web import QueryResult
 
 IS_ENFORCED = "is_enforced"
+IS_PARDONABLE = "is_pardonable"
 MAX_THREADS = "max_threads"
 NO_UNITS = "no_units"
 NO_SUGGESTION = "no_suggestion"
@@ -54,6 +56,13 @@ class AllocationPolicyConfig(Configuration):
     pass
 
 
+class QuotaAllowanceDecision(IntEnum):
+    REJECTED = 0
+    PARDONED = 1
+    ALLOWED = 2
+    THROTTLED = 3
+
+
 @dataclass(frozen=True)
 class QuotaAllowance:
     can_run: bool
@@ -78,6 +87,23 @@ class QuotaAllowance:
     def to_dict(self) -> dict[str, Any]:
         return asdict(self)
 
+    def decision(
+        self,
+        policy: AllocationPolicy,
+        load_info: LoadInfo | None,
+    ) -> QuotaAllowanceDecision:
+        if not self.can_run:
+            should_pardon = getattr(load_info, "should_pardon", lambda: False)
+            idle_pardon = policy.is_pardonable and should_pardon()
+            return (
+                QuotaAllowanceDecision.PARDONED if idle_pardon else QuotaAllowanceDecision.REJECTED
+            )
+
+        if self.max_threads < policy.max_threads:
+            return QuotaAllowanceDecision.THROTTLED
+
+        return QuotaAllowanceDecision.ALLOWED
+
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, QuotaAllowance):
             return False
@@ -221,9 +247,11 @@ class AllocationPolicy(ConfigurableComponent, ABC):
     Any configuration definition that exists in your sub class' `_additional_config_definitions()` will appear in the
     Capacity Management Snuba Admin UI for the policy. From there you can modify the live values to alter how your policy works.
 
-    The base class comes with a built in config accessible as a property of the class itself:
+    The base class comes with built in configs accessible as properties of the class itself:
     - is_enforced
         - Use this to throttle/reject queries OR just log stuff. A configured policy is always active.
+    - is_pardonable
+        - When true, rejections from this policy can be pardoned if the cluster is idle.
 
     Eg.
 
@@ -342,6 +370,12 @@ def __init__(
                 value_type=int,
                 default=kwargs.get(IS_ENFORCED, 1),
             ),
+            AllocationPolicyConfig(
+                name=IS_PARDONABLE,
+                description="Toggles whether rejections from this policy can be pardoned when the cluster is idle.",
+                value_type=int,
+                default=kwargs.get(IS_PARDONABLE, 1),
+            ),
             AllocationPolicyConfig(
                 name=MAX_THREADS,
                 description="The max threads Clickhouse can use for the query.",
@@ -375,6 +409,10 @@ def metrics(self) -> MetricsWrapper:
     def is_enforced(self) -> bool:
         return bool(self.get_config_value(IS_ENFORCED))
 
+    @property
+    def is_pardonable(self) -> bool:
+        return bool(self.get_config_value(IS_PARDONABLE))
+
     @property
     def max_threads(self) -> int:
         """Maximum number of threads run a single query on ClickHouse with."""
@@ -412,14 +450,19 @@ def _get_default_config_definitions(self) -> list[Configuration]:
         return cast(list[Configuration], self._default_config_definitions)
 
     def get_quota_allowance(
-        self, tenant_ids: dict[str, str | int], query_id: str
+        self,
+        tenant_ids: dict[str, str | int],
+        query_id: str,
+        load_info: LoadInfo | None = None,
     ) -> QuotaAllowance:
         with traces.start_span(
             name=self.__class__.__name__,
             attributes={SENTRY_OP: "allocation_policy.get_quota_allowance"},
         ) as span:
             for t, tid in tenant_ids.items():
                 span.set_attribute(f"tenant_ids.{t}", str(tid))
+
+            passthrough = _default_passthough_policy(self._resource_identifier.value)
             try:
                 allowance = self._get_quota_allowance(tenant_ids, query_id)
             except InvalidTenantsForAllocationPolicy as e:
@@ -442,31 +485,49 @@ def get_quota_allowance(
                     1,
                     tags={"method": "get_quota_allowance", "reason": type(e).__name__},
                 )
-                return DEFAULT_PASSTHROUGH_POLICY.get_quota_allowance(tenant_ids, query_id)
+                return passthrough.get_quota_allowance(tenant_ids, query_id, load_info)
             except Exception:
                 self.metrics.increment("fail_open", 1, tags={"method": "get_quota_allowance"})
                 logger.exception(
                     "Allocation policy failed to get quota allowance, this is a bug, fix it"
                 )
                 if settings.RAISE_ON_ALLOCATION_POLICY_FAILURES:
                     raise
-                return DEFAULT_PASSTHROUGH_POLICY.get_quota_allowance(tenant_ids, query_id)
-            if not allowance.can_run:
-                self.metrics.increment(
-                    "db_request_rejected",
-                    tags={"referrer": str(tenant_ids.get("referrer", "no_referrer"))},
-                )
-            elif allowance.max_threads < self.max_threads:
-                # NOTE: The elif is very intentional here. Don't count the throttling
-                # if the request was rejected.
-                self.metrics.increment(
-                    "db_request_throttled",
-                    tags={
-                        "referrer": str(tenant_ids.get("referrer", "no_referrer")),
-                        "max_threads": str(allowance.max_threads),
-                    },
-                )
-                span.set_attribute("db_request_throttled", True)
+                return passthrough.get_quota_allowance(tenant_ids, query_id, load_info)
+
+            decision = allowance.decision(self, load_info)
+            referrer = str(tenant_ids.get("referrer", "no_referrer"))
+            match decision:
+                case QuotaAllowanceDecision.PARDONED:
+                    self.metrics.increment("db_request_pardoned", tags={"referrer": referrer})
+                    allowance = QuotaAllowance(
+                        can_run=True,
+                        max_threads=self.max_threads,
+                        explanation={
+                            **allowance.explanation,
+                            "idle_pardon": load_info.to_dict() if load_info is not None else {},
+                        },
+                        is_throttled=allowance.is_throttled,
+                        throttle_threshold=allowance.throttle_threshold,
+                        rejection_threshold=allowance.rejection_threshold,
+                        quota_used=allowance.quota_used,
+                        quota_unit=allowance.quota_unit,
+                        suggestion=allowance.suggestion,
+                        max_bytes_to_read=0,
+                    )
+                case QuotaAllowanceDecision.REJECTED:
+                    self.metrics.increment("db_request_rejected", tags={"referrer": referrer})
+                case QuotaAllowanceDecision.THROTTLED:
+                    self.metrics.increment(
+                        "db_request_throttled",
+                        tags={"referrer": referrer, "max_threads": str(allowance.max_threads)},
+                    )
+                    span.set_attribute("db_request_throttled", True)
+                case QuotaAllowanceDecision.ALLOWED:
+                    pass
+                case unreachable:
+                    assert_never(unreachable)
+
             if not self.is_enforced:
                 allowance = QuotaAllowance(
                     can_run=True,
@@ -479,6 +540,7 @@ def get_quota_allowance(
                     quota_unit=allowance.quota_unit,
                     suggestion=allowance.suggestion,
                 )
+
             # make sure we always know which storage key we rejected a query from
             allowance.explanation["storage_key"] = self._resource_identifier.value
             for k, v in allowance.to_dict().items():

snuba/web/db_query.py

@@ -26,7 +26,9 @@
 from snuba.clickhouse.query import Query
 from snuba.clickhouse.query_dsl.accessors import get_time_range_estimate
 from snuba.clickhouse.query_profiler import generate_profile
+from snuba.clusters.load_info import get_cluster_loadinfo
 from snuba.configs.configuration import ResourceIdentifier
+from snuba.datasets.storages.factory import get_storage
 from snuba.datasets.storages.storage_key import StorageKey
 from snuba.downsampled_storage_tiers import Tier
 from snuba.query import ProcessableQuery
@@ -878,15 +880,23 @@ def _apply_allocation_policies_quota(
     rejection_quota_and_policy = None
     throttle_quota_and_policy = None
     min_threads_across_policies = MAX_THRESHOLD
+    load_info = None
+    if get_option("storage_routing.enable_get_cluster_loadinfo", False):
+        load_info = get_cluster_loadinfo(
+            get_storage(
+                StorageKey(allocation_policies[0].resource_identifier.value)
+            ).get_storage_set_key()
+        )
     with traces.start_span(
         name="_apply_allocation_policies_quota",
         attributes={SENTRY_OP: "allocation_policy"},
     ) as span:
         for allocation_policy in allocation_policies:
-            allowance = allocation_policy.get_quota_allowance(attribution_info.tenant_ids, query_id)
+            allowance = allocation_policy.get_quota_allowance(
+                attribution_info.tenant_ids, query_id, load_info
+            )
             can_run &= allowance.can_run
             quota_allowances[allocation_policy.class_name()] = allowance
-            # QuotaAllowance isn't a valid attribute value; serialize it.
             span.set_attribute(
                 "quota_allowance",
                 json.dumps(

snuba/web/rpc/storage_routing/routing_strategies/storage_routing.py

@@ -26,6 +26,7 @@
 from sentry_sdk import traces
 
 from snuba import environment, settings
+from snuba.clusters.load_info import LoadInfo, get_cluster_loadinfo
 from snuba.configs.configuration import (
     ConfigurableComponent,
     ConfigurableComponentData,
@@ -54,7 +55,6 @@
 from snuba.web.rpc.common.exceptions import RPCAllocationPolicyException
 from snuba.web.rpc.common.query_info import extract_query_info, extract_query_info_tags
 from snuba.web.rpc.storage_routing.common import extract_message_meta
-from snuba.web.rpc.storage_routing.load_retriever import LoadInfo, get_cluster_loadinfo
 
 _SAMPLING_IN_STORAGE_PREFIX = "sampling_in_storage_"
 _START_ESTIMATION_MARK = "start_sampling_in_storage_estimation"
@@ -462,6 +462,7 @@ def _get_recommendations_from_allocation_policies(
                 recommendations[allocation_policy_name] = allocation_policy.get_quota_allowance(
                     routing_context.tenant_ids,
                     routing_context.query_id,
+                    routing_context.cluster_load_info,
                 )
                 # QuotaAllowance isn't a valid attribute value; serialize it.
                 span.set_attribute(
@@ -482,6 +483,7 @@ def get_routing_decision(self, routing_context: RoutingContext) -> RoutingDecisi
             try:
                 routing_context.timer.mark(_START_ESTIMATION_MARK)
 
+                routing_context.cluster_load_info = get_cluster_loadinfo()
                 routing_context.allocation_policies_recommendations = (
                     self._get_recommendations_from_allocation_policies(routing_context)
                 )
@@ -500,12 +502,6 @@ def get_routing_decision(self, routing_context: RoutingContext) -> RoutingDecisi
                     is_throttled=combined_allocation_policies_recommendations["is_throttled"],
                 )
 
-                routing_context.cluster_load_info = (
-                    get_cluster_loadinfo()
-                    if get_option("storage_routing.enable_get_cluster_loadinfo", False)
-                    else None
-                )
-
                 self._update_routing_decision(routing_decision)
 
                 org_overrides = self._get_org_clickhouse_setting_overrides(