From 2d667a41233f9873cfb529aae9224205324c5f27 Mon Sep 17 00:00:00 2001 From: zz_y Date: Tue, 26 May 2026 09:47:00 -0600 Subject: [PATCH] refactor(query-routing): centralize ASAP-first routing in EngineRouter; add range support Routing decisions now live in one place (EngineRouter + compatible_storage_backends); http.rs is a thin transport layer. Both instant and range queries try the ASAP sketch tier first and fall back to the archive (thanos_query) on CapabilityMiss, except accuracy==Exact goes straight to the archive. Supersedes PR #309's inline approach. - QueryEngine: add `execute_range` (default returns CapabilityMiss). - EngineRouter: add `execute_range` mirroring `execute()`'s failover over the shared `compatible_storage_backends` sequence. - ASAPQueryEngine: `execute_range` delegates to `execute_range_promql_modern` so the router reaches the real warm-tier range path. - ThanosQueryEngine: implement `execute_range` forwarding to `/api/v1/query_range` (4xx -> CapabilityMiss, 5xx/timeout -> Backend), ported from #309. - compatible_storage_backends: ASAP-first (Approximate -> [SketchStore, GorillaObjectStore]); Exact -> archive only; PrometheusRemote unchanged. Tests updated. - http.rs: `process_range_query_request` thinned to a single `query_router.execute_range` call; inline Thanos lookup removed. accuracy is still hardcoded `Approximate` on the live path (matches the instant path); the `Exact -> archive` gate is wired but dormant until a request can declare accuracy. cargo check + routing/thanos/http/capability_matching tests green. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/asap_types/src/capability_matching.rs | 189 +++++------ data_plane/src/drivers/query/servers/http.rs | 193 +++++++++-- .../query_engines/asap_query_engine/engine.rs | 20 ++ .../routing/query_engine_routing.rs | 256 +++++++++++++- .../thanos_query_engine/forward.rs | 311 ++++++++++++++++++ 5 files changed, 820 insertions(+), 149 deletions(-) diff --git a/crates/asap_types/src/capability_matching.rs b/crates/asap_types/src/capability_matching.rs index ce853a6f..79855f5f 100644 --- a/crates/asap_types/src/capability_matching.rs +++ b/crates/asap_types/src/capability_matching.rs @@ -248,59 +248,55 @@ pub fn compatible_agg_types(stat: Statistic) -> &'static [AggregationType] { /// query when the metric is configured for `metric_storage_config`. /// /// The returned list is **ordered by preference**: the router walks it in -/// order and dispatches to the first backend whose engine is registered. +/// order and dispatches to the first backend whose engine is registered, +/// falling through on `CapabilityMiss` / `Backend` to the next entry. /// -/// **Step-1 of the JSONL deprecation refactor**: the legacy -/// `ColdJsonlFallback` failover slot was removed. Surviving -/// failover surface is ASAP-tier sketch ↔ Gorilla-S3 archive. +/// **ASAP-first centralization refactor**: the decision tree is now +/// owned here (and consumed identically by `EngineRouter::execute` and +/// `EngineRouter::execute_range`) so the HTTP transport layer never +/// re-derives routing. The policy is: /// -/// Routing rules (mirrors `docs/design-gorilla-s3-cold-engine.md` §8): -/// -/// * Metric configured for `GorillaObjectStore`: always -/// `[GorillaObjectStore]`. Exact-on-archive subsumes approximate-on-warm, -/// so even a `Statistic::Quantile` with an `Approximate` target still -/// routes to the archive when the metric is Gorilla-only — there is no -/// ASAP-tier sketch to fall back to in that deploy shape. -/// * Metric configured for `SketchStore` (or unconfigured / default): -/// `[SketchStore]`. A capability miss in the ASAP tier surfaces -/// as a 404 — the previous JSONL fallback path has been deleted. -/// * Metric configured for `DoubleWrite`: head depends on accuracy hint, -/// tail is the failover sequence (the cost-aware `EngineRouter` picks -/// the head, walks the tail on failure): -/// - `Exact` → `[GorillaObjectStore, SketchStore]` -/// - `Approximate` → `[SketchStore, GorillaObjectStore]` +/// * `accuracy == Exact` → archive only `[GorillaObjectStore]` (served +/// by the `thanos_query` engine). The caller demands an exact answer, +/// so the ε/δ-bounded ASAP-tier sketches are not eligible — go +/// straight to the archive regardless of where the metric is stored. +/// * `accuracy == Approximate` (the default) → ASAP-first failover +/// `[SketchStore, GorillaObjectStore]` for any metric stored in an +/// ASAP-managed tier (`SketchStore`, `GorillaObjectStore`, or +/// `DoubleWrite`): try the warm sketch (`asap_query`) first and fall +/// back to the archive (`thanos_query`) on a capability miss. This +/// collapses the old per-`metric_storage` sequences into one shared +/// ASAP-first-then-archive contract. +/// * `PrometheusRemote` keeps its own single-backend sequence +/// `[PrometheusRemote]` (Phase ε.2): the metric's raw samples never +/// landed in ASAP-managed storage, so there is no ASAP-tier sketch to +/// fall back on and the accuracy hint does not apply. A missing +/// engine surfaces as a `NoEngineRegistered` 503 from the HTTP +/// handler — the correct fail-loud behaviour for a misconfigured +/// deploy. pub fn compatible_storage_backends( _stat: Statistic, accuracy: AccuracyTarget, metric_storage_config: StorageBackend, ) -> Vec { match metric_storage_config { - StorageBackend::GorillaObjectStore => { - vec![StorageBackend::GorillaObjectStore] - } - StorageBackend::SketchStore => { - vec![StorageBackend::SketchStore] - } - StorageBackend::DoubleWrite => match accuracy { - AccuracyTarget::Exact => vec![ - StorageBackend::GorillaObjectStore, - StorageBackend::SketchStore, - ], + // Prometheus-remote owns its own storage; the accuracy hint does + // not apply and there is no ASAP-tier sketch to fall back on. + StorageBackend::PrometheusRemote => vec![StorageBackend::PrometheusRemote], + + // Every ASAP-managed tier shares the same ASAP-first policy, + // gated only on the accuracy target. + StorageBackend::SketchStore + | StorageBackend::GorillaObjectStore + | StorageBackend::DoubleWrite => match accuracy { + // Exact: archive only — the warm sketches are ε/δ-bounded. + AccuracyTarget::Exact => vec![StorageBackend::GorillaObjectStore], + // Approximate: ASAP-tier first, archive (Thanos) fallback. AccuracyTarget::Approximate => vec![ StorageBackend::SketchStore, StorageBackend::GorillaObjectStore, ], }, - // Phase ε.2: Prometheus-remote metrics route only to the - // Prometheus forwarder. There is no ASAP-tier sketch to fall - // back on (the metric's raw samples never landed in - // ASAP-managed storage), so the failover sequence is the - // single backend itself; a missing engine surfaces as a - // `NoEngineRegistered` 503 from the HTTP handler, which is - // the correct fail-loud behaviour for a misconfigured deploy. - StorageBackend::PrometheusRemote => { - vec![StorageBackend::PrometheusRemote] - } } } @@ -1263,71 +1259,50 @@ mod tests { // ----------------------------------------------------------------------- #[test] - fn gorilla_s3_metric_routes_to_archive() { - let backends = compatible_storage_backends( - Statistic::Sum, - AccuracyTarget::Exact, + fn exact_accuracy_routes_to_archive_only() { + // ASAP-first refactor: `Exact` goes straight to the archive + // (Thanos via the GorillaObjectStore slot) regardless of where + // the metric is stored — the warm sketches are ε/δ-bounded. + for cfg in [ StorageBackend::GorillaObjectStore, - ); - assert_eq!(backends, vec![StorageBackend::GorillaObjectStore]); - } - - #[test] - fn asap_query_metric_routes_to_simple_engine() { - let backends = compatible_storage_backends( - Statistic::Quantile, - AccuracyTarget::Approximate, StorageBackend::SketchStore, - ); - // Step-1 of the JSONL deprecation: ASAP-tier only routes - // to itself; the previous `ColdJsonlFallback` failover slot - // has been deleted. - assert_eq!(backends, vec![StorageBackend::SketchStore]); - } - - #[test] - fn double_write_metric_returns_both_options() { - // Exact: archive head, ASAP-tier failover. - let exact = compatible_storage_backends( - Statistic::Sum, - AccuracyTarget::Exact, StorageBackend::DoubleWrite, - ); - assert_eq!( - exact, - vec![ - StorageBackend::GorillaObjectStore, - StorageBackend::SketchStore, - ] - ); - // Approximate: ASAP-tier head (cheaper for ε/δ-bounded - // answers), archive failover. - let approx = compatible_storage_backends( - Statistic::Quantile, - AccuracyTarget::Approximate, - StorageBackend::DoubleWrite, - ); - assert_eq!( - approx, - vec![ - StorageBackend::SketchStore, - StorageBackend::GorillaObjectStore, - ] - ); + ] { + let backends = + compatible_storage_backends(Statistic::Sum, AccuracyTarget::Exact, cfg); + assert_eq!( + backends, + vec![StorageBackend::GorillaObjectStore], + "Exact accuracy must route to archive only for {cfg:?}", + ); + } } - /// Exact-on-archive subsumes approximate-on-warm: a metric configured - /// only for Gorilla-S3 still routes to the archive even when the caller - /// asks for an approximate answer (no ASAP-tier sketch exists to back- - /// fall to in that deploy shape). #[test] - fn gorilla_s3_with_non_exact_accuracy_still_archives() { - let backends = compatible_storage_backends( - Statistic::Quantile, - AccuracyTarget::Approximate, + fn approximate_accuracy_is_asap_first_with_archive_fallback() { + // ASAP-first refactor: every ASAP-managed tier shares the same + // `[SketchStore, GorillaObjectStore]` sequence for `Approximate` + // — try the warm sketch first, fall back to the Thanos archive + // on a capability miss. + for cfg in [ + StorageBackend::SketchStore, StorageBackend::GorillaObjectStore, - ); - assert_eq!(backends, vec![StorageBackend::GorillaObjectStore]); + StorageBackend::DoubleWrite, + ] { + let backends = compatible_storage_backends( + Statistic::Quantile, + AccuracyTarget::Approximate, + cfg, + ); + assert_eq!( + backends, + vec![ + StorageBackend::SketchStore, + StorageBackend::GorillaObjectStore, + ], + "Approximate accuracy must be ASAP-first then archive for {cfg:?}", + ); + } } #[test] @@ -1416,19 +1391,15 @@ mod tests { dispatchable failover (SketchStore, GorillaObjectStore, or \ PrometheusRemote); got {last:?}", ); - // The expected head is determined by `(metric_storage_config, accuracy)`: + // ASAP-first refactor: the head is determined by + // `(metric_storage_config, accuracy)`. PrometheusRemote + // keeps its single-backend slot; every ASAP-managed tier + // goes archive-only on `Exact` and ASAP-tier-first on + // `Approximate`. let expected_head = match (cfg, acc) { - (StorageBackend::GorillaObjectStore, _) => { - StorageBackend::GorillaObjectStore - } - (StorageBackend::SketchStore, _) => StorageBackend::SketchStore, - (StorageBackend::DoubleWrite, AccuracyTarget::Exact) => { - StorageBackend::GorillaObjectStore - } - (StorageBackend::DoubleWrite, AccuracyTarget::Approximate) => { - StorageBackend::SketchStore - } (StorageBackend::PrometheusRemote, _) => StorageBackend::PrometheusRemote, + (_, AccuracyTarget::Exact) => StorageBackend::GorillaObjectStore, + (_, AccuracyTarget::Approximate) => StorageBackend::SketchStore, }; assert_eq!( backends[0], expected_head, diff --git a/data_plane/src/drivers/query/servers/http.rs b/data_plane/src/drivers/query/servers/http.rs index 744eb1c7..5cfcb09c 100644 --- a/data_plane/src/drivers/query/servers/http.rs +++ b/data_plane/src/drivers/query/servers/http.rs @@ -1662,28 +1662,58 @@ async fn process_range_query_request( parsed_request.query, parsed_request.start, parsed_request.end, parsed_request.step ); - // B7.5 retirement — legacy `handle_range_query_promql` is gone. - // Route directly through the modern warm-tier path - // (`execute_range_promql_modern`), which classifies via the - // analyzer + ASAP-tier reducer and returns Matrix shape per the - // `/api/v1/query_range` wire-format requirement. + // ASAP-first centralization refactor — the range path is now a + // thin transport layer. All engine-selection / failover lives in + // `EngineRouter::execute_range`, which walks the shared + // `compatible_storage_backends` policy table: + // * `accuracy == Exact` → archive only (thanos_query) + // * otherwise → ASAP-tier (asap_query) first, fall + // back to the archive on CapabilityMiss. + // The handler no longer reaches into `query_engine` / + // `engine_by_id` to do its own Thanos lookup; it just resolves the + // metric's storage axis and hands the range request to the router. let start_ms = (parsed_request.start * 1000.0) as u64; let end_ms = (parsed_request.end * 1000.0) as u64; let step_ms = (parsed_request.step * 1000.0) as u64; - let modern_result = state - .query_engine - .execute_range_promql_modern( + + // Range handlers don't read the `X-ASAP-Tenant` header, so resolve + // the metric's storage axis against the `default` tenant's routing + // table — the same resolution the instant path applies when no + // tenant header is present. + let metric_storage = resolve_metric_storage(state, &parsed_request.query, "default"); + + // Match the instant path's hardcoding of `(Sum, Approximate)`. + // TODO: derive `accuracy` (and `stat`) from the request rather than + // pinning Approximate — once the request carries an accuracy hint, + // an `Exact` range query will route straight to the archive via the + // shared policy table. + let stat = Statistic::Sum; + let accuracy = AccuracyTarget::Approximate; + + debug!( + "Dispatching range query via EngineRouter: query='{}' metric_storage={:?} \ + stat={:?} accuracy={:?}", + parsed_request.query, metric_storage, stat, accuracy, + ); + + let router_result = state + .query_router + .execute_range( &parsed_request.query, + stat, + accuracy, + metric_storage, start_ms, end_ms, step_ms, ) .await; - match modern_result { + + match router_result { Ok(query_result) => { let query_duration = query_start_time.elapsed(); debug!( - "Modern range execute took: {:.2}ms", + "EngineRouter range dispatch took: {:.2}ms", query_duration.as_secs_f64() * 1000.0 ); let total_duration = start_time.elapsed(); @@ -1703,15 +1733,52 @@ async fn process_range_query_request( Err(status) => status.into_response(), } } - Err(_) => { - debug!( - "Modern range-query path returned CapabilityMiss for query='{}', \ - falling through to unsupported", - parsed_request.query + Err(EngineRouterError::NoEngineRegistered { tried, registered }) => { + warn!( + tried = ?tried, + registered = ?registered, + "EngineRouter (range): no engine registered for any compatible backend", ); - match state.adapter.format_unsupported_query_response().await { - Ok(json) => json.into_response(), - Err(status) => status.into_response()} + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "status": "error", + "errorType": "internal", + "error": format!( + "no engine registered for any compatible backend; tried {tried:?}, registered={registered:?}" + )})), + ) + .into_response() + } + Err(EngineRouterError::AllFailed { last }) => { + use crate::query_engines::EngineError; + warn!(error = %last, "EngineRouter (range): all compatible engines failed"); + // A terminal CapabilityMiss means no tier could serve the + // range query — surface the adapter's "unsupported query" + // response (the wire shape browsers / dashboards expect), + // matching the pre-refactor fall-through. A Backend error + // is a real upstream failure → 5xx. + match &last { + EngineError::CapabilityMiss { .. } => { + debug!( + "Range query CapabilityMiss across all tiers for query='{}', \ + falling through to unsupported", + parsed_request.query + ); + match state.adapter.format_unsupported_query_response().await { + Ok(json) => json.into_response(), + Err(status) => status.into_response(), + } + } + EngineError::Backend { .. } => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({ + "status": "error", + "errorType": "internal", + "error": last.to_string()})), + ) + .into_response(), + } } } } @@ -3252,16 +3319,24 @@ aggregations: #[tokio::test] async fn http_returns_503_when_no_engines_registered() { - // Pin `storage_backend = GorillaObjectStore` but register no - // archive engine (only `ASAPQueryEngine` is registered under - // `asap_query`). The router walks - // `compatible_storage_backends = [GorillaObjectStore]` and - // bails out with `NoEngineRegistered`, which the HTTP layer - // surfaces as 503. Step-1 of the JSONL deprecation removed - // the `ColdJsonlFallback` failover slot, so this is the - // canonical "engine missing" path now. + // Pin `storage_backend = PrometheusRemote` but register no + // Prometheus forwarder (only `ASAPQueryEngine` is registered + // under `asap_query`). The router walks + // `compatible_storage_backends = [PrometheusRemote]` and bails + // out with `NoEngineRegistered`, which the HTTP layer surfaces + // as 503. + // + // ASAP-first refactor note: this test used to pin + // `GorillaObjectStore` and rely on the old archive-only + // `[GorillaObjectStore]` sequence. Under the ASAP-first policy + // a `GorillaObjectStore` metric now resolves to + // `[SketchStore, GorillaObjectStore]` — the registered + // ASAP engine is tried first and CapabilityMisses, yielding a + // 404 (`AllFailed`) rather than a 503. `PrometheusRemote` keeps + // its single-backend slot, so it remains the canonical "engine + // missing → 503" path. let server_port = - setup_test_server_with_empty_router(StorageBackend::GorillaObjectStore).await; + setup_test_server_with_empty_router(StorageBackend::PrometheusRemote).await; let client = Client::new(); let resp = client .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) @@ -4255,6 +4330,70 @@ aggregations: ); } + // ── Path A2 range-query e2e test ──────────────────────────────────── + + /// GET /api/v1/query_range returns success when the ASAP sketch + /// tier misses and the `ThanosQueryEngine` is registered. Covers + /// the full centralized HTTP path: + /// browser → ASAP backend → EngineRouter::execute_range + /// → asap_query (CapabilityMiss) → thanos_query + /// → mock thanos → matrix response + #[tokio::test] + async fn http_query_range_forwards_to_thanos_when_asap_misses() { + use crate::query_engines::thanos_query_engine::forward::test_support::spawn_mock_thanos_capture_range; + use crate::query_engines::thanos_query_engine::forward::test_support::CANNED_MATRIX_BODY; + use crate::query_engines::thanos_query_engine::{ThanosQueryConfig, ThanosQueryEngine}; + + let (mock_url, _captured, _mock_handle) = + spawn_mock_thanos_capture_range(CANNED_MATRIX_BODY).await; + let cfg = ThanosQueryConfig { + base_url: mock_url, + request_timeout: std::time::Duration::from_secs(5), + }; + let engine = ThanosQueryEngine::new(cfg).expect("engine"); + let arc_engine: Arc = Arc::new(engine); + + // GorillaObjectStore metric → `Approximate` policy yields + // `[SketchStore, GorillaObjectStore]`: the ASAP tier misses + // (no sketch index for the query) and the router falls over to + // the registered thanos_query engine. + let server_port = setup_test_server_with_named_router( + StorageBackend::GorillaObjectStore, + vec![arc_engine], + ) + .await; + + let client = Client::new(); + let resp = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query_range")) + .query(&[ + ("query", "rate(http_requests_total[1m])"), + ("start", "1700000000"), + ("end", "1700003600"), + ("step", "15"), + ]) + .send() + .await + .expect("request"); + + assert!( + resp.status().is_success(), + "expected 2xx from range query forwarded to Thanos; got {}", + resp.status() + ); + let body: serde_json::Value = resp.json().await.unwrap(); + assert_eq!( + body["status"].as_str().unwrap_or(""), + "success", + "wire response must carry status=success; got {body}" + ); + assert_eq!( + body["data"]["resultType"].as_str().unwrap_or(""), + "matrix", + "wire response must carry resultType=matrix; got {body}" + ); + } + #[tokio::test] async fn http_engine_override_can_target_thanos_query_id() { // X-ASAP-Engine: thanos_query must reach the forwarder diff --git a/data_plane/src/query_engines/asap_query_engine/engine.rs b/data_plane/src/query_engines/asap_query_engine/engine.rs index 9e5ac51c..a6712e73 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -1641,6 +1641,26 @@ impl crate::query_engines::routing::query_engine_routing::QueryEngine for ASAPQu )) } + /// Range-query entry point for the [`EngineRouter`] failover loop. + /// + /// Delegates to the inherent `execute_range_promql_modern`, which + /// runs the ASAP-tier reducer over `[start_ms, end_ms]` and returns + /// a `matrix` result. Without this override the router would hit the + /// trait default (`CapabilityMiss`) and never reach the ASAP-tier + /// range path, so every range query would fall straight through to + /// the archive even when the warm sketches can answer it. + async fn execute_range( + &self, + query: &str, + start_ms: u64, + end_ms: u64, + step_ms: u64, + ) -> Result + { + self.execute_range_promql_modern(query, start_ms, end_ms, step_ms) + .await + } + fn capabilities(&self) -> crate::query_engines::routing::query_engine_routing::EngineCapabilities { crate::query_engines::routing::query_engine_routing::EngineCapabilities { data_source_id: asap_types::StorageBackend::SketchStore.data_source_id(), diff --git a/data_plane/src/query_engines/routing/query_engine_routing.rs b/data_plane/src/query_engines/routing/query_engine_routing.rs index d0dfda8b..756f7f85 100644 --- a/data_plane/src/query_engines/routing/query_engine_routing.rs +++ b/data_plane/src/query_engines/routing/query_engine_routing.rs @@ -69,6 +69,25 @@ pub trait QueryEngine: Send + Sync { /// Answer `query` against this engine's storage tier. async fn execute(&self, query: &str) -> Result; + /// Forward a range query to this engine's storage tier. + /// + /// Params are milliseconds since epoch. The default returns `CapabilityMiss` + /// so existing impls need not change; override when the engine has native + /// range-query support (e.g. Thanos, or the ASAP-tier reducer). + async fn execute_range( + &self, + query: &str, + start_ms: u64, + end_ms: u64, + step_ms: u64, + ) -> Result { + let _ = (query, start_ms, end_ms, step_ms); + Err(EngineError::capability_miss( + self.capabilities().data_source_id, + "execute_range not implemented for this engine", + )) + } + /// What this engine can serve. Cheap; the router calls it on every /// `register` and may re-call to refresh cost estimates. fn capabilities(&self) -> EngineCapabilities; @@ -236,6 +255,90 @@ impl EngineRouter { last: last_err.expect("at least one engine ran (any_engine_tried=true)"), }) } + + /// Range-query sibling of [`Self::execute`]. Walks the identical + /// compatible-backend failover sequence for `(stat, accuracy, + /// metric_storage)`, but dispatches to `engine.execute_range(query, + /// start_ms, end_ms, step_ms)` instead of `engine.execute(query)`. + /// + /// The decision tree is centralized here (not in `http.rs`): the + /// shared [`compatible_storage_backends`] policy table decides the + /// order — `accuracy == Exact` yields archive-only + /// `[GorillaObjectStore]`, otherwise the ASAP-first + /// `[SketchStore, GorillaObjectStore]` sequence. On + /// [`EngineError::CapabilityMiss`] or [`EngineError::Backend`] the + /// router falls through to the next compatible backend; the + /// `thanos_query` archive engine answers the matrix natively. + /// + /// On exhaustion returns the same [`EngineRouterError`] variants as + /// [`Self::execute`]. + pub async fn execute_range( + &self, + query: &str, + stat: Statistic, + accuracy: AccuracyTarget, + metric_storage: StorageBackend, + start_ms: u64, + end_ms: u64, + step_ms: u64, + ) -> Result { + let backends = compatible_storage_backends(stat, accuracy, metric_storage); + debug!( + query = query, + stat = ?stat, + accuracy = ?accuracy, + metric_storage = ?metric_storage, + backends = ?backends, + start_ms, + end_ms, + step_ms, + "router: dispatching range query", + ); + + let mut last_err: Option = None; + let mut any_engine_tried = false; + + for backend in &backends { + let id = backend.data_source_id(); + let Some(engine) = self.engines.get(id) else { + debug!( + backend = ?backend, + data_source_id = id, + "router: no engine registered, trying next failover (range)", + ); + continue; + }; + any_engine_tried = true; + match engine.execute_range(query, start_ms, end_ms, step_ms).await { + Ok(result) => { + debug!( + backend = ?backend, + "router: range dispatch succeeded", + ); + return Ok(result); + } + Err(e) => { + warn!( + backend = ?backend, + error = %e, + "router: engine failed on range query, falling through to next backend", + ); + last_err = Some(e); + } + } + } + + if !any_engine_tried { + return Err(EngineRouterError::NoEngineRegistered { + tried: backends, + registered: self.engines.keys().copied().collect(), + }); + } + + Err(EngineRouterError::AllFailed { + last: last_err.expect("at least one engine ran (any_engine_tried=true)"), + }) + } } // --------------------------------------------------------------------------- @@ -295,6 +398,29 @@ mod tests { )), } } + // Mirror `execute`'s outcome so the router's range failover loop + // can be exercised hermetically (the trait default would always + // CapabilityMiss, which wouldn't test Ok / Backend dispatch). + async fn execute_range( + &self, + _query: &str, + _start_ms: u64, + _end_ms: u64, + _step_ms: u64, + ) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + match self.outcome { + Outcome::Ok => Ok(QueryResult::matrix(Vec::new())), + Outcome::Backend => Err(EngineError::backend( + self.caps.data_source_id, + "simulated backend failure (range)", + )), + Outcome::CapabilityMiss => Err(EngineError::capability_miss( + self.caps.data_source_id, + "no compatible aggregation (range)", + )), + } + } fn capabilities(&self) -> EngineCapabilities { self.caps } @@ -353,31 +479,33 @@ mod tests { } #[tokio::test] - async fn router_falls_back_to_warm_when_archive_fails_on_double_write() { - // Double-write deploy with `Exact` head: archive head fails, - // router falls through to the ASAP-tier sketch (the only - // remaining failover after Step-1 deleted JSONL). + async fn router_falls_back_to_archive_when_asap_misses_on_double_write() { + // ASAP-first refactor: an `Approximate` double-write query + // tries the ASAP-tier sketch first; on a CapabilityMiss the + // router falls through to the Thanos archive (the + // `[SketchStore, GorillaObjectStore]` failover sequence). let mut router = EngineRouter::new(); + let (warm, warm_calls) = + StubEngine::new(StorageBackend::SketchStore, Outcome::CapabilityMiss); let (gorilla, gorilla_calls) = - StubEngine::new(StorageBackend::GorillaObjectStore, Outcome::Backend); - let (warm, warm_calls) = StubEngine::new(StorageBackend::SketchStore, Outcome::Ok); - router.register(gorilla); + StubEngine::new(StorageBackend::GorillaObjectStore, Outcome::Ok); router.register(warm); + router.register(gorilla); let result = router .execute( "sum_over_time(foo[5m])", Statistic::Sum, - AccuracyTarget::Exact, + AccuracyTarget::Approximate, StorageBackend::DoubleWrite, ) .await; assert!( result.is_ok(), - "router must reach ASAP-tier when the archive head fails", + "router must reach the archive when the ASAP-tier head misses", ); - assert_eq!(gorilla_calls.load(Ordering::SeqCst), 1); assert_eq!(warm_calls.load(Ordering::SeqCst), 1); + assert_eq!(gorilla_calls.load(Ordering::SeqCst), 1); } #[tokio::test] @@ -393,9 +521,15 @@ mod tests { .await; match result { Err(EngineRouterError::NoEngineRegistered { tried, registered }) => { - // Step-1 deleted the JSONL failover slot, so the - // SketchStore failover sequence is just itself. - assert_eq!(tried, vec![StorageBackend::SketchStore]); + // ASAP-first refactor: an `Approximate` query walks the + // `[SketchStore, GorillaObjectStore]` failover sequence. + assert_eq!( + tried, + vec![ + StorageBackend::SketchStore, + StorageBackend::GorillaObjectStore, + ] + ); assert!(registered.is_empty()); } other => panic!("expected NoEngineRegistered, got {other:?}"), @@ -541,4 +675,100 @@ mod tests { other => panic!("expected NoEngineRegistered, got {other:?}"), } } + + // ── range-query dispatch (`EngineRouter::execute_range`) ──────────── + + /// ASAP-first refactor: an `Approximate` range query tries the + /// ASAP-tier sketch first and answers there when it can. + #[tokio::test] + async fn router_range_dispatches_to_asap_tier_first() { + let mut router = EngineRouter::new(); + let (warm, warm_calls) = StubEngine::new(StorageBackend::SketchStore, Outcome::Ok); + let (archive, archive_calls) = + StubEngine::new(StorageBackend::GorillaObjectStore, Outcome::Ok); + router.register(warm); + router.register(archive); + + let result = router + .execute_range( + "rate(http_requests_total[1m])", + Statistic::Sum, + AccuracyTarget::Approximate, + StorageBackend::SketchStore, + 1_700_000_000_000, + 1_700_003_600_000, + 15_000, + ) + .await; + assert!(result.is_ok()); + assert_eq!(warm_calls.load(Ordering::SeqCst), 1); + assert_eq!( + archive_calls.load(Ordering::SeqCst), + 0, + "archive must not run when the ASAP-tier answers the range query", + ); + } + + /// ASAP-first refactor: on a CapabilityMiss in the ASAP tier the + /// range query fails over to the Thanos archive. + #[tokio::test] + async fn router_range_falls_over_to_archive_on_capability_miss() { + let mut router = EngineRouter::new(); + let (warm, warm_calls) = + StubEngine::new(StorageBackend::SketchStore, Outcome::CapabilityMiss); + let (archive, archive_calls) = + StubEngine::new(StorageBackend::GorillaObjectStore, Outcome::Ok); + router.register(warm); + router.register(archive); + + let result = router + .execute_range( + "rate(http_requests_total[1m])", + Statistic::Sum, + AccuracyTarget::Approximate, + StorageBackend::SketchStore, + 0, + 1_000, + 1_000, + ) + .await; + assert!(result.is_ok()); + assert_eq!(warm_calls.load(Ordering::SeqCst), 1); + assert_eq!( + archive_calls.load(Ordering::SeqCst), + 1, + "range query must fail over to the archive on ASAP-tier CapabilityMiss", + ); + } + + /// ASAP-first refactor: an `Exact` range query routes straight to + /// the archive — the ASAP-tier sketch is never consulted. + #[tokio::test] + async fn router_range_exact_goes_straight_to_archive() { + let mut router = EngineRouter::new(); + let (warm, warm_calls) = StubEngine::new(StorageBackend::SketchStore, Outcome::Ok); + let (archive, archive_calls) = + StubEngine::new(StorageBackend::GorillaObjectStore, Outcome::Ok); + router.register(warm); + router.register(archive); + + let result = router + .execute_range( + "audit_events", + Statistic::Sum, + AccuracyTarget::Exact, + StorageBackend::DoubleWrite, + 0, + 1_000, + 1_000, + ) + .await; + assert!(result.is_ok()); + assert_eq!(archive_calls.load(Ordering::SeqCst), 1); + assert_eq!( + warm_calls.load(Ordering::SeqCst), + 0, + "Exact range query must skip the ASAP-tier and archive only", + ); + } } diff --git a/data_plane/src/query_engines/thanos_query_engine/forward.rs b/data_plane/src/query_engines/thanos_query_engine/forward.rs index 3abc4770..8016d479 100644 --- a/data_plane/src/query_engines/thanos_query_engine/forward.rs +++ b/data_plane/src/query_engines/thanos_query_engine/forward.rs @@ -129,6 +129,10 @@ impl ThanosQueryConfig { fn instant_endpoint(&self) -> String { format!("{}/api/v1/query", self.base_url) } + + fn range_endpoint(&self) -> String { + format!("{}/api/v1/query_range", self.base_url) + } } /// Forwards PromQL queries to an upstream `thanos-query` sidecar @@ -254,6 +258,67 @@ impl ThanosQueryEngine { .map_err(ThanosQueryError::ParseError)?; Ok(result) } + + /// Forward `query` to `${base_url}/api/v1/query_range` with + /// the given time bounds and return the matrix as a [`QueryResult`]. + pub async fn query_range( + &self, + query: &str, + start_ms: u64, + end_ms: u64, + step_ms: u64, + ) -> Result { + let started = Instant::now(); + let url = self.config.range_endpoint(); + let start_s = format!("{:.3}", start_ms as f64 / 1000.0); + let end_s = format!("{:.3}", end_ms as f64 / 1000.0); + let step_s = format!("{:.3}", step_ms as f64 / 1000.0); + debug!( + url = %url, + query = query, + start = %start_s, + end = %end_s, + step = %step_s, + "thanos-forward: issuing range query", + ); + + let resp = self + .client + .post(&url) + .form(&[ + ("query", query), + ("start", start_s.as_str()), + ("end", end_s.as_str()), + ("step", step_s.as_str()), + ]) + .send() + .await + .map_err(|e| ThanosQueryError::Unreachable(e.to_string()))?; + + let status = resp.status(); + if status.is_server_error() { + return Err(ThanosQueryError::Unreachable(format!( + "upstream returned {status}", + ))); + } + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + return Err(ThanosQueryError::BadQuery { + status: status.as_u16(), + body, + }); + } + + let payload: ThanosResponse = resp + .json() + .await + .map_err(|e| ThanosQueryError::ParseError(e.to_string()))?; + + let elapsed_ms = started.elapsed().as_millis(); + let result = build_result_from_thanos_payload(payload, elapsed_ms) + .map_err(ThanosQueryError::ParseError)?; + Ok(result) + } } #[async_trait] @@ -295,6 +360,43 @@ impl QueryEngine for ThanosQueryEngine { } } + async fn execute_range( + &self, + query: &str, + start_ms: u64, + end_ms: u64, + step_ms: u64, + ) -> Result { + match self.query_range(query, start_ms, end_ms, step_ms).await { + Ok(result) => Ok(result), + Err(ThanosQueryError::Unreachable(reason)) => { + warn!( + engine = self.data_source_id, + error = %reason, + "thanos-forward: upstream unreachable (range query)", + ); + Err(crate::query_engines::EngineError::backend( + self.data_source_id, + format!("thanos_unreachable: {reason}"), + )) + } + Err(ThanosQueryError::BadQuery { status, body }) => { + Err(crate::query_engines::EngineError::capability_miss( + self.data_source_id, + format!("thanos rejected range query (status {status}): {body}"), + )) + } + Err(ThanosQueryError::ParseError(msg)) => Err(crate::query_engines::EngineError::backend( + self.data_source_id, + format!("thanos range response parse error: {msg}"), + )), + Err(ThanosQueryError::ConfigInvalid(msg)) => Err(crate::query_engines::EngineError::backend( + self.data_source_id, + format!("thanos client misconfigured: {msg}"), + )), + } + } + fn capabilities(&self) -> EngineCapabilities { EngineCapabilities { data_source_id: self.data_source_id, @@ -662,6 +764,97 @@ pub mod test_support { ] } }"#; + + /// A minimal canned matrix response — one series with two samples. + pub const CANNED_MATRIX_BODY: &str = r#"{ + "status": "success", + "data": { + "resultType": "matrix", + "result": [ + { + "metric": {"__name__": "http_requests_total", "job": "api"}, + "values": [[1700000000.0, "42"], [1700000015.0, "43"]] + } + ] + } + }"#; + + use std::collections::HashMap; + use std::sync::{Arc, Mutex}; + + #[derive(Debug, Clone, Default)] + pub struct CapturedParams(pub HashMap); + + /// Spawn a mock thanos that captures form params of the first POST + /// to `/api/v1/query_range`, then returns `canned_body`. + pub async fn spawn_mock_thanos_capture_range( + canned_body: &'static str, + ) -> (String, Arc>>, tokio::task::JoinHandle<()>) { + use axum::extract::Form; + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let local: SocketAddr = listener.local_addr().expect("local_addr"); + let base_url = format!("http://{local}"); + let captured: Arc>> = Arc::new(Mutex::new(None)); + let cap2 = captured.clone(); + let app = axum::Router::new() + .route( + "/api/v1/query_range", + axum::routing::post(move |Form(params): Form>| { + let captured = cap2.clone(); + async move { + *captured.lock().unwrap() = Some(CapturedParams(params)); + canned_body + } + }), + ) + .route("/api/v1/query", axum::routing::post(move || async move { canned_body })); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.expect("capture serve"); + }); + tokio::task::yield_now().await; + (base_url, captured, handle) + } + + /// Spawn a mock that sleeps 10 s before answering `/api/v1/query_range`. + pub async fn spawn_mock_thanos_slow_range() -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let local: SocketAddr = listener.local_addr().expect("local_addr"); + let base_url = format!("http://{local}"); + let app: Router = Router::new() + .route( + "/api/v1/query_range", + axum::routing::post(|| async { + tokio::time::sleep(std::time::Duration::from_secs(10)).await; + "{}" + }), + ) + .route("/api/v1/query", axum::routing::post(|| async { "{}" })); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.expect("slow serve"); + }); + tokio::task::yield_now().await; + (base_url, handle) + } + + /// Spawn a mock that returns `status_code` on `/api/v1/query_range`. + pub async fn spawn_mock_thanos_status_range(status_code: u16) -> (String, JoinHandle<()>) { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let local: SocketAddr = listener.local_addr().expect("local_addr"); + let base_url = format!("http://{local}"); + let sc = axum::http::StatusCode::from_u16(status_code) + .unwrap_or(axum::http::StatusCode::INTERNAL_SERVER_ERROR); + let app: Router = Router::new() + .route( + "/api/v1/query_range", + axum::routing::post(move || async move { sc }), + ) + .route("/api/v1/query", axum::routing::post(|| async { axum::http::StatusCode::OK })); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.expect("status serve"); + }); + tokio::task::yield_now().await; + (base_url, handle) + } } // --------------------------------------------------------------------------- @@ -853,4 +1046,122 @@ mod tests { other => panic!("expected Vector, got {other:?}"), } } + + // ── query_range unit tests ────────────────────────────────────────── + + #[tokio::test] + async fn query_range_calls_range_endpoint_not_instant() { + use test_support::{spawn_mock_thanos_capture_range, CANNED_MATRIX_BODY}; + let (url, _captured, _handle) = + spawn_mock_thanos_capture_range(CANNED_MATRIX_BODY).await; + let engine = ThanosQueryEngine::new(config_for(&url)).expect("engine"); + let result = engine + .query_range("rate(http_requests_total[1m])", 1_700_000_000_000, 1_700_003_600_000, 15_000) + .await; + assert!(result.is_ok(), "expected Ok; got {result:?}"); + match result.unwrap() { + QueryResult::Matrix(_) => {} + other => panic!("expected Matrix result, got {other:?}"), + } + } + + #[tokio::test] + async fn query_range_passes_params_correctly() { + use test_support::{spawn_mock_thanos_capture_range, CANNED_MATRIX_BODY}; + let (url, captured, _handle) = + spawn_mock_thanos_capture_range(CANNED_MATRIX_BODY).await; + let engine = ThanosQueryEngine::new(config_for(&url)).expect("engine"); + engine + .query_range("up", 1_700_000_000_000, 1_700_003_600_000, 15_000) + .await + .expect("ok"); + let params = captured.lock().unwrap().clone().expect("params captured"); + assert_eq!(params.0.get("query").map(|s| s.as_str()), Some("up"), + "query param must be forwarded verbatim"); + let start: f64 = params.0["start"].parse().unwrap(); + assert!((start - 1_700_000_000.0).abs() < 0.1, "start must be unix seconds: {start}"); + let end: f64 = params.0["end"].parse().unwrap(); + assert!((end - 1_700_003_600.0).abs() < 0.1, "end must be unix seconds: {end}"); + let step: f64 = params.0["step"].parse().unwrap(); + assert!((step - 15.0).abs() < 0.1, "step must be seconds: {step}"); + } + + #[tokio::test] + async fn query_range_parses_matrix_response() { + use test_support::{spawn_mock_thanos_capture_range, CANNED_MATRIX_BODY}; + let (url, _captured, _handle) = + spawn_mock_thanos_capture_range(CANNED_MATRIX_BODY).await; + let engine = ThanosQueryEngine::new(config_for(&url)).expect("engine"); + let result = engine + .query_range("up", 1_700_000_000_000, 1_700_003_600_000, 15_000) + .await + .expect("ok"); + match result { + QueryResult::Matrix(mv) => { + assert_eq!(mv.values.len(), 1, "one series in canned body"); + assert_eq!(mv.values[0].samples.len(), 2, "two samples in canned body"); + assert!((mv.values[0].samples[0].value - 42.0).abs() < 1e-9); + assert!((mv.values[0].samples[1].value - 43.0).abs() < 1e-9); + } + other => panic!("expected Matrix result, got {other:?}"), + } + } + + #[tokio::test] + async fn query_range_4xx_returns_capability_miss() { + use test_support::spawn_mock_thanos_status_range; + let (url, _handle) = spawn_mock_thanos_status_range(400).await; + let engine = ThanosQueryEngine::new(config_for(&url)).expect("engine"); + let raw = engine.query_range("bad[", 0, 1_000, 1_000).await; + assert!( + matches!(raw, Err(ThanosQueryError::BadQuery { .. })), + "4xx must surface as BadQuery; got {raw:?}", + ); + let trait_err = QueryEngine::execute_range(&engine, "bad[", 0, 1_000, 1_000) + .await.unwrap_err(); + assert!( + matches!(trait_err, crate::query_engines::EngineError::CapabilityMiss { .. }), + "4xx must fold into CapabilityMiss via trait; got {trait_err:?}", + ); + } + + #[tokio::test] + async fn query_range_5xx_returns_backend_error() { + use test_support::spawn_mock_thanos_status_range; + let (url, _handle) = spawn_mock_thanos_status_range(503).await; + let engine = ThanosQueryEngine::new(config_for(&url)).expect("engine"); + let raw = engine.query_range("up", 0, 1_000, 1_000).await; + assert!( + matches!(raw, Err(ThanosQueryError::Unreachable(_))), + "5xx must surface as Unreachable; got {raw:?}", + ); + let trait_err = QueryEngine::execute_range(&engine, "up", 0, 1_000, 1_000) + .await.unwrap_err(); + assert!( + matches!(trait_err, crate::query_engines::EngineError::Backend { .. }), + "5xx must fold into Backend via trait; got {trait_err:?}", + ); + } + + #[tokio::test] + async fn query_range_timeout_returns_backend_error() { + use test_support::spawn_mock_thanos_slow_range; + let (url, _handle) = spawn_mock_thanos_slow_range().await; + let cfg = ThanosQueryConfig { + base_url: url.trim_end_matches('/').to_string(), + request_timeout: Duration::from_millis(100), + }; + let engine = ThanosQueryEngine::new(cfg).expect("engine"); + let raw = engine.query_range("up", 0, 1_000_000_000, 1_000).await; + assert!( + matches!(raw, Err(ThanosQueryError::Unreachable(_))), + "timeout must surface as Unreachable; got {raw:?}", + ); + let trait_err = QueryEngine::execute_range(&engine, "up", 0, 1_000_000_000, 1_000) + .await.unwrap_err(); + assert!( + matches!(trait_err, crate::query_engines::EngineError::Backend { .. }), + "timeout must fold into Backend via trait; got {trait_err:?}", + ); + } }