diff --git a/src/sentry/api/endpoints/organization_trace_item_stats.py b/src/sentry/api/endpoints/organization_trace_item_stats.py index 5d9781aefddf..e38d4ec86f9c 100644 --- a/src/sentry/api/endpoints/organization_trace_item_stats.py +++ b/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 diff --git a/src/sentry/options/defaults.py b/src/sentry/options/defaults.py index aa13c482ac0d..fac6f02af086 100644 --- a/src/sentry/options/defaults.py +++ b/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 diff --git a/tests/snuba/api/endpoints/test_organization_trace_item_stats.py b/tests/snuba/api/endpoints/test_organization_trace_item_stats.py index b0ef5a3edf70..e221eb8fd3a6 100644 --- a/tests/snuba/api/endpoints/test_organization_trace_item_stats.py +++ b/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)