From 36ca4af44c82408746cb442e4572f6ca45edd18a Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Wed, 6 May 2026 01:07:21 -0400 Subject: [PATCH] =?UTF-8?q?feat(engine):=20GorillaQueryEngine=20=E2=80=94?= =?UTF-8?q?=20exact=20PromQL=20over=20Gorilla-S3=20chunks=20(Phase=204)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 4 of the Gorilla-S3-cold-engine — adds a sibling engine to `SimpleEngine` that executes PromQL exactly against the Phase-3 `GorillaS3ColdStore`. Returns answers carrying an `AccuracyEnvelope { kind: Exact, ε: 0, δ: 0 }` and a `data_source: gorilla_archive` info marker so callers / dashboards can distinguish cold-tier exact answers from warm-tier sketch answers. Architecture ------------ `asap-query-engine/src/engines/gorilla_engine/` (new): * `mod.rs` — `GorillaQueryEngine` (holds `Arc` so tests can inject mocks; production constructor `with_gorilla_s3` keeps the design.md type signature). Public surface: `execute(query)` + `execute_at(query, now_ms)`. Wraps results via `wrap_result` (exact accuracy envelope + window). `GorillaEngineConfig { max_buffered_samples, query_timeout_secs }` with sensible defaults (10M / 30s). `EngineError` enum covers Plan / ColdStore / TooManySamples / Timeout. `ExecutionOutcome` carries `(value, samples_scanned, chunks_fetched)` + an `info_lines` builder pinning the on-wire info strings. * `query_planner.rs` — minimal PromQL → `QueryPlan` translator (metric, half-open `[start_ms, end_ms)`, `QueryStatistic`). Supports `sum/count/avg/min/max_over_time`, `rate`, `increase`, `quantile_over_time(φ, m[range])`, and `topk(k, )` (PromQL grammar requires the inner be a vector so `topk(k, sum_over_time(m[range]))` is the legal spelling). Tests + caller can use `plan_query_at` to pin `now_ms`. * `exact_executor.rs` — `ExactExecutor` dispatches per `QueryStatistic`: - **Streaming additive** (`Sum/Count/Avg/Min/Max/Rate/Increase`): `list_chunks` → for each chunk `read_chunk` → fold into a bounded `AdditiveAccumulator` → drop the decoded chunk before fetching the next. Memory is O(1) per query. Rate / Increase track first/last `(ts, value)` and divide by `range_seconds` at finalisation. - **Buffered** (`Quantile / TopK`): `collect_buffered_samples` materialises every in-range sample up to `max_buffered_samples` (errors with `TooManySamples` otherwise). Quantile sorts + nearest-rank index; TopK sorts descending and returns sum of the top-k values. * `tests.rs` — 20 unit tests via an in-process `MockColdStore` satisfying `ColdStore` (no S3 / disk dependency). Pinned NOW via `execute_at` so time math is deterministic. Covers every `QueryStatistic` happy path, the empty-data sentinel, the outside-range filter, the buffered-budget guard, the result wrapping (accuracy envelope + `data_source: gorilla_archive` marker), and the timeout path. Test additions (15 spec'd + 5 planner) -------------------------------------- * `execute_sum_over_time_streaming` * `execute_count_over_time` * `execute_avg_over_time` * `execute_min_over_time` / `execute_max_over_time` * `execute_rate_basic` * `execute_increase_basic` * `execute_quantile_buffered_basic` * `execute_quantile_too_many_samples_errors` * `execute_topk_basic` * `execute_empty_chunks_returns_zero_or_nan` * `execute_chunks_partially_outside_range_filtered` * `result_carries_exact_accuracy_envelope` * `result_includes_data_source_gorilla_archive` * `engine_respects_config_timeout` * `query_planner::tests::{plans_sum_over_time, plans_quantile_over_time, plans_topk, rejects_binary_expression, streaming_classification}` Touches only the new directory + `engines/mod.rs` (re-exports). Zero changes to `simple_engine.rs`, the cold store, or the sketch warm-tier. No new Cargo deps — `asap-gorilla` / `tokio` were already pulled in by PR #84. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../engines/gorilla_engine/exact_executor.rs | 316 +++++++++ .../src/engines/gorilla_engine/mod.rs | 266 ++++++++ .../engines/gorilla_engine/query_planner.rs | 307 +++++++++ .../src/engines/gorilla_engine/tests.rs | 615 ++++++++++++++++++ asap-query-engine/src/engines/mod.rs | 4 + 5 files changed, 1508 insertions(+) create mode 100644 asap-query-engine/src/engines/gorilla_engine/exact_executor.rs create mode 100644 asap-query-engine/src/engines/gorilla_engine/mod.rs create mode 100644 asap-query-engine/src/engines/gorilla_engine/query_planner.rs create mode 100644 asap-query-engine/src/engines/gorilla_engine/tests.rs diff --git a/asap-query-engine/src/engines/gorilla_engine/exact_executor.rs b/asap-query-engine/src/engines/gorilla_engine/exact_executor.rs new file mode 100644 index 000000000..6398b0dd1 --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/exact_executor.rs @@ -0,0 +1,316 @@ +//! Per-statistic executors for the Phase-4 Gorilla engine. +//! +//! Two strategies, picked by [`super::query_planner::QueryStatistic::is_streaming_additive`]: +//! +//! * **Streaming-additive** — `Sum / Count / Avg / Min / Max / +//! Rate / Increase`. Walk chunks one at a time, fold each +//! sample into a tiny accumulator, drop the decoded chunk +//! before fetching the next one. Memory is O(1) per query. +//! * **Buffered** — `Quantile / TopK / Cardinality`. Materialise +//! every in-range sample, then sort or otherwise post-process. +//! Bounded by [`super::GorillaEngineConfig::max_buffered_samples`]; +//! over-budget queries fail fast with +//! [`super::EngineError::TooManySamples`] rather than OOM. + +use std::sync::Arc; + +use tracing::debug; + +use crate::drivers::query::fallback::cold_store::{ColdStore, RawSample}; + +use super::query_planner::{QueryPlan, QueryStatistic}; +use super::{EngineError, ExecutionOutcome, GorillaEngineConfig}; + +/// Streaming-additive operation tag — what the per-sample fold +/// does. Pulled out so [`ExactExecutor::execute_streaming_additive`] +/// is a single function regardless of which stat is being computed. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AdditiveOp { + Sum, + Count, + /// `(sum, count)` — the engine divides at the end. + Avg, + Min, + Max, + /// `last - first` over the time-ordered samples. + Increase, + /// `(last - first) / range_seconds`. + Rate, +} + +/// Per-statistic executor. Holds an `Arc` so the +/// engine + executor share the same cold-tier handle without +/// re-implementing trait dispatch. +pub struct ExactExecutor { + cold_store: Arc, + config: GorillaEngineConfig, +} + +impl ExactExecutor { + pub fn new(cold_store: Arc, config: GorillaEngineConfig) -> Self { + Self { cold_store, config } + } + + /// Top-level dispatch — picks streaming vs buffered based on + /// the plan's statistic. + pub async fn execute_plan(&self, plan: &QueryPlan) -> Result { + match &plan.statistic { + QueryStatistic::SumOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Sum).await + } + QueryStatistic::CountOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Count) + .await + } + QueryStatistic::AvgOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Avg).await + } + QueryStatistic::MinOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Min).await + } + QueryStatistic::MaxOverTime => { + self.execute_streaming_additive(plan, AdditiveOp::Max).await + } + QueryStatistic::Rate => self.execute_streaming_additive(plan, AdditiveOp::Rate).await, + QueryStatistic::Increase => { + self.execute_streaming_additive(plan, AdditiveOp::Increase) + .await + } + QueryStatistic::QuantileOverTime { phi } => self.execute_quantile(plan, *phi).await, + QueryStatistic::TopK { k } => self.execute_topk(plan, *k).await, + } + } + + /// Streaming additive path. Reads chunks one at a time, applies + /// the per-sample fold, drops the decoded chunk before fetching + /// the next one. Memory is O(1) per query, regardless of how + /// many samples the time range covers. + pub async fn execute_streaming_additive( + &self, + plan: &QueryPlan, + op: AdditiveOp, + ) -> Result { + let (start_ms, end_ms) = plan.time_range_ms; + let chunks = self + .cold_store + .list_chunks(&plan.metric, start_ms, end_ms) + .await?; + let chunks_fetched = chunks.len(); + debug!( + metric = plan.metric.as_str(), + chunks = chunks_fetched, + op = ?op, + "gorilla-engine: streaming additive over chunks" + ); + + let mut acc = AdditiveAccumulator::new(op); + let mut samples_scanned: usize = 0; + for chunk in chunks { + let samples = self.cold_store.read_chunk(&chunk).await?; + for s in samples { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + acc.observe(&s); + samples_scanned += 1; + } + } + } + + let value = acc.finalize(plan, op); + Ok(ExecutionOutcome { + value, + samples_scanned, + chunks_fetched, + }) + } + + /// Buffered quantile path. Materialises every in-range sample + /// up to [`GorillaEngineConfig::max_buffered_samples`], sorts + /// the value column, and picks the φ-rank using a + /// nearest-rank rule (matches Prometheus's + /// `quantile_over_time` semantics for the linear-interp-free + /// midpoint case — the float index rounds to nearest). + pub async fn execute_quantile( + &self, + plan: &QueryPlan, + phi: f64, + ) -> Result { + let (samples, chunks_fetched) = self.collect_buffered_samples(plan).await?; + if samples.is_empty() { + return Ok(ExecutionOutcome { + value: f64::NAN, + samples_scanned: 0, + chunks_fetched, + }); + } + + let mut values: Vec = samples.iter().map(|s| s.value).collect(); + values.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + + let n = values.len() as f64; + let phi = phi.clamp(0.0, 1.0); + let raw_idx = ((n - 1.0) * phi).round() as i64; + let idx = raw_idx.clamp(0, values.len() as i64 - 1) as usize; + Ok(ExecutionOutcome { + value: values[idx], + samples_scanned: values.len(), + chunks_fetched, + }) + } + + /// Buffered top-k path. Materialises every in-range sample, + /// sorts the value column descending, and returns the SUM of + /// the top-`k` values. The Phase-5 capability router will + /// extend this to per-group top-k once spatial grouping + /// lands; the MVP scalar return value matches the existing + /// `ExecutionOutcome` shape. + pub async fn execute_topk( + &self, + plan: &QueryPlan, + k: usize, + ) -> Result { + let (samples, chunks_fetched) = self.collect_buffered_samples(plan).await?; + if samples.is_empty() { + return Ok(ExecutionOutcome { + value: f64::NAN, + samples_scanned: 0, + chunks_fetched, + }); + } + + let mut values: Vec = samples.iter().map(|s| s.value).collect(); + values.sort_by(|a, b| b.partial_cmp(a).unwrap_or(std::cmp::Ordering::Equal)); + let take = k.min(values.len()); + let topk_sum: f64 = values.iter().take(take).sum(); + Ok(ExecutionOutcome { + value: topk_sum, + samples_scanned: values.len(), + chunks_fetched, + }) + } + + /// Shared helper for the buffered paths: walk every chunk, + /// keep in-range samples, enforce the + /// [`GorillaEngineConfig::max_buffered_samples`] ceiling. + async fn collect_buffered_samples( + &self, + plan: &QueryPlan, + ) -> Result<(Vec, usize), EngineError> { + let (start_ms, end_ms) = plan.time_range_ms; + let chunks = self + .cold_store + .list_chunks(&plan.metric, start_ms, end_ms) + .await?; + let chunks_fetched = chunks.len(); + let mut buffer: Vec = Vec::new(); + let limit = self.config.max_buffered_samples; + for chunk in chunks { + let samples = self.cold_store.read_chunk(&chunk).await?; + for s in samples { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + if buffer.len() >= limit { + // Surface the over-budget sample count + // (limit + 1) so callers can pin + // `count > limit` in tests; we don't + // bother walking the rest of the chunks + // just to tighten the number. + return Err(EngineError::TooManySamples { + count: buffer.len() + 1, + limit, + }); + } + buffer.push(s); + } + } + } + Ok((buffer, chunks_fetched)) + } +} + +/// Per-sample fold for the streaming-additive path. Fields are +/// kept in raw f64 (sum / min / max) + i64 (count) so the +/// finaliser can pick the right arithmetic per op. +/// +/// `op` is taken at `new` time rather than carried as a struct +/// field — the finaliser receives it as an argument, keeping the +/// struct itself op-agnostic and shrinking the per-fold footprint. +#[derive(Debug, Clone, Copy)] +struct AdditiveAccumulator { + sum: f64, + count: i64, + min: f64, + max: f64, + /// Earliest observed `(ts_ms, value)` — used by Rate / Increase. + first: Option<(i64, f64)>, + /// Latest observed `(ts_ms, value)` — used by Rate / Increase. + last: Option<(i64, f64)>, +} + +impl AdditiveAccumulator { + fn new(_op: AdditiveOp) -> Self { + Self { + sum: 0.0, + count: 0, + min: f64::INFINITY, + max: f64::NEG_INFINITY, + first: None, + last: None, + } + } + + fn observe(&mut self, s: &RawSample) { + self.sum += s.value; + self.count += 1; + if s.value < self.min { + self.min = s.value; + } + if s.value > self.max { + self.max = s.value; + } + match self.first { + None => self.first = Some((s.ts_ms, s.value)), + Some((ts, _)) if s.ts_ms < ts => self.first = Some((s.ts_ms, s.value)), + _ => {} + } + match self.last { + None => self.last = Some((s.ts_ms, s.value)), + Some((ts, _)) if s.ts_ms > ts => self.last = Some((s.ts_ms, s.value)), + _ => {} + } + } + + /// Convert the running accumulator into a final scalar. + /// Returns `NaN` for the empty-time-range case so downstream + /// formatting stays consistent. + fn finalize(&self, plan: &QueryPlan, op: AdditiveOp) -> f64 { + if self.count == 0 { + return match op { + AdditiveOp::Count => 0.0, + _ => f64::NAN, + }; + } + match op { + AdditiveOp::Sum => self.sum, + AdditiveOp::Count => self.count as f64, + AdditiveOp::Avg => self.sum / self.count as f64, + AdditiveOp::Min => self.min, + AdditiveOp::Max => self.max, + AdditiveOp::Increase => match (self.first, self.last) { + (Some((_, fv)), Some((_, lv))) => lv - fv, + _ => f64::NAN, + }, + AdditiveOp::Rate => match (self.first, self.last) { + (Some((_, fv)), Some((_, lv))) => { + let (start_ms, end_ms) = plan.time_range_ms; + let range_secs = ((end_ms - start_ms).max(1)) as f64 / 1000.0; + if range_secs <= 0.0 { + f64::NAN + } else { + (lv - fv) / range_secs + } + } + _ => f64::NAN, + }, + } + } +} + diff --git a/asap-query-engine/src/engines/gorilla_engine/mod.rs b/asap-query-engine/src/engines/gorilla_engine/mod.rs new file mode 100644 index 000000000..bc2c3ce2e --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/mod.rs @@ -0,0 +1,266 @@ +//! Phase 4: `GorillaQueryEngine` — exact PromQL execution over the +//! Gorilla-S3 cold tier. +//! +//! This engine is a SIBLING of [`crate::engines::simple_engine::SimpleEngine`]. +//! Both consume the same PromQL surface, but where `SimpleEngine` +//! answers from warm-tier sketches (approximate, ε/δ-bounded), the +//! `GorillaQueryEngine` answers exactly from per-hour Gorilla +//! chunks landed on S3 / MinIO via the Phase-3 +//! [`crate::drivers::query::fallback::cold_store::gorilla_s3::GorillaS3ColdStore`]. +//! +//! Result wrapping pins three things: +//! +//! 1. an [`crate::stores::sketch_db::AccuracyEnvelope`] with +//! `kind = Exact`, ε = 0, δ = 0, +//! 2. a `data_source: gorilla_archive` info line, +//! 3. cheap diagnostics (`samples_scanned`, `chunks_fetched`). +//! +//! See `docs/design-gorilla-s3-cold-engine.md` §6. +//! +//! ## Two execution strategies +//! +//! Per-statistic dispatch in [`exact_executor`]: +//! +//! * **Streaming-additive** — `Sum`, `Count`, `Min`, `Max`, `Rate`, +//! `Increase` (and `Avg` derived as Sum/Count). One chunk at a +//! time, fold into a small accumulator, drop the decoded samples +//! before fetching the next chunk. Memory cost is O(1) per group. +//! * **Buffered** — `Quantile`, `TopK`, `Cardinality`. Materialise +//! every in-range sample, then sort / count. Bounded by +//! [`GorillaEngineConfig::max_buffered_samples`]; over-budget +//! queries fail fast with [`EngineError::TooManySamples`]. + +pub mod exact_executor; +pub mod query_planner; + +#[cfg(test)] +mod tests; + +use std::sync::Arc; +use std::time::Duration; + +use thiserror::Error; +use tokio::time::error::Elapsed; +use tracing::debug; + +use crate::data_model::KeyByLabelValues; +use crate::drivers::query::fallback::cold_store::{ColdStore, ColdStoreError}; +use crate::engines::query_result::{InstantVectorElement, QueryResult}; +use crate::stores::sketch_db::accuracy::{AccuracyEnvelope, AccuracyProfile}; + +pub use exact_executor::{AdditiveOp, ExactExecutor}; +pub use query_planner::{plan_query, plan_query_at, QueryPlan, QueryStatistic}; + +/// Marker line that every `GorillaQueryEngine` answer carries on +/// its `infos` array. Pinned so dashboards / Phase-5 capability +/// routers can byte-compare without parsing. +pub const DATA_SOURCE_GORILLA_ARCHIVE: &str = "data_source: gorilla_archive"; + +/// Tunable runtime knobs for the Gorilla query engine. +/// +/// Call sites typically construct via `Default::default()`; tests +/// override `max_buffered_samples` to exercise the bounded-buffer +/// guard. +#[derive(Debug, Clone)] +pub struct GorillaEngineConfig { + /// Hard cap on the number of samples a buffered-aggregate + /// query (quantile / topk / cardinality) is allowed to + /// materialise in memory. Default `10_000_000` + /// (~160 MB at 16 B per `(ts, value)` pair). + pub max_buffered_samples: usize, + /// Wall-clock query timeout, in seconds. Default `30`. + pub query_timeout_secs: u64, +} + +impl Default for GorillaEngineConfig { + fn default() -> Self { + Self { + max_buffered_samples: 10_000_000, + query_timeout_secs: 30, + } + } +} + +/// Error surface returned by [`GorillaQueryEngine::execute`]. +#[derive(Debug, Error)] +pub enum EngineError { + /// PromQL string failed to parse, or used a construct outside + /// the engine's supported surface (see [`query_planner`]). + #[error("query planning failed: {0}")] + Plan(String), + /// Cold-store fetch / decode failed. + #[error("cold-store error: {0}")] + ColdStore(#[from] ColdStoreError), + /// Buffered-aggregate budget exceeded — query asked for more + /// samples than [`GorillaEngineConfig::max_buffered_samples`] + /// will allow. The user should narrow the time range or + /// lower the cardinality. + #[error( + "buffered-aggregate budget exceeded: {count} samples > limit {limit}; \ + narrow the time range or lower the metric cardinality" + )] + TooManySamples { + /// Samples the engine attempted to materialise. + count: usize, + /// Configured ceiling. + limit: usize, + }, + /// Wall-clock timeout fired before the query finished. + #[error("query timed out after {0:?}")] + Timeout(Duration), +} + +impl From for EngineError { + fn from(_: Elapsed) -> Self { + Self::Timeout(Duration::from_secs(0)) + } +} + +/// Phase-4 cold-tier exact engine. +/// +/// Holds an `Arc` rather than a concrete +/// `Arc` so tests can inject in-memory mocks +/// and so future cold backends (local-FS chunks, multi-region +/// fan-out) drop in without changing the engine surface. The +/// production constructor [`GorillaQueryEngine::with_gorilla_s3`] +/// keeps the design.md type signature working at the call site. +pub struct GorillaQueryEngine { + cold_store: Arc, + config: GorillaEngineConfig, +} + +impl GorillaQueryEngine { + /// Build with an arbitrary cold-store implementation. Used by + /// tests + the Phase-5 capability router (which may swap the + /// concrete impl based on routing decisions). + pub fn new(cold_store: Arc, config: GorillaEngineConfig) -> Self { + Self { + cold_store, + config, + } + } + + /// Convenience constructor for the production + /// [`crate::drivers::query::fallback::cold_store::gorilla_s3::GorillaS3ColdStore`] + /// path. Mirrors the design.md type signature. + pub fn with_gorilla_s3( + cold_store: Arc, + config: GorillaEngineConfig, + ) -> Self { + Self::new(cold_store as Arc, config) + } + + /// Read-only access to the configured limits — useful for + /// diagnostics + the Phase-5 router's cost estimator. + pub fn config(&self) -> &GorillaEngineConfig { + &self.config + } + + /// Execute a parsed PromQL query against the cold tier. + /// + /// The query string is parsed via [`query_planner::plan_query`], + /// the resulting plan dispatches to either the streaming + /// additive or the buffered execution path, and the answer is + /// wrapped with the exact-accuracy envelope + the + /// `data_source: gorilla_archive` annotation. + pub async fn execute(&self, query: &str) -> Result { + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0); + self.execute_at(query, now_ms).await + } + + /// Like [`Self::execute`], with a caller-supplied `now_ms` + /// pinning the right edge of the request window. Used by + /// tests + by the (future) Phase-5 router that wants to back- + /// date a query against historical chunks. + pub async fn execute_at( + &self, + query: &str, + now_ms: i64, + ) -> Result { + let timeout = Duration::from_secs(self.config.query_timeout_secs.max(1)); + tokio::time::timeout(timeout, self.execute_inner(query, now_ms)) + .await + .map_err(|_| EngineError::Timeout(timeout))? + } + + async fn execute_inner( + &self, + query: &str, + now_ms: i64, + ) -> Result { + let plan = query_planner::plan_query_at(query, now_ms).map_err(EngineError::Plan)?; + debug!( + metric = plan.metric.as_str(), + stat = ?plan.statistic, + start_ms = plan.time_range_ms.0, + end_ms = plan.time_range_ms.1, + "gorilla-engine: executing plan" + ); + + let executor = ExactExecutor::new(self.cold_store.clone(), self.config.clone()); + let outcome = executor.execute_plan(&plan).await?; + + Ok(wrap_result(&plan, outcome)) + } +} + +/// Wrap a finished `(scalar value, sample / chunk counts)` into a +/// `QueryResult` with the exact-accuracy envelope + the +/// `data_source: gorilla_archive` info line. Pulled out so tests +/// can pin the wrapping shape independently of the executor. +pub fn wrap_result(plan: &QueryPlan, outcome: ExecutionOutcome) -> QueryResult { + // Result timestamp is the right edge of the requested range — + // mirrors `SimpleEngine`'s convention for instant-vector queries + // against a closed time window. + let result_ts = plan.time_range_ms.1.max(0) as u64; + + let labels = KeyByLabelValues::new_with_labels(Vec::new()); + let element = InstantVectorElement::new(labels, outcome.value); + let envelope = AccuracyEnvelope::single(AccuracyProfile::exact()); + QueryResult::vector(vec![element], result_ts).with_accuracy(envelope) + // Window is the requested range, expressed in u64 ms. + .with_window_used(( + plan.time_range_ms.0.max(0) as u64, + plan.time_range_ms.1.max(0) as u64, + )) +} + +/// Output of an executed plan. Kept narrow on purpose — Phase 4's +/// MVP returns a single scalar per query. Higher-cardinality +/// (per-group) shapes will land in Phase 5+ once capability +/// routing decides which engine answers grouped queries. +#[derive(Debug, Clone, PartialEq)] +pub struct ExecutionOutcome { + /// Final scalar (e.g. `sum_over_time` total, `quantile_over_time` + /// φ-quantile). NaN when the time range carries no samples. + pub value: f64, + /// Number of raw samples that contributed to `value`. + pub samples_scanned: usize, + /// Number of chunks the executor fetched from the cold store. + pub chunks_fetched: usize, +} + +impl ExecutionOutcome { + /// "no data" sentinel — used when the time range is empty. + pub fn empty() -> Self { + Self { + value: f64::NAN, + samples_scanned: 0, + chunks_fetched: 0, + } + } + + /// Build the `infos` array surfaced on the wire response. + /// Pulled out so tests can pin the exact strings. + pub fn info_lines(&self) -> Vec { + vec![ + AccuracyProfile::exact().summary(), + DATA_SOURCE_GORILLA_ARCHIVE.to_string(), + format!("samples_scanned: {}", self.samples_scanned), + format!("chunks_fetched: {}", self.chunks_fetched), + ] + } +} diff --git a/asap-query-engine/src/engines/gorilla_engine/query_planner.rs b/asap-query-engine/src/engines/gorilla_engine/query_planner.rs new file mode 100644 index 000000000..6334f6fd1 --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/query_planner.rs @@ -0,0 +1,307 @@ +//! PromQL → [`QueryPlan`] translator for the Phase-4 Gorilla engine. +//! +//! The Phase-4 surface is intentionally narrow: instant-vector +//! queries that wrap a single matrix selector with one of the +//! supported `*_over_time` / `rate` / `increase` functions, OR +//! a top-level `quantile_over_time(φ, m[range])` / +//! `topk(k, m[range])`-style aggregation. +//! +//! Time range is `(now - lookback_ms, now)` where `now` is the +//! caller-supplied "query time" — fixed to `chrono::Utc::now()` +//! when not specified, so callers that don't care about backdating +//! a query don't need to thread a clock through. + +use std::time::SystemTime; + +use chrono::Utc; +use promql_parser::parser::{ + AggregateExpr, Call, Expr, FunctionArgs, MatrixSelector, NumberLiteral, ParenExpr, + VectorSelector, +}; + +/// Statistic to compute, alongside any extra parameters +/// (quantile φ, top-k k). +#[derive(Debug, Clone, PartialEq)] +pub enum QueryStatistic { + /// `sum_over_time(m[range])` + SumOverTime, + /// `count_over_time(m[range])` + CountOverTime, + /// `avg_over_time(m[range])` (= sum / count) + AvgOverTime, + /// `min_over_time(m[range])` + MinOverTime, + /// `max_over_time(m[range])` + MaxOverTime, + /// `rate(m[range])` — `(last - first) / range_seconds` + Rate, + /// `increase(m[range])` — `last - first` + Increase, + /// `quantile_over_time(φ, m[range])` + QuantileOverTime { phi: f64 }, + /// `topk(k, sum_over_time(m[range]))`-style aggregation. The + /// MVP Phase 4 implementation returns the sum of the top-`k` + /// sample values in the range — once Phase 5 adds spatial + /// grouping the executor will return a per-group vector. + TopK { k: usize }, +} + +impl QueryStatistic { + /// True iff the executor can answer this statistic via the + /// streaming-additive path; false → buffered path (everything + /// has to be in memory before producing the answer). + pub fn is_streaming_additive(&self) -> bool { + matches!( + self, + Self::SumOverTime + | Self::CountOverTime + | Self::AvgOverTime + | Self::MinOverTime + | Self::MaxOverTime + | Self::Rate + | Self::Increase + ) + } +} + +/// Output of [`plan_query`]. +#[derive(Debug, Clone, PartialEq)] +pub struct QueryPlan { + pub metric: String, + /// Half-open `[start_ms, end_ms)` request window. Computed as + /// `(now_ms - range_ms, now_ms)` from the matrix selector's + /// `[range]` duration. + pub time_range_ms: (i64, i64), + pub statistic: QueryStatistic, +} + +/// Parse `query` and produce a [`QueryPlan`]. `now` defaults to +/// the system clock; the [`plan_query_at`] variant lets tests pin +/// a deterministic timestamp. +pub fn plan_query(query: &str) -> Result { + let now_ms = SystemTime::now() + .duration_since(SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or_else(|_| Utc::now().timestamp_millis()); + plan_query_at(query, now_ms) +} + +/// As [`plan_query`], with a caller-supplied `now_ms`. +pub fn plan_query_at(query: &str, now_ms: i64) -> Result { + let ast = promql_parser::parser::parse(query).map_err(|e| format!("parse: {e}"))?; + plan_from_ast(&ast, now_ms) +} + +fn plan_from_ast(ast: &Expr, now_ms: i64) -> Result { + match ast { + Expr::Paren(ParenExpr { expr }) => plan_from_ast(expr, now_ms), + Expr::Call(call) => plan_from_call(call, now_ms), + Expr::Aggregate(agg) => plan_from_aggregate(agg, now_ms), + other => Err(format!( + "unsupported top-level expression: {:?}; the Gorilla engine \ + expects a single function call (rate/increase/*_over_time) \ + or topk(k, ...) aggregation", + std::mem::discriminant(other) + )), + } +} + +fn plan_from_call(call: &Call, now_ms: i64) -> Result { + let name = call.func.name.to_lowercase(); + match name.as_str() { + "sum_over_time" + | "count_over_time" + | "avg_over_time" + | "min_over_time" + | "max_over_time" + | "rate" + | "increase" => { + let ms = expect_single_matrix_arg(&call.args, &name)?; + let (metric, range_ms) = matrix_metric_and_range_ms(ms); + let stat = match name.as_str() { + "sum_over_time" => QueryStatistic::SumOverTime, + "count_over_time" => QueryStatistic::CountOverTime, + "avg_over_time" => QueryStatistic::AvgOverTime, + "min_over_time" => QueryStatistic::MinOverTime, + "max_over_time" => QueryStatistic::MaxOverTime, + "rate" => QueryStatistic::Rate, + "increase" => QueryStatistic::Increase, + _ => unreachable!(), + }; + Ok(QueryPlan { + metric, + time_range_ms: (now_ms - range_ms, now_ms), + statistic: stat, + }) + } + "quantile_over_time" => { + // quantile_over_time(φ, m[range]) + if call.args.args.len() != 2 { + return Err(format!( + "quantile_over_time expects 2 args, got {}", + call.args.args.len() + )); + } + let phi = expect_number(&call.args.args[0], "quantile_over_time φ")?; + let ms = expect_matrix_selector(&call.args.args[1], "quantile_over_time")?; + let (metric, range_ms) = matrix_metric_and_range_ms(ms); + Ok(QueryPlan { + metric, + time_range_ms: (now_ms - range_ms, now_ms), + statistic: QueryStatistic::QuantileOverTime { phi }, + }) + } + other => Err(format!( + "unsupported PromQL function: {other}; the Gorilla engine \ + supports rate/increase/*_over_time/quantile_over_time" + )), + } +} + +fn plan_from_aggregate(agg: &AggregateExpr, now_ms: i64) -> Result { + // PromQL grammar requires aggregation operators to take a + // vector — so the legal Phase-4 spellings are e.g. + // `topk(2, sum_over_time(m[10s]))`. We strip the outer + // aggregation, recurse into the inner call to recover the + // `(metric, range)` pair, then overlay the TopK statistic. + let op_str = format!("{}", agg.op); + if !op_str.eq_ignore_ascii_case("topk") { + return Err(format!( + "unsupported top-level aggregation: {op_str}; only `topk(k, ...)` \ + is supported in Phase 4" + )); + } + let k_expr = agg + .param + .as_deref() + .ok_or_else(|| "topk requires a numeric parameter (k)".to_string())?; + let k = expect_number(k_expr, "topk k")?; + if !k.is_finite() || k <= 0.0 { + return Err(format!("topk k must be positive, got {k}")); + } + // Recurse into the inner expression — it can be a matrix + // selector (handled by [`matrix_metric_and_range_ms`] + // directly) OR a vector-returning function call (the legal + // PromQL spelling). Either way we end up with a + // `(metric, range_ms)` pair we can overlay TopK on. + let inner_plan = match &*agg.expr { + Expr::MatrixSelector(ms) => { + let (metric, range_ms) = matrix_metric_and_range_ms(ms); + QueryPlan { + metric, + time_range_ms: (now_ms - range_ms, now_ms), + statistic: QueryStatistic::SumOverTime, // overlay below + } + } + _ => plan_from_ast(&agg.expr, now_ms)?, + }; + Ok(QueryPlan { + metric: inner_plan.metric, + time_range_ms: inner_plan.time_range_ms, + statistic: QueryStatistic::TopK { k: k as usize }, + }) +} + +fn expect_single_matrix_arg<'a>( + args: &'a FunctionArgs, + fname: &str, +) -> Result<&'a MatrixSelector, String> { + if args.args.len() != 1 { + return Err(format!( + "{fname} expects 1 matrix-selector arg, got {}", + args.args.len() + )); + } + expect_matrix_selector(&args.args[0], fname) +} + +fn expect_matrix_selector<'a>(expr: &'a Expr, ctx: &str) -> Result<&'a MatrixSelector, String> { + match expr { + Expr::MatrixSelector(ms) => Ok(ms), + Expr::Paren(ParenExpr { expr }) => expect_matrix_selector(expr, ctx), + other => Err(format!( + "{ctx}: expected matrix selector `metric[range]`, got {:?}", + std::mem::discriminant(other) + )), + } +} + +fn expect_number(expr: &Expr, ctx: &str) -> Result { + match expr { + Expr::NumberLiteral(NumberLiteral { val }) => Ok(*val), + Expr::Paren(ParenExpr { expr }) => expect_number(expr, ctx), + other => Err(format!( + "{ctx}: expected numeric literal, got {:?}", + std::mem::discriminant(other) + )), + } +} + +fn matrix_metric_and_range_ms(ms: &MatrixSelector) -> (String, i64) { + let metric = vector_selector_metric(&ms.vs); + let range_ms = ms.range.as_millis() as i64; + (metric, range_ms) +} + +fn vector_selector_metric(vs: &VectorSelector) -> String { + if let Some(name) = &vs.name { + return name.clone(); + } + // Fallback: inspect matchers for an `__name__` exact match. + for m in vs.matchers.matchers.iter() { + if m.name == "__name__" { + return m.value.clone(); + } + } + String::new() +} + +#[cfg(test)] +mod tests { + use super::*; + + const NOW: i64 = 1_715_000_000_000; + + #[test] + fn plans_sum_over_time() { + let plan = plan_query_at("sum_over_time(http_requests_total[5m])", NOW).unwrap(); + assert_eq!(plan.metric, "http_requests_total"); + assert_eq!(plan.statistic, QueryStatistic::SumOverTime); + assert_eq!(plan.time_range_ms, (NOW - 5 * 60_000, NOW)); + } + + #[test] + fn plans_quantile_over_time() { + let plan = plan_query_at("quantile_over_time(0.99, latency_ms[1m])", NOW).unwrap(); + assert_eq!(plan.metric, "latency_ms"); + assert!(matches!( + plan.statistic, + QueryStatistic::QuantileOverTime { phi } if (phi - 0.99).abs() < 1e-12 + )); + } + + #[test] + fn plans_topk() { + // Legal PromQL spelling: aggregation wraps a vector-returning + // function call. The Phase-4 planner peels off the outer + // `topk` and recovers the `(metric, range)` pair from the + // inner `sum_over_time(...)`. + let plan = plan_query_at("topk(3, sum_over_time(m[10s]))", NOW).unwrap(); + assert!(matches!(plan.statistic, QueryStatistic::TopK { k } if k == 3)); + assert_eq!(plan.metric, "m"); + assert_eq!(plan.time_range_ms, (NOW - 10_000, NOW)); + } + + #[test] + fn rejects_binary_expression() { + assert!(plan_query_at("foo + bar", NOW).is_err()); + } + + #[test] + fn streaming_classification() { + assert!(QueryStatistic::SumOverTime.is_streaming_additive()); + assert!(QueryStatistic::Rate.is_streaming_additive()); + assert!(!QueryStatistic::QuantileOverTime { phi: 0.5 }.is_streaming_additive()); + assert!(!QueryStatistic::TopK { k: 1 }.is_streaming_additive()); + } +} diff --git a/asap-query-engine/src/engines/gorilla_engine/tests.rs b/asap-query-engine/src/engines/gorilla_engine/tests.rs new file mode 100644 index 000000000..70b68d151 --- /dev/null +++ b/asap-query-engine/src/engines/gorilla_engine/tests.rs @@ -0,0 +1,615 @@ +//! Phase-4 unit tests for the Gorilla query engine. +//! +//! Tests exercise the engine end-to-end via a `MockColdStore` +//! injected in place of the production `GorillaS3ColdStore`. The +//! mock is intentionally minimal: it owns a `Vec<(ChunkRef, +//! Vec)>` and answers `list_chunks` / `read_chunk` +//! straight off it, with optional latency injection for the +//! timeout test. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use tokio::time::sleep; + +use crate::drivers::query::fallback::cold_store::{ + ChunkRef, ColdStore, ColdStoreError, RawSample, +}; +use crate::engines::query_result::QueryResult; +use crate::stores::sketch_db::accuracy::{AccuracyKind, AccuracyProfile}; + +use super::query_planner::{plan_query_at, QueryStatistic}; +use super::{ + wrap_result, EngineError, ExactExecutor, ExecutionOutcome, GorillaEngineConfig, + GorillaQueryEngine, DATA_SOURCE_GORILLA_ARCHIVE, +}; + +// ───────────────────────────────────────────────────────────────────── +// Mock cold store +// ───────────────────────────────────────────────────────────────────── + +/// In-process mock that satisfies the [`ColdStore`] trait without +/// any S3 / disk roundtrip. Built once in each test from a list of +/// `(ChunkRef, samples)` pairs. +#[derive(Default)] +struct MockColdStore { + chunks: Vec<(ChunkRef, Vec)>, + /// If set, every `read_chunk` call sleeps for this duration — + /// used by the timeout test. + read_delay: Option, +} + +impl MockColdStore { + fn new(chunks: Vec<(ChunkRef, Vec)>) -> Self { + Self { + chunks, + read_delay: None, + } + } + + fn with_read_delay(mut self, d: Duration) -> Self { + self.read_delay = Some(d); + self + } +} + +#[async_trait] +impl ColdStore for MockColdStore { + async fn scan( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + let mut out = Vec::new(); + let chunks = self.list_chunks(metric, start_ms, end_ms).await?; + for c in chunks { + for s in self.read_chunk(&c).await? { + if s.ts_ms >= start_ms && s.ts_ms < end_ms { + out.push(s); + } + } + } + Ok(out) + } + + async fn list_chunks( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + Ok(self + .chunks + .iter() + .filter(|(c, _)| { + c.metric == metric + && c.time_range_ms.0 < end_ms + && c.time_range_ms.1 >= start_ms + }) + .map(|(c, _)| c.clone()) + .collect()) + } + + async fn read_chunk(&self, chunk: &ChunkRef) -> Result, ColdStoreError> { + if let Some(d) = self.read_delay { + sleep(d).await; + } + for (c, samples) in &self.chunks { + if c.key == chunk.key { + return Ok(samples.clone()); + } + } + Err(ColdStoreError::Backend(format!( + "mock: no such chunk {}", + chunk.key + ))) + } +} + +// ───────────────────────────────────────────────────────────────────── +// Fixture helpers +// ───────────────────────────────────────────────────────────────────── + +const NOW_MS: i64 = 1_715_000_000_000; +const METRIC: &str = "http_requests_total"; + +fn raw(ts_ms: i64, value: f64) -> RawSample { + RawSample { + ts_ms, + labels: BTreeMap::new(), + value, + } +} + +/// One chunk covering `[start, start + n*step]` with a +/// monotonically-increasing value column (`base + i*step_v`). +fn linear_chunk( + key: &str, + start_ms: i64, + step_ms: i64, + n: usize, + base: f64, + step_v: f64, +) -> (ChunkRef, Vec) { + let samples: Vec = (0..n) + .map(|i| raw(start_ms + (i as i64) * step_ms, base + (i as f64) * step_v)) + .collect(); + let last_ts = samples.last().map(|s| s.ts_ms).unwrap_or(start_ms); + let chunk = ChunkRef { + key: key.to_string(), + metric: METRIC.to_string(), + time_range_ms: (start_ms, last_ts + 1), + label_hash: 0, + sample_count: n as u32, + size_bytes: 0, + }; + (chunk, samples) +} + +fn cfg() -> GorillaEngineConfig { + GorillaEngineConfig { + max_buffered_samples: 1_000_000, + query_timeout_secs: 30, + } +} + +fn engine_with(chunks: Vec<(ChunkRef, Vec)>) -> GorillaQueryEngine { + GorillaQueryEngine::new(Arc::new(MockColdStore::new(chunks)), cfg()) +} + +fn engine_with_config( + chunks: Vec<(ChunkRef, Vec)>, + config: GorillaEngineConfig, +) -> GorillaQueryEngine { + GorillaQueryEngine::new(Arc::new(MockColdStore::new(chunks)), config) +} + +// ───────────────────────────────────────────────────────────────────── +// Streaming-additive happy paths +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn execute_sum_over_time_streaming() { + // 60 samples × value 2.0 = 120.0 + let chunks = vec![linear_chunk( + "c1", + NOW_MS - 60_000, + 1_000, + 60, + 2.0, + 0.0, + )]; + let engine = engine_with(chunks); + let plan = plan_query_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS).unwrap(); + assert_eq!(plan.statistic, QueryStatistic::SumOverTime); + + let exec = ExactExecutor::new( + Arc::new(MockColdStore::new(vec![linear_chunk( + "c1", + NOW_MS - 60_000, + 1_000, + 60, + 2.0, + 0.0, + )])), + cfg(), + ); + let outcome = exec.execute_plan(&plan).await.unwrap(); + assert_eq!(outcome.value, 120.0); + assert_eq!(outcome.samples_scanned, 60); + assert_eq!(outcome.chunks_fetched, 1); + + // Also verify via the high-level engine. + let result = engine + .execute_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS) + .await + .unwrap(); + assert!(matches!(result, QueryResult::Vector(_))); +} + +#[tokio::test] +async fn execute_count_over_time() { + let chunks = vec![linear_chunk("c1", NOW_MS - 30_000, 1_000, 30, 0.0, 0.0)]; + let engine = engine_with(chunks); + let plan = plan_query_at(&format!("count_over_time({METRIC}[1m])"), NOW_MS).unwrap(); + let exec = ExactExecutor::new( + Arc::new(MockColdStore::new(vec![linear_chunk( + "c1", + NOW_MS - 30_000, + 1_000, + 30, + 0.0, + 0.0, + )])), + cfg(), + ); + let outcome = exec.execute_plan(&plan).await.unwrap(); + assert_eq!(outcome.value, 30.0); + + let result = engine + .execute_at(&format!("count_over_time({METRIC}[1m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 30.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_avg_over_time() { + // 4 samples: 1, 2, 3, 4 → avg = 2.5 + let chunks = vec![linear_chunk("c1", NOW_MS - 4_000, 1_000, 4, 1.0, 1.0)]; + let engine = engine_with(chunks); + let result = engine + .execute_at(&format!("avg_over_time({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert!((iv.values[0].value - 2.5).abs() < 1e-12); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_min_over_time() { + // values 5, 1, 3, 4 → min = 1 + let samples = vec![ + raw(NOW_MS - 4_000, 5.0), + raw(NOW_MS - 3_000, 1.0), + raw(NOW_MS - 2_000, 3.0), + raw(NOW_MS - 1_000, 4.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 4_000, NOW_MS), + label_hash: 0, + sample_count: 4, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("min_over_time({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 1.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_max_over_time() { + let samples = vec![ + raw(NOW_MS - 4_000, 5.0), + raw(NOW_MS - 3_000, 1.0), + raw(NOW_MS - 2_000, 3.0), + raw(NOW_MS - 1_000, 4.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 4_000, NOW_MS), + label_hash: 0, + sample_count: 4, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("max_over_time({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 5.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_rate_basic() { + // Counter goes from 100 at t=NOW-10s to 200 at t=NOW-1s. + // rate over 10s window = (200 - 100) / 10s = 10.0 + let samples = vec![ + raw(NOW_MS - 10_000, 100.0), + raw(NOW_MS - 5_000, 150.0), + raw(NOW_MS - 1_000, 200.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 10_000, NOW_MS), + label_hash: 0, + sample_count: 3, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("rate({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert!((iv.values[0].value - 10.0).abs() < 1e-9); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_increase_basic() { + let samples = vec![ + raw(NOW_MS - 10_000, 100.0), + raw(NOW_MS - 5_000, 150.0), + raw(NOW_MS - 1_000, 250.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 10_000, NOW_MS), + label_hash: 0, + sample_count: 3, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("increase({METRIC}[10s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert!((iv.values[0].value - 150.0).abs() < 1e-9); + } else { + panic!("expected Vector"); + } +} + +// ───────────────────────────────────────────────────────────────────── +// Buffered quantile + topk +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn execute_quantile_buffered_basic() { + // Values 0..100; q0.99 → index round((100-1)*0.99) = round(98.01) = 98 → value 98. + let mut samples: Vec = (0..100) + .map(|i| raw(NOW_MS - 100_000 + (i as i64) * 1_000, i as f64)) + .collect(); + // Shuffle the value order so the executor must sort. + samples.sort_by_key(|s| s.value as i64); + samples.reverse(); + + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 100_000, NOW_MS), + label_hash: 0, + sample_count: 100, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("quantile_over_time(0.99, {METRIC}[2m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 98.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_quantile_too_many_samples_errors() { + // Generate 100 samples but cap the buffered budget at 5. + let samples: Vec = (0..100) + .map(|i| raw(NOW_MS - 100_000 + (i as i64) * 1_000, i as f64)) + .collect(); + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 100_000, NOW_MS), + label_hash: 0, + sample_count: 100, + size_bytes: 0, + }; + let cfg = GorillaEngineConfig { + max_buffered_samples: 5, + query_timeout_secs: 30, + }; + let engine = engine_with_config(vec![(chunk, samples)], cfg); + let res = engine + .execute_at(&format!("quantile_over_time(0.5, {METRIC}[2m])"), NOW_MS) + .await; + match res { + Err(EngineError::TooManySamples { count, limit }) => { + assert_eq!(limit, 5); + assert!(count > limit); + } + other => panic!("expected TooManySamples, got {other:?}"), + } +} + +#[tokio::test] +async fn execute_topk_basic() { + // Values [1, 2, 3, 10, 20]; topk(2) → 30 + let samples = vec![ + raw(NOW_MS - 5_000, 1.0), + raw(NOW_MS - 4_000, 2.0), + raw(NOW_MS - 3_000, 3.0), + raw(NOW_MS - 2_000, 10.0), + raw(NOW_MS - 1_000, 20.0), + ]; + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 5_000, NOW_MS), + label_hash: 0, + sample_count: 5, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("topk(2, sum_over_time({METRIC}[10s]))"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!(iv.values[0].value, 30.0); + } else { + panic!("expected Vector"); + } +} + +// ───────────────────────────────────────────────────────────────────── +// Edge cases +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn execute_empty_chunks_returns_zero_or_nan() { + let engine = engine_with(Vec::new()); + let sum = engine + .execute_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = sum { + assert!(iv.values[0].value.is_nan(), "sum on empty should be NaN"); + } else { + panic!("expected Vector"); + } + let count = engine + .execute_at(&format!("count_over_time({METRIC}[5m])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = count { + assert_eq!(iv.values[0].value, 0.0); + } else { + panic!("expected Vector"); + } +} + +#[tokio::test] +async fn execute_chunks_partially_outside_range_filtered() { + // Chunk has 100 samples spanning [NOW-100s, NOW]; request + // covers the latter half [NOW-50s, NOW] → exactly 50 samples + // contribute. + let samples: Vec = (0..100) + .map(|i| raw(NOW_MS - 100_000 + (i as i64) * 1_000, 1.0)) + .collect(); + let chunk = ChunkRef { + key: "c".into(), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 100_000, NOW_MS), + label_hash: 0, + sample_count: 100, + size_bytes: 0, + }; + let engine = engine_with(vec![(chunk, samples)]); + let result = engine + .execute_at(&format!("count_over_time({METRIC}[50s])"), NOW_MS) + .await + .unwrap(); + if let QueryResult::Vector(iv) = result { + assert_eq!( + iv.values[0].value, 50.0, + "exactly 50 samples should match a 50s window" + ); + } else { + panic!("expected Vector"); + } +} + +// ───────────────────────────────────────────────────────────────────── +// Result wrapping +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn result_carries_exact_accuracy_envelope() { + let chunks = vec![linear_chunk("c", NOW_MS - 1_000, 100, 10, 1.0, 0.0)]; + let engine = engine_with(chunks); + let result = engine + .execute_at(&format!("sum_over_time({METRIC}[5s])"), NOW_MS) + .await + .unwrap(); + let env = result + .accuracy() + .expect("Gorilla engine result must carry an accuracy envelope"); + assert_eq!(env.profile.kind, AccuracyKind::Exact); + assert_eq!(env.profile.epsilon, 0.0); + assert_eq!(env.profile.delta, 0.0); + // And the summary string the dashboards parse: + assert_eq!(env.profile.summary(), AccuracyProfile::exact().summary()); +} + +#[tokio::test] +async fn result_includes_data_source_gorilla_archive() { + // The wrapping fn surfaces the data_source line on + // ExecutionOutcome::info_lines — pin both the marker constant + // and the assembled info strings. + let outcome = ExecutionOutcome { + value: 42.0, + samples_scanned: 7, + chunks_fetched: 2, + }; + let infos = outcome.info_lines(); + assert!( + infos.contains(&DATA_SOURCE_GORILLA_ARCHIVE.to_string()), + "infos must include `{DATA_SOURCE_GORILLA_ARCHIVE}`; got {infos:?}" + ); + assert!( + infos.iter().any(|i| i == "samples_scanned: 7"), + "infos must report the scanned-samples count" + ); + assert!( + infos.iter().any(|i| i == "chunks_fetched: 2"), + "infos must report the chunk-fetch count" + ); + + // And via the wrap_result path, the QueryResult itself carries + // the exact-accuracy envelope (data_source line is on the + // info-array which is assembled at the HTTP-driver layer; see + // wrap_result docs). + let plan = plan_query_at(&format!("sum_over_time({METRIC}[5s])"), NOW_MS).unwrap(); + let qr = wrap_result(&plan, outcome.clone()); + let env = qr.accuracy().expect("wrap_result must attach envelope"); + assert_eq!(env.profile.kind, AccuracyKind::Exact); +} + +// ───────────────────────────────────────────────────────────────────── +// Timeout +// ───────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn engine_respects_config_timeout() { + // 1 chunk + 250 ms read delay; engine timeout = 1 s ceil. We + // configure the timeout to 1s (the floor) and force the chunk + // count up so the cumulative read time > 1 s. + let mut chunks = Vec::new(); + for i in 0..10 { + let chunk = ChunkRef { + key: format!("k-{i}"), + metric: METRIC.into(), + time_range_ms: (NOW_MS - 60_000, NOW_MS), + label_hash: 0, + sample_count: 1, + size_bytes: 0, + }; + chunks.push((chunk, vec![raw(NOW_MS - 1_000, 1.0)])); + } + let mock = MockColdStore::new(chunks).with_read_delay(Duration::from_millis(250)); + let cfg = GorillaEngineConfig { + max_buffered_samples: 1_000_000, + query_timeout_secs: 1, + }; + let engine = GorillaQueryEngine::new(Arc::new(mock), cfg); + let res = engine + .execute_at(&format!("sum_over_time({METRIC}[5m])"), NOW_MS) + .await; + match res { + Err(EngineError::Timeout(_)) => {} + other => panic!("expected Timeout, got {other:?}"), + } +} diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 06ed0db14..8395ad2a8 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -1,3 +1,4 @@ +pub mod gorilla_engine; pub mod logical; pub mod physical; pub mod query_result; @@ -5,6 +6,9 @@ pub mod simple_engine; pub mod timeline_dispatch; pub mod window_merger; +pub use gorilla_engine::{ + EngineError as GorillaEngineError, GorillaEngineConfig, GorillaQueryEngine, +}; pub use query_result::{InstantVector, QueryResult, RangeVector, RangeVectorElement, Sample}; pub use simple_engine::SimpleEngine; pub use timeline_dispatch::{combine_statistic, CombinedResult};