[CapMan visibility] rejected queries are ran with 0 threads (#6545)

Snuba payload should correctly reflect that rejected queries are ran
with 0 threads

---------

Co-authored-by: Rachel Chen <rachelchen@PL6VFX9HP4.attlocal.net>

snuba/query/allocation_policies/bytes_scanned_rejecting_policy.py

@@ -247,7 +247,7 @@ def _get_quota_allowance(
             )
             return QuotaAllowance(
                 can_run=False,
-                max_threads=self.max_threads,
+                max_threads=0,
                 explanation=explanation,
                 is_throttled=True,
                 throttle_threshold=throttle_threshold,

snuba/query/allocation_policies/concurrent_rate_limit.py

@@ -260,20 +260,16 @@ def _get_quota_allowance(
             rate_limit_params,
         )
 
-        if within_rate_limit:
-            suggestion = NO_SUGGESTION
-        else:
-            suggestion = SUGGESTION
         return QuotaAllowance(
             can_run=within_rate_limit,
-            max_threads=self.max_threads,
+            max_threads=self.max_threads if within_rate_limit else 0,
             explanation={"reason": why, "overrides": overrides},
             is_throttled=False,
             throttle_threshold=typing.cast(int, rate_limit_params.concurrent_limit),
             rejection_threshold=typing.cast(int, rate_limit_params.concurrent_limit),
             quota_used=rate_limit_stats.concurrent,
             quota_unit=QUOTA_UNIT,
-            suggestion=suggestion,
+            suggestion=NO_SUGGESTION if within_rate_limit else SUGGESTION,
         )
 
     def _update_quota_balance(

snuba/query/allocation_policies/cross_org.py

@@ -194,20 +194,16 @@ def _get_quota_allowance(
                 "cross_org_query"
             ] = f"This referrer is not registered for the current storage {self._storage_key.value}, if you want to increase its limits, register it in the yaml of the CrossOrgQueryAllocationPolicy"
 
-        if can_run:
-            suggestion = NO_SUGGESTION
-        else:
-            suggestion = SUGGESTION
         return QuotaAllowance(
             can_run=can_run,
-            max_threads=self._get_max_threads(referrer),
+            max_threads=self._get_max_threads(referrer) if can_run else 0,
             explanation=decision_explanation,
             is_throttled=False,
             throttle_threshold=typing.cast(int, rate_limit_params.concurrent_limit),
             rejection_threshold=typing.cast(int, rate_limit_params.concurrent_limit),
             quota_used=rate_limit_stats.concurrent,
             quota_unit=QUOTA_UNIT,
-            suggestion=suggestion,
+            suggestion=NO_SUGGESTION if can_run else SUGGESTION,
         )
 
     def _update_quota_balance(

snuba/query/allocation_policies/per_referrer.py

@@ -156,20 +156,16 @@ def _get_quota_allowance(
             "referrer": referrer,
         }
 
-        if can_run:
-            suggestion = NO_SUGGESTION
-        else:
-            suggestion = SUGGESTION
         return QuotaAllowance(
             can_run=can_run,
-            max_threads=num_threads,
+            max_threads=num_threads if can_run else 0,
             explanation=decision_explanation,
             is_throttled=is_throttled,
             throttle_threshold=requests_throttle_threshold,
             rejection_threshold=rate_limit_params.concurrent_limit,
             quota_used=rate_limit_stats.concurrent,
             quota_unit=QUOTA_UNIT,
-            suggestion=suggestion,
+            suggestion=NO_SUGGESTION if can_run else SUGGESTION,
         )
 
     def _update_quota_balance(

tests/query/allocation_policies/test_allocation_policy_base.py

@@ -65,7 +65,7 @@ def _get_quota_allowance(
     ) -> QuotaAllowance:
         return QuotaAllowance(
             can_run=False,
-            max_threads=10,
+            max_threads=0,
             explanation={},
             is_throttled=True,
             throttle_threshold=MAX_THRESHOLD,
@@ -503,7 +503,8 @@ def test_is_not_enforced() -> None:
         "organization_id": 123,
         "referrer": "some_referrer",
     }
-    assert not reject_policy.get_quota_allowance(tenant_ids, "deadbeef").can_run
+    quota_allowance = reject_policy.get_quota_allowance(tenant_ids, "deadbeef")
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
 
     reject_policy.set_config_value(config_key="is_enforced", value=0)
     # policy not enforced so we don't reject the query

tests/query/allocation_policies/test_bytes_scanned_rejecting_policy.py

@@ -89,7 +89,7 @@ def test_consume_quota(
         ),
     )
     allowance = policy.get_quota_allowance(tenant_ids, QUERY_ID)
-    assert not allowance.can_run
+    assert not allowance.can_run and allowance.max_threads == 0
     assert {
         "granted_quota": 0,
         "limit": limit,
@@ -149,7 +149,7 @@ def test_invalid_tenants(
     policy: BytesScannedRejectingPolicy, tenant_ids: dict[str, str | int]
 ) -> None:
     allowance = policy.get_quota_allowance(tenant_ids, QUERY_ID)
-    assert not allowance.can_run
+    assert not allowance.can_run and allowance.max_threads == 0
     assert allowance.explanation["__name__"] == "InvalidTenantsForAllocationPolicy"
 
 
@@ -206,7 +206,7 @@ def test_overrides(
         ),
     )
     allowance = policy.get_quota_allowance(tenant_ids, QUERY_ID)
-    assert not allowance.can_run
+    assert not allowance.can_run and allowance.max_threads == 0
     assert allowance.explanation["limit"] == limit
 
 
@@ -243,7 +243,7 @@ def test_penalize_timeout(policy: BytesScannedRejectingPolicy) -> None:
     )
 
     allowance = policy.get_quota_allowance(tenant_ids, QUERY_ID)
-    assert not allowance.can_run
+    assert not allowance.can_run and allowance.max_threads == 0
 
 
 @pytest.mark.redis_db

tests/query/allocation_policies/test_bytes_scanned_window_allocation_policy.py

@@ -148,12 +148,15 @@ def test_enforcement_switch(policy: AllocationPolicy) -> None:
 @pytest.mark.redis_db
 def test_reject_queries_without_tenant_ids(policy: AllocationPolicy) -> None:
     _configure_policy(policy)
-    assert not policy.get_quota_allowance(
+    quota_allowance = policy.get_quota_allowance(
         tenant_ids={"organization_id": 1234}, query_id=QUERY_ID
-    ).can_run
-    assert not policy.get_quota_allowance(
+    )
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
+
+    quota_allowance = policy.get_quota_allowance(
         tenant_ids={"referrer": "bloop"}, query_id=QUERY_ID
-    ).can_run
+    )
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
     # These should not fail because we know they don't have an org id
     for referrer in _ORG_LESS_REFERRERS:
         tenant_ids: dict[str, str | int] = {"referrer": referrer}

tests/query/allocation_policies/test_concurrent_rate_limit_policy.py

@@ -53,9 +53,10 @@ def test_rate_limit_concurrent(policy: ConcurrentRateLimitAllocationPolicy) -> N
             tenant_ids={"organization_id": 123}, query_id=f"abc{i}"
         )
 
-    assert not policy.get_quota_allowance(
+    quota_allowance = policy.get_quota_allowance(
         tenant_ids={"organization_id": 123}, query_id=f"abc{MAX_CONCURRENT_QUERIES}"
-    ).can_run
+    )
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
 
 
 @pytest.mark.redis_db
@@ -69,10 +70,11 @@ def test_rate_limit_concurrent_diff_tenants(
             tenant_ids={"organization_id": RATE_LIMITED_ORG_ID}, query_id=f"abc{i}"
         )
 
-    assert not policy.get_quota_allowance(
+    quota_allowance = policy.get_quota_allowance(
         tenant_ids={"organization_id": RATE_LIMITED_ORG_ID},
         query_id=f"abc{MAX_CONCURRENT_QUERIES}",
-    ).can_run
+    )
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
     policy.get_quota_allowance(
         tenant_ids={"organization_id": OTHER_ORG_ID},
         query_id=f"abc{MAX_CONCURRENT_QUERIES}",
@@ -106,9 +108,10 @@ def test_rate_limit_concurrent_complete_query(
         )
 
     # cant submit anymore
-    assert not policy.get_quota_allowance(
+    quota_allowance = policy.get_quota_allowance(
         tenant_ids={"organization_id": 123}, query_id=f"abc{MAX_CONCURRENT_QUERIES}"
-    ).can_run
+    )
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
 
     # one query finishes
     policy.update_quota_balance(
@@ -123,10 +126,11 @@ def test_rate_limit_concurrent_complete_query(
     )
 
     # but no more than that
-    assert not policy.get_quota_allowance(
+    quota_allowance = policy.get_quota_allowance(
         tenant_ids={"organization_id": 123},
         query_id="some_query_id",
-    ).can_run
+    )
+    assert not quota_allowance.can_run and quota_allowance.max_threads == 0
 
 
 @pytest.mark.redis_db
@@ -260,7 +264,7 @@ def test_apply_overrides(
     allowance = policy.get_quota_allowance(
         tenant_ids=tenant_ids, query_id=f"{expected_concurrent_limit+1}"
     )
-    assert not allowance.can_run
+    assert not allowance.can_run and allowance.max_threads == 0
     assert allowance.explanation["overrides"] == expected_overrides
 
 
@@ -346,7 +350,8 @@ def test_cross_org(policy: ConcurrentRateLimitAllocationPolicy) -> None:
 def test_bad_tenants(policy: ConcurrentRateLimitAllocationPolicy):
     bad_tenant_ids: dict[str, str | int] = {"referrer": "abcd"}
     with mock.patch("snuba.settings.RAISE_ON_ALLOCATION_POLICY_FAILURES", False):
-        assert not policy.get_quota_allowance(bad_tenant_ids, "1234").can_run
+        quota_allowance = policy.get_quota_allowance(bad_tenant_ids, "1234")
+        assert not quota_allowance.can_run and quota_allowance.max_threads == 0
 
         # does not raise
         policy.update_quota_balance(bad_tenant_ids, "1234", _RESULT_SUCCESS)

tests/query/allocation_policies/test_cross_org_policy.py

@@ -43,9 +43,10 @@ def test_policy_pass_basic(self):
         assert cross_org_allowance.can_run is True
         assert cross_org_allowance.max_threads == 1
 
-        assert not policy.get_quota_allowance(
+        quota_allowance = policy.get_quota_allowance(
             tenant_ids={"referrer": "statistical_detectors"}, query_id="3"
-        ).can_run
+        )
+        assert not quota_allowance.can_run and quota_allowance.max_threads == 0
         policy.update_quota_balance(
             tenant_ids={"referrer": "statistical_detectors"},
             query_id="2",
@@ -124,9 +125,11 @@ def test_override(self):
             0,
             {"referrer": "statistical_detectors"},
         )
-        assert not policy.get_quota_allowance(
+
+        quota_allowance = policy.get_quota_allowance(
             tenant_ids={"referrer": "statistical_detectors"}, query_id="2"
-        ).can_run
+        )
+        assert not quota_allowance.can_run and quota_allowance.max_threads == 0
 
     @pytest.mark.redis_db
     def test_override_unregistered_referrer(self):

tests/query/allocation_policies/test_per_referrer.py

@@ -32,9 +32,10 @@ def test_policy_pass_basic(self):
             tenant_ids={"referrer": "statistical_detectors"}, query_id="2"
         )
 
-        assert not policy.get_quota_allowance(
+        quota_allowance = policy.get_quota_allowance(
             tenant_ids={"referrer": "statistical_detectors"}, query_id="3"
-        ).can_run
+        )
+        assert not quota_allowance.can_run and quota_allowance.max_threads == 0
         # clean up the failed request
         policy.update_quota_balance(
             tenant_ids={"referrer": "statistical_detectors"},
@@ -108,6 +109,8 @@ def test_override(self):
             0,
             {"referrer": "statistical_detectors"},
         )
-        assert not policy.get_quota_allowance(
+
+        quota_allowance = policy.get_quota_allowance(
             tenant_ids={"referrer": "statistical_detectors"}, query_id="2"
-        ).can_run
+        )
+        assert not quota_allowance.can_run and quota_allowance.max_threads == 0