perf(explore): Add a request deadline to trace item stats

src/sentry/api/endpoints/organization_trace_item_stats.py

@@ -1,12 +1,14 @@
 import dataclasses
 import logging
+import time
 from collections import defaultdict
 from collections.abc import Callable, Sequence
-from concurrent.futures import as_completed
+from concurrent.futures import FIRST_EXCEPTION, wait
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from typing import Any, Literal, cast, get_args
+from typing import Any, Literal, NoReturn, cast, get_args
 
+import sentry_sdk
 from drf_spectacular.utils import extend_schema
 from rest_framework import serializers
 from rest_framework.request import Request
@@ -26,7 +28,7 @@
     translate_escape_sequences,
 )
 from sentry.api.serializers.base import serialize
-from sentry.api.utils import handle_query_errors
+from sentry.api.utils import TimeoutException, handle_query_errors
 from sentry.apidocs.constants import (
     RESPONSE_BAD_REQUEST,
     RESPONSE_FORBIDDEN,
@@ -57,11 +59,13 @@
 from sentry.search.eap.spans.definitions import SPAN_DEFINITIONS
 from sentry.search.eap.types import EAPResponse, SearchResolverConfig, SupportedTraceItemType
 from sentry.search.eap.utils import can_expose_attribute_to_api
+from sentry.search.events.constants import TIMEOUT_RPC_ERROR_MESSAGE
 from sentry.search.events.types import SnubaParams
 from sentry.snuba import rpc_dataset_common
 from sentry.snuba.occurrences_rpc import Occurrences
 from sentry.snuba.referrer import Referrer
 from sentry.snuba.spans_rpc import Spans
+from sentry.utils import metrics
 from sentry.utils.concurrent import ContextPropagatingThreadPoolExecutor
 from sentry.utils.cursors import Cursor, CursorResult
 
@@ -107,6 +111,37 @@ def get_trace_item_stats_config(item_type: SupportedItemType) -> TraceItemStatsC
     )
 
 
+class StatsRequestDeadline:
+    """
+    Overall time budget for one stats request. Each Snuba RPC has its own timeout,
+    but the request chains several of them, so without a shared budget a slow
+    request can run for several RPC timeouts back to back.
+    """
+
+    def __init__(self, seconds: float) -> None:
+        self._deadline = time.monotonic() + seconds if seconds > 0 else None
+
+    def remaining(self, stage: str) -> float | None:
+        """
+        Seconds left in the budget, or None when the budget is disabled. Raises a
+        timeout instead of returning once the budget is spent, so `stage` never starts.
+        """
+        if self._deadline is None:
+            return None
+        remaining = self._deadline - time.monotonic()
+        if remaining <= 0:
+            raise_deadline_exceeded(stage)
+        return remaining
+
+
+def raise_deadline_exceeded(stage: str) -> NoReturn:
+    # Match what handle_query_errors returns for a Snuba RPC timeout.
+    sentry_sdk.set_tag("query.error_reason", "Timeout")
+    sentry_sdk.set_attribute("query.error_reason", "Timeout")
+    metrics.incr("api.trace_item_stats.deadline_exceeded", tags={"stage": stage})
+    raise TimeoutException(detail=TIMEOUT_RPC_ERROR_MESSAGE)
+
+
 def _find_aggregate_filters(terms: Sequence[QueryToken]) -> list[AggregateFilter]:
     found: list[AggregateFilter] = []
     for term in terms:
@@ -260,6 +295,7 @@ def get(
         value_substring_match = translate_escape_sequences(substring_match)
 
         max_attributes = options.get("explore.trace-items.keys.max")
+        deadline = StatsRequestDeadline(options.get("explore.trace-items.stats.deadline-seconds"))
         include_internal = is_active_superuser(request) or is_active_staff(request)
         include_internal_convention_attributes = request.user.is_staff or request.user.is_superuser
 
@@ -332,6 +368,7 @@ def data_fn(offset: int, limit: int) -> tuple[dict[str, Any], int]:
             item_ids = [row[stats_config.id_column] for row in table_results["data"]]
             item_id_list = ",".join(item_ids)
 
+            deadline.remaining("preflight")
             preflight_stats = run_stats_query_with_item_ids(
                 f"id:[{item_id_list}]", get_preflight_params(table_results["data"])
             )
@@ -392,21 +429,28 @@ def data_fn(offset: int, limit: int) -> tuple[dict[str, Any], int]:
                     AttributeKey(name=attr.name, type=AttributeKey.TYPE_STRING)
                 )
 
+            stats_timeout = deadline.remaining("stats")
             stats_results: dict[str, dict[str, dict[str, Any]]] = defaultdict(lambda: {"data": {}})
-            with ContextPropagatingThreadPoolExecutor(
+            query_thread_pool = ContextPropagatingThreadPoolExecutor(
                 thread_name_prefix=__name__,
                 max_workers=MAX_THREADS,
-            ) as query_thread_pool:
+            )
+            try:
                 futures = [
                     query_thread_pool.submit(run_stats_query_with_error_handling, attributes)
                     for attributes in chunked_attributes.values()
                 ]
-
-                for future in as_completed(futures):
+                done, not_done = wait(futures, timeout=stats_timeout, return_when=FIRST_EXCEPTION)
+                for future in done:
                     result = future.result()
                     for stats in result:
                         for stats_type, data in stats.items():
                             stats_results[stats_type]["data"].update(data["data"])
+                if not_done:
+                    raise_deadline_exceeded("stats")
+            finally:
+                # If the request fails or runs out of budget, don't wait on RPCs still in flight.
+                query_thread_pool.shutdown(wait=False, cancel_futures=True)
 
             return {"data": [{k: v} for k, v in stats_results.items()]}, total_attributes
 

src/sentry/options/defaults.py

@@ -2732,6 +2732,15 @@
     default=1000,
     flags=FLAG_AUTOMATOR_MODIFIABLE,
 )
+# Overall time budget, in seconds, for one trace item stats request across all of
+# its Snuba RPCs. 0 disables the budget. Defaults to SENTRY_SNUBA_TIMEOUT so it only
+# cuts off requests that chain several slow RPCs.
+register(
+    "explore.trace-items.stats.deadline-seconds",
+    type=Float,
+    default=30.0,
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
 
 
 # Used for the z-score when calculating the margin of error in performance

tests/snuba/api/endpoints/test_organization_trace_item_stats.py

@@ -1,3 +1,4 @@
+import threading
 from datetime import timedelta
 from typing import Any
 from unittest import mock
@@ -483,6 +484,78 @@ def test_duplicate_tags(self) -> None:
         assert "user.geo.subregion" not in attribute_distribution
         assert "user.geo.city" not in attribute_distribution
 
+    def _mock_clock(self, *readings: float) -> None:
+        patcher = mock.patch("sentry.api.endpoints.organization_trace_item_stats.time")
+        mock_time = patcher.start()
+        self.addCleanup(patcher.stop)
+        mock_time.monotonic.side_effect = list(readings)
+
+    @override_options({"explore.trace-items.stats.deadline-seconds": 15.0})
+    def test_deadline_exceeded_before_preflight(self) -> None:
+        self._store_span(sentry_tags={"browser": "chrome"})
+        # Request start, then the check before the preflight query.
+        self._mock_clock(0.0, 15.0)
+
+        with mock.patch.object(Spans, "run_stats_query") as mock_stats:
+            response = self.do_request(query={"statsType": ["attributeDistributions"]})
+
+        assert response.status_code == 504, response.data
+        assert "Query timeout" in response.data["detail"]
+        mock_stats.assert_not_called()
+
+    @override_options({"explore.trace-items.stats.deadline-seconds": 15.0})
+    def test_deadline_exceeded_before_stats(self) -> None:
+        self._store_span(sentry_tags={"browser": "chrome"})
+        # Request start, the preflight check, then the check before the stats queries.
+        self._mock_clock(0.0, 1.0, 16.0)
+        real_run_stats_query = Spans.run_stats_query
+
+        with mock.patch.object(
+            Spans, "run_stats_query", side_effect=real_run_stats_query
+        ) as mock_stats:
+            response = self.do_request(query={"statsType": ["attributeDistributions"]})
+
+        assert response.status_code == 504, response.data
+        # Only the preflight query ran; the per-attribute stats queries never started.
+        assert mock_stats.call_count == 1
+        assert "attributes" not in mock_stats.call_args.kwargs
+
+    @override_options({"explore.trace-items.stats.deadline-seconds": 15.0})
+    def test_deadline_exceeded_during_stats(self) -> None:
+        self._store_span(sentry_tags={"browser": "chrome"})
+        # Leave 10ms of budget for the stats queries, which then never finish in time.
+        self._mock_clock(0.0, 1.0, 14.99)
+        real_run_stats_query = Spans.run_stats_query
+        release = threading.Event()
+        finished = threading.Event()
+        self.addCleanup(release.set)
+
+        def run_stats_query(*args: Any, **kwargs: Any) -> Any:
+            if "attributes" in kwargs:
+                release.wait(timeout=10)
+                finished.set()
+                return []
+            return real_run_stats_query(*args, **kwargs)
+
+        with mock.patch.object(Spans, "run_stats_query", side_effect=run_stats_query):
+            response = self.do_request(query={"statsType": ["attributeDistributions"]})
+
+        assert response.status_code == 504, response.data
+        assert "Query timeout" in response.data["detail"]
+        # The response did not wait for the stats queries still in flight.
+        assert not finished.is_set()
+
+    @override_options({"explore.trace-items.stats.deadline-seconds": 15.0})
+    def test_within_deadline_returns_stats(self) -> None:
+        self._store_span(sentry_tags={"device": "mobile"})
+        self._mock_clock(0.0, 5.0, 10.0)
+
+        response = self.do_request(query={"statsType": ["attributeDistributions"]})
+
+        assert response.status_code == 200, response.data
+        attribute_distribution = response.data["data"][0]["attributeDistributions"]["data"]
+        assert {"label": "mobile", "value": 1.0} in attribute_distribution["sentry.device"]
+
     def test_preflight_narrows_time_range_to_sampled_items(self) -> None:
         oldest = before_now(hours=3).replace(microsecond=750_000)
         newest = before_now(minutes=10).replace(microsecond=750_000)