squash

Author: Rachel Chen

snuba/datasets/configuration/events_analytics_platform/storages/eap_items.yaml

@@ -109,31 +109,6 @@ schema:
   local_table_name: eap_items_1_local
   dist_table_name: eap_items_1_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor

snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_512.yaml

@@ -108,31 +108,6 @@ schema:
   local_table_name: eap_items_1_downsample_512_local
   dist_table_name: eap_items_1_downsample_512_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor

snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_64.yaml

@@ -109,31 +109,6 @@ schema:
   local_table_name: eap_items_1_downsample_64_local
   dist_table_name: eap_items_1_downsample_64_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor

snuba/datasets/configuration/events_analytics_platform/storages/eap_items_downsample_8.yaml

@@ -108,31 +108,6 @@ schema:
   local_table_name: eap_items_1_downsample_8_local
   dist_table_name: eap_items_1_downsample_8_dist
   partition_format: [date]
-allocation_policies:
-  - name: ConcurrentRateLimitAllocationPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - referrer
-        - project_id
-      default_config_overrides:
-        is_enforced: 0
-  - name: ReferrerGuardRailPolicy
-    args:
-      required_tenant_types:
-        - referrer
-      default_config_overrides:
-        is_enforced: 0
-        is_active: 0
-  - name: BytesScannedRejectingPolicy
-    args:
-      required_tenant_types:
-        - organization_id
-        - project_id
-        - referrer
-      default_config_overrides:
-        is_active: 0
-        is_enforced: 0
 
 query_processors:
   - processor: UniqInSelectAndHavingProcessor

snuba/query/allocation_policies/utils.py

@@ -0,0 +1,11 @@
+from snuba.query.allocation_policies import QuotaAllowance
+
+
+def get_max_bytes_to_read(quota_allowances: dict[str, QuotaAllowance]) -> int:
+    max_bytes_to_read = min(
+        [qa.max_bytes_to_read for qa in quota_allowances.values()],
+        key=lambda mb: float("inf") if mb == 0 else mb,
+    )
+    if max_bytes_to_read != 0:
+        return max_bytes_to_read
+    return 0

snuba/web/db_query.py

@@ -33,6 +33,7 @@
     QueryResultOrError,
     QuotaAllowance,
 )
+from snuba.query.allocation_policies.utils import get_max_bytes_to_read
 from snuba.query.composite import CompositeQuery
 from snuba.query.data_source.join import IndividualNode, JoinClause, JoinVisitor
 from snuba.query.data_source.simple import Table
@@ -808,7 +809,11 @@ def _apply_allocation_policies_quota(
     Sets the resource quota in the query_settings object to the minimum of all available
     quota allowances from the given allocation policies.
     """
-    quota_allowances: dict[str, Any] = {}
+
+    if len(allocation_policies) == 0:
+        return
+
+    quota_allowances: dict[str, QuotaAllowance] = {}
     can_run = True
     rejection_quota_and_policy = None
     throttle_quota_and_policy = None
@@ -855,10 +860,7 @@ def _apply_allocation_policies_quota(
         summary: dict[str, Any] = {}
         summary["threads_used"] = min_threads_across_policies
 
-        max_bytes_to_read = min(
-            [qa.max_bytes_to_read for qa in quota_allowances.values()],
-            key=lambda mb: float("inf") if mb == 0 else mb,
-        )
+        max_bytes_to_read = get_max_bytes_to_read(quota_allowances)
         if max_bytes_to_read != 0:
             query_settings.push_clickhouse_setting("max_bytes_to_read", max_bytes_to_read)
             summary["max_bytes_to_read"] = max_bytes_to_read

snuba/web/rpc/__init__.py

@@ -31,10 +31,11 @@
     RPCRequestException,
     convert_rpc_exception_to_proto,
 )
-from snuba.web.rpc.storage_routing.load_retriever import get_cluster_loadinfo
+from snuba.web.rpc.storage_routing.common import extract_message_meta
 from snuba.web.rpc.storage_routing.routing_strategies.storage_routing import (
     RoutingContext,
     RoutingDecision,
+    get_stats_dict,
 )
 from snuba.web.rpc.storage_routing.routing_strategy_selector import (
     RoutingStrategySelector,
@@ -191,7 +192,21 @@ def execute(self, in_msg: Tin) -> Tout:
         span = scope.span
         if span is not None:
             span.description = self.config_key()
-        self.routing_context = RoutingContext(timer=self._timer, in_msg=in_msg)
+        request_meta = extract_message_meta(in_msg)
+        self.routing_context = RoutingContext(
+            timer=self._timer,
+            in_msg=in_msg,
+            query_id=uuid.uuid4().hex,
+            tenant_ids={
+                "organization_id": request_meta.organization_id,
+                "referrer": request_meta.referrer,
+                **(
+                    {"project_id": request_meta.project_ids[0]}
+                    if hasattr(request_meta, "project_ids") and len(request_meta.project_ids) == 1
+                    else {}
+                ),
+            },
+        )
 
         self.__before_execute(in_msg)
         error: Exception | None = None
@@ -202,7 +217,29 @@ def execute(self, in_msg: Tin) -> Tout:
                     span.set_data("selected_tier", self.routing_decision.tier)
                     out = self._execute(in_msg)
             else:
-                raise AllocationPolicyViolations
+                self.metrics.increment(
+                    "request_rate_limited",
+                    tags=self._timer.tags,
+                )
+                error = QueryException.from_args(
+                    AllocationPolicyViolations.__name__,
+                    "Query cannot be run due to routing strategy deciding it cannot run, most likely due to allocation policies",
+                    extra={
+                        "stats": get_stats_dict(self.routing_decision),
+                        "sql": "no sql run",
+                        "experiments": {},
+                    },
+                )
+                error.__cause__ = AllocationPolicyViolations.from_args(
+                    {
+                        "summary": {},
+                        "details": {
+                            key: quota_allowance.to_dict()
+                            for key, quota_allowance in self.routing_decision.routing_context.allocation_policies_recommendations.items()
+                        },
+                    }
+                )
+                raise error
         except QueryException as e:
             out = self.response_class()()
             if (
@@ -253,16 +290,10 @@ def __before_execute(self, in_msg: Tin) -> None:
         meta = getattr(in_msg, "meta", None)
         if meta is not None:
             if not hasattr(meta, "request_id") or not meta.request_id:
-                meta.request_id = str(uuid.uuid4())
+                meta.request_id = self.routing_context.query_id
 
         self._timer.update_tags(self.__extract_request_tags(in_msg))
 
-        # we're calling this function to get the cluster load info to emit metrics and to prevent dead code
-        # the result is currently not used in storage routing
-        # can turn off on Snuba Admin
-        if state.get_config("storage_routing.enable_get_cluster_loadinfo", True):
-            get_cluster_loadinfo()
-
         selected_strategy = RoutingStrategySelector().select_routing_strategy(self.routing_context)
         self.routing_decision = selected_strategy.get_routing_decision(self.routing_context)
         self._timer.mark("rpc_start")
@@ -322,11 +353,11 @@ def _execute(self, in_msg: Tin) -> Tout:
 
     def __after_execute(self, in_msg: Tin, out_msg: Tout, error: Exception | None) -> Tout:
         res = self._after_execute(in_msg, out_msg, error)
-        output_metrics_error = None
+        routing_strategy_after_execute_error = None
         try:
-            self.routing_decision.strategy.output_metrics(self.routing_decision)
+            self.routing_decision.strategy.after_execute(self.routing_decision, error)
         except Exception as e:
-            output_metrics_error = e
+            routing_strategy_after_execute_error = e
 
         self._timer.mark("rpc_end")
         self._timer.send_metrics_to(self.metrics)
@@ -336,13 +367,11 @@ def __after_execute(self, in_msg: Tin, out_msg: Tout, error: Exception | None) -
                     "request_invalid",
                     tags=self._timer.tags,
                 )
-            elif isinstance(error, AllocationPolicyViolations):
-                sentry_sdk.capture_exception(error)
-                self.metrics.increment(
-                    "request_rate_limited",
-                    tags=self._timer.tags,
-                )
-            else:
+            # AllocationPolicyViolations is not a request_error
+            elif not (
+                isinstance(error, QueryException)
+                and error.exception_type == AllocationPolicyViolations.__name__
+            ):
                 sentry_sdk.capture_exception(error)
                 self.metrics.increment(
                     "request_error",
@@ -355,13 +384,13 @@ def __after_execute(self, in_msg: Tin, out_msg: Tout, error: Exception | None) -
                 tags=self._timer.tags,
             )
 
-        if output_metrics_error is not None:
-            self.metrics.increment("metrics_failure")
+        if routing_strategy_after_execute_error is not None:
+            self.metrics.increment("after_execute_failure")
             sentry_sdk.capture_message(
-                f"Error in routing strategy output metrics: {output_metrics_error}"
+                f"Error in routing strategy after execute: {routing_strategy_after_execute_error}"
             )
             if settings.RAISE_ON_ROUTING_STRATEGY_FAILURES:
-                raise output_metrics_error
+                raise routing_strategy_after_execute_error
 
         return res
 

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

@@ -149,14 +149,26 @@ def _get_min_timerange_to_query_outcomes(self) -> int:
             or default
         )
 
-    def _get_routing_decision(self, routing_context: RoutingContext) -> RoutingDecision:
+    def _get_routing_decision(
+        self,
+        routing_context: RoutingContext,
+    ) -> RoutingDecision:
+        combined_allocation_policies_recommendations = (
+            self._get_combined_allocation_policies_recommendations(routing_context)
+        )
+
         routing_decision = RoutingDecision(
             routing_context=routing_context,
             strategy=self,
             tier=Tier.TIER_1,
-            clickhouse_settings={},
-            can_run=True,
+            clickhouse_settings=combined_allocation_policies_recommendations["settings"],
+            can_run=combined_allocation_policies_recommendations["can_run"],
+            is_throttled=combined_allocation_policies_recommendations["is_throttled"],
         )
+
+        if not routing_decision.can_run:
+            return routing_decision
+
         in_msg_meta = extract_message_meta(routing_decision.routing_context.in_msg)
         sentry_sdk.update_current_span(
             attributes={