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/web/db_query.py

@@ -808,6 +808,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.
     """
+
+    # if it's an EAP query, then allocation policies were already applied in the routing strategy layer, so we bypass it here
+    if attribution_info.app_id.key == "eap":
+        return
+
     quota_allowances: dict[str, Any] = {}
     can_run = True
     rejection_quota_and_policy = None

snuba/web/rpc/__init__.py

@@ -35,6 +35,7 @@
 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,7 @@ 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)
+        self.routing_context = RoutingContext(timer=self._timer, in_msg=in_msg, query_id=uuid.uuid4().hex)
 
         self.__before_execute(in_msg)
         error: Exception | None = None
@@ -202,7 +203,20 @@ 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,
+                )
+                # thoughts on making AlocationPolicyViolations a subclass of QueryException?
+                raise 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": {},
+                    },
+                )
         except QueryException as e:
             out = self.response_class()()
             if (
@@ -257,12 +271,6 @@ def __before_execute(self, in_msg: Tin) -> None:
 
         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")
@@ -325,6 +333,10 @@ def __after_execute(self, in_msg: Tin, out_msg: Tout, error: Exception | None) -
         output_metrics_error = None
         try:
             self.routing_decision.strategy.output_metrics(self.routing_decision)
+            if self.routing_decision.routing_context.query_result is not None or isinstance(
+                error, QueryException
+            ):
+                self.routing_decision.strategy.update_allocation_policies_balances(self.routing_decision, error)  # type: ignore
         except Exception as e:
             output_metrics_error = e
 
@@ -336,12 +348,6 @@ 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:
                 sentry_sdk.capture_exception(error)
                 self.metrics.increment(

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

@@ -30,6 +30,7 @@
 from snuba.web.rpc.storage_routing.common import extract_message_meta
 from snuba.web.rpc.storage_routing.routing_strategies.storage_routing import (
     BaseRoutingStrategy,
+    CombinedAllocationPoliciesRecommendations,
     RoutingContext,
     RoutingDecision,
 )
@@ -165,14 +166,55 @@ def _get_min_timerange_to_query_outcomes(self) -> int:
             or default
         )
 
-    def _get_routing_decision(self, routing_context: RoutingContext) -> RoutingDecision:
+    def _get_combined_allocation_policies_recommendations(
+        self, routing_context: RoutingContext
+    ) -> CombinedAllocationPoliciesRecommendations:
+        settings = {}
+        max_bytes_to_read = min(
+            [
+                qa.max_bytes_to_read
+                for qa in routing_context.allocation_policies_recommendations.values()
+            ],
+            key=lambda mb: float("inf") if mb == 0 else mb,
+        )
+        if max_bytes_to_read != 0:
+            settings["max_bytes_to_read"] = max_bytes_to_read
+
+        settings["max_threads"] = min(
+            [qa.max_threads for qa in routing_context.allocation_policies_recommendations.values()],
+        )
+
+        return CombinedAllocationPoliciesRecommendations(
+            can_run=all(
+                qa.can_run for qa in routing_context.allocation_policies_recommendations.values()
+            ),
+            is_throttled=any(
+                qa.is_throttled
+                for qa in routing_context.allocation_policies_recommendations.values()
+            ),
+            settings=settings,
+        )
+
+    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={