diff --git a/asap-query-engine/src/drivers/query/servers/http.rs b/asap-query-engine/src/drivers/query/servers/http.rs index 94ce7fd7..2e3806b7 100644 --- a/asap-query-engine/src/drivers/query/servers/http.rs +++ b/asap-query-engine/src/drivers/query/servers/http.rs @@ -14,8 +14,15 @@ use std::time::Instant; use tokio::net::TcpListener; use tracing::{debug, info, warn}; -use crate::drivers::query::adapters::{create_http_adapter, AdapterConfig, HttpProtocolAdapter}; -use crate::engines::{QueryExecutionError, SimpleEngine}; +use crate::drivers::query::adapters::{ + create_http_adapter, AdapterConfig, HttpProtocolAdapter, PrometheusResponse, +}; +use crate::engines::{QueryExecutionError, RqeQueryError, SimpleEngine}; +use crate::utils::http::convert_query_result_to_prometheus; +use asap_types::rqe_expr::RqeExpr; + +/// Answers an RQE expression tree, whatever the adapter's query language. +pub const RQE_QUERY_ENDPOINT: &str = "/api/v1/rqe_query"; use crate::query_tracker::QueryTracker; use crate::stores::Store; @@ -106,6 +113,7 @@ impl HttpServer { .route(range_query_endpoint, post(handle_range_query_post)) .route(runtime_info_path, get(handle_runtime_info)) .route(runtime_info_path, post(handle_runtime_info)) + .route(RQE_QUERY_ENDPOINT, post(handle_rqe_query)) .route("/metrics", get(handle_metrics)) .with_state(app_state); @@ -143,6 +151,7 @@ impl HttpServer { .route(range_query_endpoint, get(handle_range_query)) .route(range_query_endpoint, post(handle_range_query_post)) .route(runtime_info_path, get(handle_runtime_info)) + .route(RQE_QUERY_ENDPOINT, post(handle_rqe_query)) .with_state(app_state); let listener = TcpListener::bind("127.0.0.1:0").await?; @@ -445,6 +454,75 @@ async fn handle_instant_query_post( result } +/// Body of an RQE query: the tree and its evaluation time in unix seconds. +#[derive(serde::Deserialize)] +#[serde(deny_unknown_fields)] +struct RqeQueryRequest { + expr: RqeExpr, + time: f64, +} + +/// A Prometheus-format error response with an explicit HTTP status. +fn prometheus_error(status: StatusCode, error_type: &str, error: &str) -> Response { + (status, Json(PrometheusResponse::error(error_type, error))).into_response() +} + +/// `POST /api/v1/rqe_query`: answers with Prometheus instant-vector JSON. Never +/// forwards to a fallback backend: an invalid request is a 400, a failed +/// execution a 500, and a window without data yet an empty result. With +/// request handling disabled it answers 500 `internal`, as the instant route +/// does when it has no fallback either. +async fn handle_rqe_query(State(state): State, body: Bytes) -> Response { + if !state.config.handle_http_requests { + debug!("HTTP request handling is disabled; RQE queries have no fallback"); + return prometheus_error( + StatusCode::INTERNAL_SERVER_ERROR, + "internal", + "RQE query handling is disabled", + ); + } + let start_time = Instant::now(); + let request: RqeQueryRequest = match serde_json::from_slice(&body) { + Ok(request) => request, + Err(error) => { + return prometheus_error( + StatusCode::BAD_REQUEST, + "bad_data", + &format!("invalid RQE query body: {error}"), + ); + } + }; + match state + .query_engine + .handle_rqe_query(&request.expr, request.time) + { + Ok((labels, result)) => { + info!( + "rqe_query time={} destination=asap total_latency_ms={:.2}", + request.time, + start_time.elapsed().as_secs_f64() * 1000.0 + ); + match convert_query_result_to_prometheus(&result, &labels) { + Ok(data) => Json(PrometheusResponse::success(data)).into_response(), + Err(error) => { + prometheus_error(StatusCode::INTERNAL_SERVER_ERROR, "internal", error) + } + } + } + Err(RqeQueryError::BadRequest(error)) => { + prometheus_error(StatusCode::BAD_REQUEST, "bad_data", &error) + } + Err(RqeQueryError::Execution(error)) => { + warn!(error = %error, "RQE query execution failed"); + prometheus_error( + StatusCode::INTERNAL_SERVER_ERROR, + "internal", + &error.to_string(), + ) + } + } +} + async fn handle_runtime_info( State(state): State, headers: axum::http::HeaderMap, diff --git a/asap-query-engine/src/engines/mod.rs b/asap-query-engine/src/engines/mod.rs index 5d1c8033..15efd232 100644 --- a/asap-query-engine/src/engines/mod.rs +++ b/asap-query-engine/src/engines/mod.rs @@ -9,5 +9,5 @@ pub mod window_merger; pub use query_result::{InstantVector, QueryResult, RangeVector, RangeVectorElement, Sample}; #[cfg(feature = "native_query_legacy_test_support")] pub use simple_engine::NativeRangeExecutionMode; -pub use simple_engine::{QueryExecutionError, SimpleEngine}; +pub use simple_engine::{QueryExecutionError, RqeQueryError, SimpleEngine}; pub use window_merger::{create_window_merger, NaiveMerger, WindowMerger}; diff --git a/asap-query-engine/src/engines/simple_engine/elastic.rs b/asap-query-engine/src/engines/simple_engine/elastic.rs index 6c974a65..c3d5f461 100644 --- a/asap-query-engine/src/engines/simple_engine/elastic.rs +++ b/asap-query-engine/src/engines/simple_engine/elastic.rs @@ -44,7 +44,7 @@ impl SimpleEngine { // TODO: Figure out how to handle query configuration for ElasticSearch queries. let query_config = self.find_query_config(&query)?; let agg_info = self - .get_aggregation_id_info(&query_config) + .get_aggregation_id_info(&query_config.aggregations) .map_err(|e| { warn!("{}", e); e diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index e2f8990f..5f237069 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -1,12 +1,14 @@ mod elastic; mod promql; +mod rqe; mod sql; use promql::METRIC_NAME_LABEL; +pub use rqe::RqeQueryError; use crate::data_model::{ - AggregationIdInfo, InferenceConfig, KeyByLabelValues, QueryBounds, QueryConfig, QueryLanguage, - StreamingConfig, + AggregationIdInfo, AggregationReference, InferenceConfig, KeyByLabelValues, QueryBounds, + QueryConfig, QueryLanguage, StreamingConfig, }; use crate::engines::query_plan::{ NodeId, PlanOptions, QueryPlan, QueryPlanExecutionError, QueryPlanNode, QueryPlanRuntime, @@ -691,6 +693,47 @@ impl SimpleEngine { )) } + /// Builds a `QueryExecutionContext` from already-resolved plain values: + /// plans the store reads and derives the grouping/aggregated labels from + /// the aggregation configs. Shared by the SQL and RQE context builders. + #[allow(clippy::too_many_arguments)] + fn build_execution_context_tail( + &self, + metric: &str, + timestamps: &QueryTimestamps, + metadata: QueryMetadata, + agg_info: AggregationIdInfo, + spatial_filter: String, + query_time: u64, + ) -> Result { + let (query_plan, do_merge, value_window_type) = + self.create_store_query_plan(metric, timestamps, &agg_info)?; + + let sc = self.streaming_config.read().unwrap().clone(); + let grouping_labels = sc + .get_aggregation_config(agg_info.aggregation_id_for_value) + .map(|config| config.grouping_labels.clone()) + .unwrap_or_else(|| metadata.query_output_labels.clone()); + + let aggregated_labels = sc + .get_aggregation_config(agg_info.aggregation_id_for_key) + .map(|config| config.aggregated_labels.clone()) + .unwrap_or_else(KeyByLabelNames::empty); + + Ok(QueryExecutionContext { + metric: metric.to_string(), + metadata, + store_plan: query_plan, + agg_info, + value_window_type, + do_merge, + spatial_filter, + query_time, + grouping_labels, + aggregated_labels, + }) + } + /// The bucket-map grid width for scanning an aggregation's stored /// buckets: `slide_interval_ms`, not `window_size_ms`. /// `precompute_engine/window_manager.rs` persists buckets on the @@ -1771,12 +1814,12 @@ impl SimpleEngine { stats.into_iter().next() } + /// Resolves the key and value aggregations a query reads. With two + /// aggregations, the set/delta-set one supplies keys and the other values. fn get_aggregation_id_info( &self, - query_config: &QueryConfig, + query_config_aggregations: &[AggregationReference], ) -> Result { - let query_config_aggregations = &query_config.aggregations; - if query_config_aggregations.is_empty() { return Err("Query config has no aggregations defined".to_string()); } diff --git a/asap-query-engine/src/engines/simple_engine/promql.rs b/asap-query-engine/src/engines/simple_engine/promql.rs index 0cdabfd3..f9f7bd29 100644 --- a/asap-query-engine/src/engines/simple_engine/promql.rs +++ b/asap-query-engine/src/engines/simple_engine/promql.rs @@ -51,7 +51,7 @@ fn detect_scalar_arm<'a>( /// two arms don't share the same label set — mirroring DataFusion's /// `build_binary_vector_plan`, which fails to resolve a join column that /// only exists on one side. -fn combine_vector_vector( +pub(super) fn combine_vector_vector( lhs_results: Vec, lhs_labels: &[String], rhs_results: Vec, @@ -83,7 +83,7 @@ fn combine_vector_vector( /// Scalar combiner for one binary-expr level (instant query): applies /// `op(scalar, value)` or `op(value, scalar)` per `scalar_on_left` to every /// element of the vector arm's results. -fn combine_scalar( +pub(super) fn combine_scalar( vector_results: Vec, scalar: f64, op: &promql_parser::parser::token::TokenType, @@ -119,14 +119,14 @@ impl SimpleEngine { /// boundary, unconditionally — mirroring SQL's `align_end_timestamp_sql`. /// A no-op in the common case (already-aligned timestamps), a safety net /// otherwise, for every PromQL query shape, not just a subset. - fn align_end_timestamp_promql(&self, end_timestamp: u64) -> u64 { + pub(super) fn align_end_timestamp_promql(&self, end_timestamp: u64) -> u64 { let interval_ms = self.data_ingestion_interval_ms; if end_timestamp.is_multiple_of(interval_ms) { return end_timestamp; } let aligned = (end_timestamp / interval_ms) * interval_ms; warn!( - "PromQL query end timestamp {} is not aligned with data ingestion interval of {} ms; \ + "Query end timestamp {} is not aligned with data ingestion interval of {} ms; \ aligning down to {}.", end_timestamp, interval_ms, aligned ); @@ -286,7 +286,7 @@ impl SimpleEngine { let match_result = self.find_matching_controller_pattern(arm_ast, planned_subquery)?; let agg_info = self - .get_aggregation_id_info(query_config) + .get_aggregation_id_info(&query_config.aggregations) .map_err(|e| { warn!("{}", e); e @@ -1323,7 +1323,7 @@ impl SimpleEngine { // Resolve aggregation: try pre-configured query_configs first, fall back to capability matching. let agg_info: AggregationIdInfo = if let Some(config) = self.find_query_config(query) { - self.get_aggregation_id_info(&config) + self.get_aggregation_id_info(&config.aggregations) .map_err(|e| { warn!("{}", e); e diff --git a/asap-query-engine/src/engines/simple_engine/rqe.rs b/asap-query-engine/src/engines/simple_engine/rqe.rs new file mode 100644 index 00000000..63065672 --- /dev/null +++ b/asap-query-engine/src/engines/simple_engine/rqe.rs @@ -0,0 +1,557 @@ +//! Answers an RQE expression tree at one query time. A leaf reads the +//! aggregations its `rqe_configs` entry names (no query text is parsed), an +//! `agg` node regroups its child, and a `binary` node combines its sides the +//! way PromQL arithmetic does. + +use super::promql::{combine_scalar, combine_vector_vector}; +use super::{QueryExecutionError, QueryMetadata, QueryTimestamps, SimpleEngine}; +use crate::data_model::{RqeConfig, StreamingConfig}; +use crate::engines::query_result::{InstantVectorElement, QueryResult}; +use crate::engines::query_time_aggregation::{ + apply_instant_pipeline, output_labels_for_aggregation, pipeline_supports_labels, +}; +use crate::precompute_engine::accumulator_factory::cms_heap_params; +use asap_types::aggregation_config::AggregationConfig; +use asap_types::aggregation_mode::{AggregationMode, CountMode}; +use asap_types::capability_matching::{compatible_agg_types, window_compatible}; +use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeGrouping, QueryTimeGroupingMode, +}; +use asap_types::rqe_expr::{RqeAggOp, RqeBinaryOp, RqeCapability, RqeExpr, RqeLeaf}; +use promql_parser::parser::token::{TokenType, T_ADD, T_DIV, T_MUL, T_SUB}; +use promql_utilities::data_model::KeyByLabelNames; +use promql_utilities::query_logics::enums::Statistic; +use std::collections::HashMap; + +#[derive(Debug, thiserror::Error)] +pub enum RqeQueryError { + /// The request is wrong: an invalid tree, an unknown RQE id, a leaf its + /// RQE's plan can't serve, or a bad time. + #[error("{0}")] + BadRequest(String), + /// The request is valid but the engine could not answer it. + #[error(transparent)] + Execution(#[from] QueryExecutionError), +} + +/// A node's result: an instant vector with its label names, or a literal. +enum RqeValue { + Vector { + labels: KeyByLabelNames, + values: Vec, + }, + Scalar(f64), +} + +impl SimpleEngine { + /// Evaluates `expr` at `time` (unix seconds) into an instant vector. + /// + /// Output labels are the root's labels (see `RqeExpr::output_labels`), + /// never `__name__`: every leaf reads a range like `sum_over_time`, whose + /// result Prometheus strips of its name. A leaf whose window holds no + /// data yet contributes an empty vector. + pub fn handle_rqe_query( + &self, + expr: &RqeExpr, + time: f64, + ) -> Result<(KeyByLabelNames, QueryResult), RqeQueryError> { + expr.validate() + .map_err(|error| RqeQueryError::BadRequest(error.to_string()))?; + if !time.is_finite() || time < 0.0 { + return Err(RqeQueryError::BadRequest(format!( + "time must be a non-negative number of unix seconds, got {time}" + ))); + } + let leaves = expr.leaves(); + // One snapshot (a shared pointer, not a copy) for every leaf to read. + let streaming_config = self.streaming_config.read().unwrap().clone(); + // The leaves' entries, copied so the lock isn't held while reading. + // Configs are loaded once and don't change while a tree is answered. + let configs: HashMap<&str, RqeConfig> = { + let inference_config = self.inference_config.read().unwrap(); + leaves + .iter() + .map(|leaf| { + inference_config + .rqe_configs + .get(&leaf.id) + .map(|config| (leaf.id.as_str(), config.clone())) + .ok_or_else(|| { + RqeQueryError::BadRequest(format!("unknown RQE id {:?}", leaf.id)) + }) + }) + .collect::>()? + }; + + let query_time = Self::convert_query_time_to_data_time(time); + match self.evaluate_rqe_node(expr, &configs, &streaming_config, query_time)? { + RqeValue::Vector { labels, values } => { + Ok((labels, QueryResult::vector(values, query_time))) + } + RqeValue::Scalar(_) => unreachable!("a valid tree's root reads a leaf"), + } + } + + fn evaluate_rqe_node( + &self, + expr: &RqeExpr, + configs: &HashMap<&str, RqeConfig>, + streaming_config: &StreamingConfig, + query_time: u64, + ) -> Result { + match expr { + RqeExpr::Rqe { rqe } => { + let (labels, values) = self.evaluate_rqe_leaf( + rqe, + &configs[rqe.id.as_str()], + streaming_config, + query_time, + )?; + Ok(RqeValue::Vector { labels, values }) + } + RqeExpr::Agg { agg, by, of } => { + let RqeValue::Vector { labels, values } = + self.evaluate_rqe_node(of, configs, streaming_config, query_time)? + else { + unreachable!("a valid agg operand reads a leaf"); + }; + let (labels, values) = + apply_instant_pipeline(labels, values, &[regroup_stage(*agg, by)]) + .map_err(QueryExecutionError::Native)?; + Ok(RqeValue::Vector { labels, values }) + } + RqeExpr::Binary { op, lhs, rhs } => { + let op = binary_token(*op); + let lhs = self.evaluate_rqe_node(lhs, configs, streaming_config, query_time)?; + let rhs = self.evaluate_rqe_node(rhs, configs, streaming_config, query_time)?; + match (lhs, rhs) { + ( + RqeValue::Vector { + labels, + values: lhs_values, + }, + RqeValue::Vector { + labels: rhs_labels, + values: rhs_values, + }, + ) => { + let values = combine_vector_vector( + lhs_values, + &labels.labels, + rhs_values, + &rhs_labels.labels, + &op, + ) + .ok_or_else(|| { + RqeQueryError::BadRequest(format!( + "binary op sides have different labels: {:?} vs {:?}", + labels.labels, rhs_labels.labels + )) + })?; + Ok(RqeValue::Vector { labels, values }) + } + (RqeValue::Vector { labels, values }, RqeValue::Scalar(scalar)) => { + Ok(RqeValue::Vector { + labels, + values: combine_scalar(values, scalar, &op, false), + }) + } + (RqeValue::Scalar(scalar), RqeValue::Vector { labels, values }) => { + Ok(RqeValue::Vector { + labels, + values: combine_scalar(values, scalar, &op, true), + }) + } + (RqeValue::Scalar(_), RqeValue::Scalar(_)) => { + unreachable!("a valid binary op reads a leaf on one side") + } + } + } + RqeExpr::Scalar { scalar } => Ok(RqeValue::Scalar(*scalar)), + } + } + + /// Reads `leaf` from the aggregations `config` names, over the window + /// ending at `query_time` aligned as a PromQL instant query aligns it. + /// A roll-up config reads at the aggregation's own grouping, then its + /// query-time stages regroup to the leaf's. A top-k leaf ranks within + /// each bucket of its aggregation's grouping, as PromQL `topk by` does, + /// so that grouping must be the leaf's `topk_by`. + fn evaluate_rqe_leaf( + &self, + leaf: &RqeLeaf, + config: &RqeConfig, + streaming_config: &StreamingConfig, + query_time: u64, + ) -> Result<(KeyByLabelNames, Vec), RqeQueryError> { + let failed = + |error: String| QueryExecutionError::Native(format!("RQE {:?}: {error}", leaf.id)); + let agg_info = self + .get_aggregation_id_info(&config.aggregations) + .map_err(failed)?; + let value_config = streaming_config + .get_aggregation_config(agg_info.aggregation_id_for_value) + .ok_or_else(|| { + failed(format!( + "no streaming config for aggregation_id {}", + agg_info.aggregation_id_for_value + )) + })?; + + let (statistic_to_compute, query_kwargs) = leaf_statistic(leaf); + check_leaf_against_aggregation(leaf, statistic_to_compute, value_config)?; + check_leaf_retention(leaf, config, value_config)?; + + let end_timestamp = self.align_end_timestamp_promql(query_time); + let start_timestamp = end_timestamp.checked_sub(leaf.lookback_ms).ok_or_else(|| { + RqeQueryError::BadRequest(format!( + "RQE {:?}: its {} ms lookback starts before the unix epoch", + leaf.id, leaf.lookback_ms + )) + })?; + + let leaf_labels = KeyByLabelNames::new(leaf.grouping.clone()); + let read_labels = if config.query_time_aggregations.is_empty() { + check_read_grouping(leaf, &leaf_labels, statistic_to_compute, value_config)?; + leaf_labels.clone() + } else { + let mut labels = value_config.grouping_labels.labels.clone(); + labels.extend(value_config.aggregated_labels.labels.iter().cloned()); + let labels = KeyByLabelNames::new(labels); + check_roll_up_grouping(leaf, &leaf_labels, &labels, &config.query_time_aggregations)?; + labels + }; + let metadata = QueryMetadata { + query_output_labels: read_labels, + statistic_to_compute, + query_kwargs, + keep_metric_name: false, + }; + let context = self + .build_execution_context_tail( + &value_config.metric, + &QueryTimestamps { + start_timestamp, + end_timestamp, + }, + metadata, + agg_info, + String::new(), + query_time, + ) + .map_err(failed)?; + + let Some((labels, result)) = self.execute_context_result(context, true, true)? else { + return Ok((leaf_labels, Vec::new())); + }; + let QueryResult::Vector(vector) = result else { + return Err(failed("an instant read returned a range result".to_string()).into()); + }; + let (labels, values) = if config.query_time_aggregations.is_empty() { + (labels, vector.values) + } else { + apply_instant_pipeline(labels, vector.values, &config.query_time_aggregations) + .map_err(failed)? + }; + debug_assert_eq!(labels, leaf_labels, "checked before reading"); + Ok((labels, values)) + } +} + +fn bad_leaf(leaf: &RqeLeaf, reason: String) -> RqeQueryError { + RqeQueryError::BadRequest(format!("RQE {:?}: {reason}", leaf.id)) +} + +/// Whether the aggregation `leaf`'s RQE reads can answer it: same metric +/// (`table.value` for SQL), a type serving its capability, and windows that +/// tile its lookback. A top-k leaf must rank within the aggregation's heaps, +/// and ask for no more than a heap holds. +fn check_leaf_against_aggregation( + leaf: &RqeLeaf, + statistic: Statistic, + aggregation: &AggregationConfig, +) -> Result<(), RqeQueryError> { + // A SQL aggregation's metric is `table.value_column`. + let metric = match &leaf.value { + Some(value) => format!("{}.{value}", leaf.metric), + None => leaf.metric.clone(), + }; + if metric != aggregation.metric { + return Err(bad_leaf( + leaf, + format!( + "reads {metric:?}, but its aggregation {} is on {:?}", + aggregation.aggregation_id, aggregation.metric + ), + )); + } + // RQE sums and top-k rank values, so a SUM/COUNT sketch must weight by value. + let value_weighted = match aggregation.mode() { + Ok(AggregationMode::SumOrCount(mode)) => mode == CountMode::Sum, + Ok(_) => true, + Err(error) => { + return Err(QueryExecutionError::Native(format!("RQE {:?}: {error}", leaf.id)).into()) + } + }; + if !compatible_agg_types(statistic).contains(&aggregation.aggregation_type) || !value_weighted { + return Err(bad_leaf( + leaf, + format!( + "{:?} can't be answered by its aggregation {} ({} {:?})", + leaf.capability, + aggregation.aggregation_id, + aggregation.aggregation_type, + aggregation.aggregation_sub_type + ), + )); + } + if !window_compatible(aggregation, leaf.lookback_ms) { + return Err(bad_leaf( + leaf, + format!( + "its {} ms lookback is not a whole number of its aggregation {}'s {} ms windows \ + (sliding every {} ms)", + leaf.lookback_ms, + aggregation.aggregation_id, + aggregation.window_size_ms, + aggregation.slide_interval_ms + ), + )); + } + if leaf.capability == RqeCapability::TopkByValue { + // Top-k is served only by a CountMinSketchWithHeap (checked above). + let (_, _, heap_size) = cms_heap_params(aggregation) + .map_err(|error| QueryExecutionError::Native(format!("RQE {:?}: {error}", leaf.id)))?; + let k = leaf.k.expect("a valid top-k leaf has k"); + if k > heap_size as u64 { + return Err(bad_leaf( + leaf, + format!( + "asks for the top {k}, but its aggregation {} keeps a heap of {heap_size}", + aggregation.aggregation_id + ), + )); + } + let buckets = KeyByLabelNames::new(leaf.topk_by.clone().unwrap_or_default()); + if aggregation.grouping_labels != buckets { + return Err(bad_leaf( + leaf, + format!( + "ranks within topk_by buckets {:?}, but its aggregation keeps a heap per {:?}", + buckets.labels, aggregation.grouping_labels.labels + ), + )); + } + } + Ok(()) +} + +/// The closed windows a lookback reads: those wholly inside it, which start +/// in `[end - lookback, end - window]`, so `(lookback - window) / slide + 1`. +/// The planner retains this many for the lookback (its `retained_windows`). +/// The lookback must already be `window_compatible`. +fn windows_needed(aggregation: &AggregationConfig, lookback_ms: u64) -> u64 { + (lookback_ms - aggregation.window_size_ms) / SimpleEngine::bucket_step_ms(aggregation) + 1 +} + +/// A leaf may look back only as many windows as its RQE's plan retains of +/// the aggregation it reads: older ones may already be gone. +fn check_leaf_retention( + leaf: &RqeLeaf, + config: &RqeConfig, + aggregation: &AggregationConfig, +) -> Result<(), RqeQueryError> { + let retained = config + .aggregations + .iter() + .find(|reference| reference.aggregation_id == aggregation.aggregation_id) + .and_then(|reference| reference.num_aggregates_to_retain); + let Some(retained) = retained else { + return Ok(()); + }; + let needed = windows_needed(aggregation, leaf.lookback_ms); + if needed > retained { + return Err(bad_leaf( + leaf, + format!( + "its {} ms lookback reads {needed} windows of its aggregation {}, \ + but its plan retains {retained}", + leaf.lookback_ms, aggregation.aggregation_id + ), + )); + } + Ok(()) +} + +/// A leaf read with no query-time stage gets the aggregation's own labels +/// (its groups and keys), or, from a single-population aggregation that +/// merges across groups, a subset of its groups. +fn check_read_grouping( + leaf: &RqeLeaf, + leaf_labels: &KeyByLabelNames, + statistic: Statistic, + aggregation: &AggregationConfig, +) -> Result<(), RqeQueryError> { + let groups = &aggregation.grouping_labels.labels; + let mut served = groups.clone(); + served.extend(aggregation.aggregated_labels.labels.iter().cloned()); + let served = KeyByLabelNames::new(served); + let rolls_up = aggregation.aggregated_labels.labels.is_empty() + && aggregation.aggregation_type.merges_across_groups() + && statistic != Statistic::Topk + && leaf_labels + .labels + .iter() + .all(|label| groups.contains(label)); + if *leaf_labels == served || rolls_up { + return Ok(()); + } + Err(bad_leaf( + leaf, + format!( + "groups by {:?}, but its aggregation {} serves {:?}", + leaf_labels.labels, aggregation.aggregation_id, served.labels + ), + )) +} + +/// A roll-up's query-time stages must regroup the labels read into the +/// leaf's grouping. +fn check_roll_up_grouping( + leaf: &RqeLeaf, + leaf_labels: &KeyByLabelNames, + read_labels: &KeyByLabelNames, + stages: &[QueryTimeAggregation], +) -> Result<(), RqeQueryError> { + let invalid_stage = + |error: String| QueryExecutionError::Native(format!("RQE {:?}: {error}", leaf.id)); + if !pipeline_supports_labels(read_labels, stages).map_err(invalid_stage)? { + return Err(bad_leaf( + leaf, + format!( + "its query-time aggregations can't regroup labels {:?}", + read_labels.labels + ), + )); + } + let output = stages + .iter() + .try_fold(read_labels.clone(), |labels, stage| { + output_labels_for_aggregation(&labels, stage) + }) + .map_err(invalid_stage)?; + if output != *leaf_labels { + return Err(bad_leaf( + leaf, + format!( + "groups by {:?}, but its query-time aggregations regroup to {:?}", + leaf_labels.labels, output.labels + ), + )); + } + Ok(()) +} + +/// The statistic and its kwargs a validated leaf's capability computes. +fn leaf_statistic(leaf: &RqeLeaf) -> (Statistic, HashMap) { + match leaf.capability { + RqeCapability::Sum => (Statistic::Sum, HashMap::new()), + RqeCapability::Cardinality => (Statistic::Cardinality, HashMap::new()), + RqeCapability::Quantile => { + let phi = leaf.phi.expect("a valid quantile leaf has phi"); + ( + Statistic::Quantile, + HashMap::from([("quantile".to_string(), phi.to_string())]), + ) + } + RqeCapability::TopkByValue => { + let k = leaf.k.expect("a valid top-k leaf has k"); + ( + Statistic::Topk, + HashMap::from([("k".to_string(), k.to_string())]), + ) + } + } +} + +/// The query-time stage that regroups a vector by `by` (all into one group +/// when `by` is empty). +fn regroup_stage(op: RqeAggOp, by: &[String]) -> QueryTimeAggregation { + let operator = match op { + RqeAggOp::Avg => QueryTimeAggregationOperator::Avg, + RqeAggOp::Min => QueryTimeAggregationOperator::Min, + RqeAggOp::Max => QueryTimeAggregationOperator::Max, + RqeAggOp::Sum => QueryTimeAggregationOperator::Sum, + }; + let mode = if by.is_empty() { + QueryTimeGroupingMode::All + } else { + QueryTimeGroupingMode::By + }; + QueryTimeAggregation { + operator, + grouping: QueryTimeGrouping { + mode, + labels: by.to_vec(), + }, + parameter: None, + } +} + +fn binary_token(op: RqeBinaryOp) -> TokenType { + TokenType::new(match op { + RqeBinaryOp::Add => T_ADD, + RqeBinaryOp::Sub => T_SUB, + RqeBinaryOp::Mul => T_MUL, + RqeBinaryOp::Div => T_DIV, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_types::enums::{AggregationType, WindowType}; + + fn aggregation(window_size_ms: u64, slide_interval_ms: u64) -> AggregationConfig { + let window_type = if window_size_ms == slide_interval_ms { + WindowType::Tumbling + } else { + WindowType::Sliding + }; + AggregationConfig::new( + 1, + AggregationType::Sum, + String::new(), + HashMap::new(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + String::new(), + window_size_ms, + slide_interval_ms, + window_type, + String::new(), + "m".to_string(), + None, + None, + None, + None, + ) + } + + #[test] + fn windows_needed_counts_the_closed_windows_inside_the_lookback() { + const MINUTE_MS: u64 = 60_000; + // Tumbling: one window per window width. + assert_eq!( + windows_needed(&aggregation(MINUTE_MS, MINUTE_MS), 3 * MINUTE_MS), + 3 + ); + // A tumbling config may leave the slide unset. + assert_eq!(windows_needed(&aggregation(MINUTE_MS, 0), 3 * MINUTE_MS), 3); + // 5m windows sliding every minute: the newest for 5m, eleven for 15m. + let sliding = aggregation(5 * MINUTE_MS, MINUTE_MS); + assert_eq!(windows_needed(&sliding, 5 * MINUTE_MS), 1); + assert_eq!(windows_needed(&sliding, 15 * MINUTE_MS), 11); + } +} diff --git a/asap-query-engine/src/engines/simple_engine/sql.rs b/asap-query-engine/src/engines/simple_engine/sql.rs index 953bf533..22330fff 100644 --- a/asap-query-engine/src/engines/simple_engine/sql.rs +++ b/asap-query-engine/src/engines/simple_engine/sql.rs @@ -440,55 +440,6 @@ impl SimpleEngine { )) } - /// Shared context-building tail for both SQL context builders. - /// - /// Called by `build_query_execution_context_sql` and `build_spatiotemporal_context` - /// after labels, statistic, metadata, timestamps, and `agg_info` are resolved. - /// Builds the query plan, derives grouping/aggregated labels, and returns the - /// final `QueryExecutionContext`. - #[allow(clippy::too_many_arguments)] - fn build_sql_execution_context_tail( - &self, - metric: &str, - timestamps: &QueryTimestamps, - metadata: QueryMetadata, - agg_info: AggregationIdInfo, - spatial_filter: String, - query_time: u64, - ) -> Option { - let (query_plan, do_merge, value_window_type) = self - .create_store_query_plan(metric, timestamps, &agg_info) - .map_err(|e| { - warn!("Failed to create store query plan: {}", e); - e - }) - .ok()?; - - let sc = self.streaming_config.read().unwrap().clone(); - let grouping_labels = sc - .get_aggregation_config(agg_info.aggregation_id_for_value) - .map(|config| config.grouping_labels.clone()) - .unwrap_or_else(|| metadata.query_output_labels.clone()); - - let aggregated_labels = sc - .get_aggregation_config(agg_info.aggregation_id_for_key) - .map(|config| config.aggregated_labels.clone()) - .unwrap_or_else(KeyByLabelNames::empty); - - Some(QueryExecutionContext { - metric: metric.to_string(), - metadata, - store_plan: query_plan, - agg_info, - value_window_type, - do_merge, - spatial_filter, - query_time, - grouping_labels, - aggregated_labels, - }) - } - /// Build execution context for SpatioTemporal queries. /// These queries span multiple scrape intervals but GROUP BY a subset of labels. /// `config` is the query's planned config; without one the aggregation @@ -564,7 +515,7 @@ impl SimpleEngine { // Resolve aggregation: try pre-configured query_configs first, fall back to capability matching. let agg_info: AggregationIdInfo = if let Some(config) = config { - self.get_aggregation_id_info(&config) + self.get_aggregation_id_info(&config.aggregations) .map_err(|e| { warn!("{}", e); e @@ -584,7 +535,7 @@ impl SimpleEngine { }; let metric = &match_result.outer_data()?.metric; - self.build_sql_execution_context_tail( + self.build_execution_context_tail( metric, ×tamps, metadata, @@ -592,6 +543,8 @@ impl SimpleEngine { String::new(), query_time, ) + .map_err(|e| warn!("Failed to create store query plan: {}", e)) + .ok() } } diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index e55149b5..a03023b8 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -926,7 +926,7 @@ fn hydra_kll_params(config: &AggregationConfig) -> Result<(usize, usize, u16), S /// then falls back to the `row_num`/`col_num`/`heap_size` aliases. All three /// parameters are required — the planner always emits them and their absence /// indicates a malformed config. -fn cms_heap_params(config: &AggregationConfig) -> Result<(usize, usize, usize), String> { +pub(crate) fn cms_heap_params(config: &AggregationConfig) -> Result<(usize, usize, usize), String> { let read = |names: &[&str]| -> Result { names .iter() diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index 9bb4a15d..a737299f 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -13,6 +13,7 @@ pub mod query_equivalence_tests; pub mod query_time_aggregation_public_tests; pub mod range_query_arithmetic_tests; pub mod rollup_tests; +pub mod rqe_query_tests; pub mod sql_nested_aggregation_tests; pub mod sql_pattern_matching_tests; pub mod stage_e_instant_range_equivalence_tests; diff --git a/asap-query-engine/src/tests/rqe_query_tests.rs b/asap-query-engine/src/tests/rqe_query_tests.rs new file mode 100644 index 00000000..7f9b13f7 --- /dev/null +++ b/asap-query-engine/src/tests/rqe_query_tests.rs @@ -0,0 +1,1128 @@ +//! RQE expression trees answered by `SimpleEngine::handle_rqe_query` and the +//! `POST /api/v1/rqe_query` route. + +#[cfg(test)] +mod tests { + use crate::data_model::{ + AggregationConfig, AggregationReference, AggregationType, CleanupPolicy, InferenceConfig, + KeyByLabelValues, PrecomputedOutput, PromQLSchema, QueryConfig, QueryLanguage, RqeConfig, + SchemaConfig, StreamingConfig, WindowType, + }; + use crate::drivers::query::adapters::AdapterConfig; + use crate::drivers::query::servers::http::{HttpServer, HttpServerConfig, RQE_QUERY_ENDPOINT}; + use crate::engines::{QueryResult, RqeQueryError, SimpleEngine}; + use crate::precompute_operators::{ + CountMinSketchWithHeapAccumulator, DDSketchAccumulator, HllAccumulator, + MultipleSumAccumulator, SumAccumulator, + }; + use crate::stores::simple_map_store::SimpleMapStore; + use crate::stores::Store; + use crate::AggregateCore; + use asap_types::query_config::{ + QueryTimeAggregation, QueryTimeAggregationOperator, QueryTimeGrouping, + QueryTimeGroupingMode, + }; + use asap_types::rqe_expr::RqeExpr; + use promql_utilities::data_model::KeyByLabelNames; + use serde_json::{json, Value}; + use std::collections::HashMap; + use std::sync::Arc; + + const WINDOW_MS: u64 = 1_000; + const QUERY_TIME_S: f64 = 1_000.0; + const QUERY_TIME_MS: u64 = 1_000_000; + /// Two tumbling windows: ending one window before the query time, and at it. + const TWO_WINDOWS_MS: u64 = 2 * WINDOW_MS; + + const SUM_AGG: u64 = 1; + const DDSKETCH_AGG: u64 = 2; + const HLL_AGG: u64 = 3; + const TOPK_AGG: u64 = 4; + const MULTI_SUM_AGG: u64 = 5; + /// One heap per `job`, keyed by `srcip`: serves `topk by (job)`. + const TOPK_BY_JOB_AGG: u64 = 6; + const TOPK_BY_JOB_QUERY: &str = "topk by (job) (2, sum_over_time(flows[1s]))"; + /// A SQL table's `latency` column, as the planner names it. + const SQL_DDSKETCH_AGG: u64 = 7; + /// Each top-k aggregation's heap capacity. + const HEAP_SIZE: usize = 32; + const CMS_DEPTH: usize = 3; + const CMS_WIDTH: usize = 1024; + /// In `rqe_configs` but not in the streaming config. + const MISSING_AGG: u64 = 99; + + type Row = (Vec, f64); + + fn strings(values: &[&str]) -> Vec { + values.iter().map(|value| value.to_string()).collect() + } + + fn aggregation( + id: u64, + aggregation_type: AggregationType, + metric: &str, + grouping: &[&str], + aggregated: &[&str], + ) -> (u64, AggregationConfig) { + ( + id, + AggregationConfig { + aggregation_id: id, + aggregation_type, + // Sums and top-k heaps weight by value. + aggregation_sub_type: match aggregation_type { + AggregationType::MultipleSum | AggregationType::CountMinSketchWithHeap => { + "sum".to_string() + } + _ => String::new(), + }, + parameters: match aggregation_type { + AggregationType::CountMinSketchWithHeap => HashMap::from([ + ("depth".to_string(), json!(CMS_DEPTH)), + ("width".to_string(), json!(CMS_WIDTH)), + ("heapsize".to_string(), json!(HEAP_SIZE)), + ]), + _ => HashMap::new(), + }, + grouping_labels: KeyByLabelNames::new(strings(grouping)), + aggregated_labels: KeyByLabelNames::new(strings(aggregated)), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size_ms: WINDOW_MS, + slide_interval_ms: WINDOW_MS, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: metric.to_string(), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }, + ) + } + + fn sum_by_job() -> QueryTimeAggregation { + QueryTimeAggregation { + operator: QueryTimeAggregationOperator::Sum, + grouping: QueryTimeGrouping { + mode: QueryTimeGroupingMode::By, + labels: strings(&["job"]), + }, + parameter: None, + } + } + + fn rqe_config(aggregation_id: u64, stages: Vec) -> RqeConfig { + RqeConfig { + aggregations: vec![AggregationReference::new(aggregation_id, None)], + query_time_aggregations: stages, + } + } + + fn query_config(query: &str, aggregation_id: u64) -> QueryConfig { + QueryConfig::new(query.to_string()) + .add_aggregation(AggregationReference::new(aggregation_id, None)) + } + + /// DDSketch, not KLL: its merge is deterministic, so two reads of the + /// same windows agree exactly. + fn ddsketch(values: impl IntoIterator) -> DDSketchAccumulator { + let mut sketch = DDSketchAccumulator::new(0.01); + values.into_iter().for_each(|value| sketch.update(value)); + sketch + } + + fn hll(values: std::ops::Range) -> HllAccumulator { + let mut sketch = HllAccumulator::with_default_precision(); + values.for_each(|value| sketch.update(f64::from(value))); + sketch + } + + fn multiple_sum(sums: &[(&[&str], f64)]) -> MultipleSumAccumulator { + let mut accumulator = MultipleSumAccumulator::new(); + for (labels, sum) in sums { + accumulator.add_sum(KeyByLabelValues::new_with_labels(strings(labels)), *sum); + } + accumulator + } + + /// One engine holding every aggregation the tests read, each RQE id in + /// `rqe_configs`, and the PromQL queries the equivalence tests compare + /// against. + fn engine() -> SimpleEngine { + let streaming_config = Arc::new(StreamingConfig { + aggregation_configs: HashMap::from([ + aggregation( + SUM_AGG, + AggregationType::Sum, + "http_requests", + &["instance", "job"], + &[], + ), + aggregation( + DDSKETCH_AGG, + AggregationType::DDSketch, + "latency", + &["job"], + &[], + ), + aggregation(HLL_AGG, AggregationType::HLL, "users", &["job"], &[]), + aggregation( + TOPK_AGG, + AggregationType::CountMinSketchWithHeap, + "transfers", + &[], + &["srcip"], + ), + aggregation( + MULTI_SUM_AGG, + AggregationType::MultipleSum, + "events", + &[], + &["instance", "job"], + ), + aggregation( + TOPK_BY_JOB_AGG, + AggregationType::CountMinSketchWithHeap, + "flows", + &["job"], + &["srcip"], + ), + aggregation( + SQL_DDSKETCH_AGG, + AggregationType::DDSketch, + "http.latency", + &["job"], + &[], + ), + ]), + }); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + let insert = |aggregation_id: u64, + end_ms: u64, + key: Option<&[&str]>, + accumulator: Box| { + let key = key.map(|labels| KeyByLabelValues::new_with_labels(strings(labels))); + store + .insert_precomputed_output( + PrecomputedOutput::new(end_ms - WINDOW_MS, end_ms, key, aggregation_id), + accumulator, + ) + .unwrap(); + }; + let previous_window_ms = QUERY_TIME_MS - WINDOW_MS; + + // Sums: (a, api) 1 then 2, (b, api) 3 then 4, (a, web) 5 then 6. + for (labels, earlier, latest) in [ + (["a", "api"], 1.0, 2.0), + (["b", "api"], 3.0, 4.0), + (["a", "web"], 5.0, 6.0), + ] { + insert( + SUM_AGG, + previous_window_ms, + Some(&labels), + Box::new(SumAccumulator::with_sum(earlier)), + ); + insert( + SUM_AGG, + QUERY_TIME_MS, + Some(&labels), + Box::new(SumAccumulator::with_sum(latest)), + ); + } + + for (job, offset) in [("api", 0.0), ("web", 1_000.0)] { + insert( + DDSKETCH_AGG, + previous_window_ms, + Some(&[job]), + Box::new(ddsketch((0..500).map(|v| f64::from(v) + offset))), + ); + insert( + DDSKETCH_AGG, + QUERY_TIME_MS, + Some(&[job]), + Box::new(ddsketch((500..1_000).map(|v| f64::from(v) + offset))), + ); + } + + insert( + HLL_AGG, + QUERY_TIME_MS, + Some(&["api"]), + Box::new(hll(0..1_000)), + ); + insert(HLL_AGG, QUERY_TIME_MS, Some(&["web"]), Box::new(hll(0..50))); + + let mut heap = CountMinSketchWithHeapAccumulator::new(CMS_DEPTH, CMS_WIDTH, HEAP_SIZE); + for i in 1..=5u64 { + heap.inner.update(&format!("10.0.0.{i}"), (i * 10) as f64); + } + insert(TOPK_AGG, QUERY_TIME_MS, None, Box::new(heap)); + + // `api`'s smallest sources outweigh `web`'s largest, so a global + // ranking would keep only `api` rows. + for (job, sums) in [("api", [100.0, 200.0, 300.0]), ("web", [10.0, 20.0, 30.0])] { + let mut heap = CountMinSketchWithHeapAccumulator::new(CMS_DEPTH, CMS_WIDTH, HEAP_SIZE); + for (i, sum) in sums.into_iter().enumerate() { + heap.inner.update(&format!("10.0.0.{}", i + 1), sum); + } + insert(TOPK_BY_JOB_AGG, QUERY_TIME_MS, Some(&[job]), Box::new(heap)); + } + + insert( + MULTI_SUM_AGG, + previous_window_ms, + None, + Box::new(multiple_sum(&[(&["a", "api"], 1.0), (&["b", "api"], 2.0)])), + ); + insert( + MULTI_SUM_AGG, + QUERY_TIME_MS, + None, + Box::new(multiple_sum(&[(&["a", "api"], 4.0), (&["a", "web"], 8.0)])), + ); + + let schema = PromQLSchema::new() + .add_metric( + "http_requests".to_string(), + KeyByLabelNames::new(strings(&["instance", "job"])), + ) + .add_metric( + "latency".to_string(), + KeyByLabelNames::new(strings(&["job"])), + ) + .add_metric("users".to_string(), KeyByLabelNames::new(strings(&["job"]))) + .add_metric( + "transfers".to_string(), + KeyByLabelNames::new(strings(&["srcip"])), + ) + .add_metric( + "events".to_string(), + KeyByLabelNames::new(strings(&["instance", "job"])), + ) + .add_metric( + "flows".to_string(), + KeyByLabelNames::new(strings(&["job", "srcip"])), + ); + let inference_config = InferenceConfig { + schema: SchemaConfig::PromQL(schema), + query_configs: vec![ + query_config("sum_over_time(http_requests[2s])", SUM_AGG), + query_config("sum by (job) (sum_over_time(http_requests[2s]))", SUM_AGG), + query_config("quantile_over_time(0.9, latency[2s])", DDSKETCH_AGG), + query_config("topk(3, sum_over_time(transfers[1s]))", TOPK_AGG), + query_config(TOPK_BY_JOB_QUERY, TOPK_BY_JOB_AGG), + QueryConfig::with_plan( + "sum by (job) (sum_over_time(events[2s]))".to_string(), + "sum by (instance, job) (sum_over_time(events[2s]))".to_string(), + vec![sum_by_job()], + ) + .add_aggregation(AggregationReference::new(MULTI_SUM_AGG, None)), + ], + cleanup_policy: CleanupPolicy::NoCleanup, + rqe_configs: HashMap::from([ + ("sum".to_string(), rqe_config(SUM_AGG, vec![])), + ("sum_by_job".to_string(), rqe_config(SUM_AGG, vec![])), + ("latest".to_string(), rqe_config(SUM_AGG, vec![])), + ( + // Planned for a lookback of two windows. + "sum_retains_two".to_string(), + RqeConfig { + aggregations: vec![AggregationReference::new(SUM_AGG, Some(2))], + query_time_aggregations: vec![], + }, + ), + ("p90".to_string(), rqe_config(DDSKETCH_AGG, vec![])), + ("distinct".to_string(), rqe_config(HLL_AGG, vec![])), + ("top3".to_string(), rqe_config(TOPK_AGG, vec![])), + ( + "top2_by_job".to_string(), + rqe_config(TOPK_BY_JOB_AGG, vec![]), + ), + ( + "rolled_up".to_string(), + rqe_config(MULTI_SUM_AGG, vec![sum_by_job()]), + ), + ( + // Rolls up to `job` though its leaves group by `instance`. + "misrolled".to_string(), + rqe_config(MULTI_SUM_AGG, vec![sum_by_job()]), + ), + // The keyed sums at their own labels, with no stage. + ("events".to_string(), rqe_config(MULTI_SUM_AGG, vec![])), + ("sql_p90".to_string(), rqe_config(SQL_DDSKETCH_AGG, vec![])), + ("dangling".to_string(), rqe_config(MISSING_AGG, vec![])), + ]), + }; + SimpleEngine::new( + store, + inference_config, + streaming_config, + WINDOW_MS, + QueryLanguage::promql, + ) + } + + fn leaf( + id: &str, + metric: &str, + capability: &str, + lookback_ms: u64, + grouping: &[&str], + ) -> Value { + json!({"rqe": {"id": id, "metric": metric, "capability": capability, + "lookback_ms": lookback_ms, "grouping": grouping}}) + } + + /// `sum_over_time(http_requests[2s])`, or its `sum by` when `grouping` + /// drops a label. + fn sum_leaf(id: &str, lookback_ms: u64, grouping: &[&str]) -> Value { + leaf(id, "http_requests", "sum", lookback_ms, grouping) + } + + fn series_sum() -> Value { + sum_leaf("sum", TWO_WINDOWS_MS, &["instance", "job"]) + } + + fn latest_sum() -> Value { + sum_leaf("latest", WINDOW_MS, &["instance", "job"]) + } + + fn quantile_leaf() -> Value { + let mut quantile = leaf("p90", "latency", "quantile", TWO_WINDOWS_MS, &["job"]); + quantile["rqe"]["phi"] = json!(0.9); + quantile + } + + fn topk_leaf() -> Value { + let mut topk = leaf("top3", "transfers", "topk_by_value", WINDOW_MS, &["srcip"]); + topk["rqe"]["k"] = json!(3); + topk + } + + /// `topk by (job) (2, sum_over_time(flows[1s]))`, or with other buckets. + fn topk_by_job_leaf(topk_by: Option<&[&str]>) -> Value { + let mut topk = leaf( + "top2_by_job", + "flows", + "topk_by_value", + WINDOW_MS, + &["job", "srcip"], + ); + topk["rqe"]["k"] = json!(2); + if let Some(topk_by) = topk_by { + topk["rqe"]["topk_by"] = json!(topk_by); + } + topk + } + + fn expr(tree: Value) -> RqeExpr { + serde_json::from_value(tree).expect("a well-formed RQE tree") + } + + /// Label names and rows, sorted by label values. + fn sorted(labels: KeyByLabelNames, result: QueryResult) -> (Vec, Vec) { + let QueryResult::Vector(vector) = result else { + panic!("expected an instant vector"); + }; + let mut rows: Vec = vector + .values + .into_iter() + .map(|element| (element.labels.labels, element.value)) + .collect(); + rows.sort_by(|left, right| left.0.cmp(&right.0)); + (labels.labels, rows) + } + + fn rqe(engine: &SimpleEngine, tree: Value) -> (Vec, Vec) { + let (labels, result) = engine + .handle_rqe_query(&expr(tree), QUERY_TIME_S) + .expect("the RQE query should be answered"); + sorted(labels, result) + } + + fn promql(engine: &SimpleEngine, query: &str) -> (Vec, Vec) { + let (labels, result) = engine + .handle_query_promql(query.to_string(), QUERY_TIME_S) + .expect("the PromQL query should execute") + .expect("the PromQL query should be served locally"); + sorted(labels, result) + } + + fn rqe_error(engine: &SimpleEngine, tree: Value, time: f64) -> RqeQueryError { + engine + .handle_rqe_query(&expr(tree), time) + .expect_err("the RQE query should be refused") + } + + fn row(labels: &[&str], value: f64) -> Row { + (strings(labels), value) + } + + // --- leaves --- + + #[test] + fn sum_leaf_sums_its_lookback_per_series_without_metric_name() { + assert_eq!( + rqe(&engine(), series_sum()), + ( + strings(&["instance", "job"]), + vec![ + row(&["a", "api"], 3.0), + row(&["a", "web"], 11.0), + row(&["b", "api"], 7.0), + ] + ) + ); + } + + #[test] + fn sum_leaf_equals_promql_sum_over_time() { + let engine = engine(); + assert_eq!( + rqe(&engine, series_sum()), + promql(&engine, "sum_over_time(http_requests[2s])") + ); + } + + #[test] + fn grouped_sum_leaf_equals_promql_sum_by_of_sum_over_time() { + let engine = engine(); + let answer = rqe(&engine, sum_leaf("sum_by_job", TWO_WINDOWS_MS, &["job"])); + assert_eq!( + answer, + ( + strings(&["job"]), + vec![row(&["api"], 10.0), row(&["web"], 11.0)] + ) + ); + assert_eq!( + answer, + promql(&engine, "sum by (job) (sum_over_time(http_requests[2s]))") + ); + } + + #[test] + fn quantile_leaf_equals_promql_quantile_over_time() { + let engine = engine(); + let answer = rqe(&engine, quantile_leaf()); + assert_eq!(answer.0, strings(&["job"])); + assert_eq!(answer.1.len(), 2); + for ((labels, value), expected) in answer.1.iter().zip([900.0, 1_900.0]) { + assert!( + (value - expected).abs() <= 1_000.0 * 0.02, + "{labels:?}: p90 {value} should be near {expected}" + ); + } + assert_eq!( + answer, + promql(&engine, "quantile_over_time(0.9, latency[2s])") + ); + } + + #[test] + fn cardinality_leaf_estimates_distinct_values_per_group() { + let answer = rqe( + &engine(), + leaf("distinct", "users", "cardinality", WINDOW_MS, &["job"]), + ); + assert_eq!( + answer, + ( + strings(&["job"]), + vec![ + row(&["api"], hll(0..1_000).estimate()), + row(&["web"], hll(0..50).estimate()), + ] + ) + ); + } + + #[test] + fn topk_leaf_ranks_globally_and_equals_promql_topk_of_sum_over_time() { + let engine = engine(); + let (labels, result) = engine + .handle_rqe_query(&expr(topk_leaf()), QUERY_TIME_S) + .unwrap(); + let QueryResult::Vector(vector) = &result else { + panic!("expected an instant vector"); + }; + let ranked: Vec = vector + .values + .iter() + .map(|element| (element.labels.labels.clone(), element.value)) + .collect(); + assert_eq!(labels.labels, strings(&["srcip"])); + assert_eq!( + ranked, + vec![ + row(&["10.0.0.5"], 50.0), + row(&["10.0.0.4"], 40.0), + row(&["10.0.0.3"], 30.0), + ], + "top-k rows come back ranked, highest first" + ); + assert_eq!( + sorted(labels, result), + promql(&engine, "topk(3, sum_over_time(transfers[1s]))") + ); + } + + #[test] + fn topk_by_leaf_ranks_within_each_bucket_and_equals_promql_topk_by() { + let engine = engine(); + let answer = rqe(&engine, topk_by_job_leaf(Some(&["job"]))); + assert_eq!( + answer, + ( + strings(&["job", "srcip"]), + vec![ + row(&["api", "10.0.0.2"], 200.0), + row(&["api", "10.0.0.3"], 300.0), + row(&["web", "10.0.0.2"], 20.0), + row(&["web", "10.0.0.3"], 30.0), + ] + ), + "each job keeps its own top 2" + ); + assert_eq!(answer, promql(&engine, TOPK_BY_JOB_QUERY)); + } + + /// The buckets are the aggregation's heaps: a leaf ranking other buckets + /// than its aggregation keeps fails instead of answering another ranking. + #[test] + fn a_topk_leaf_whose_buckets_are_not_its_heaps_fails_loudly() { + let engine = engine(); + for tree in [topk_by_job_leaf(None), topk_by_job_leaf(Some(&["srcip"]))] { + let error = rqe_error(&engine, tree, QUERY_TIME_S); + assert!(matches!(&error, RqeQueryError::BadRequest(_)), "{error:?}"); + assert!(error.to_string().contains("top2_by_job"), "{error}"); + assert!(error.to_string().contains("topk_by"), "{error}"); + } + let mut global = topk_leaf(); + global["rqe"]["topk_by"] = json!(["srcip"]); + let error = rqe_error(&engine, global, QUERY_TIME_S); + assert!(error.to_string().contains("top3"), "{error}"); + } + + #[test] + fn rolled_up_leaf_applies_its_query_time_stage_like_the_promql_plan() { + let engine = engine(); + let answer = rqe( + &engine, + leaf("rolled_up", "events", "sum", TWO_WINDOWS_MS, &["job"]), + ); + assert_eq!( + answer, + ( + strings(&["job"]), + vec![row(&["api"], 7.0), row(&["web"], 8.0)] + ) + ); + assert_eq!( + answer, + promql(&engine, "sum by (job) (sum_over_time(events[2s]))") + ); + } + + #[test] + fn a_roll_up_to_other_labels_than_the_leaf_fails_loudly() { + let error = rqe_error( + &engine(), + leaf("misrolled", "events", "sum", TWO_WINDOWS_MS, &["instance"]), + QUERY_TIME_S, + ); + assert!(matches!(&error, RqeQueryError::BadRequest(_)), "{error:?}"); + assert!(error.to_string().contains("misrolled"), "{error}"); + assert!(error.to_string().contains("regroup to"), "{error}"); + } + + // --- leaves checked against their RQE's aggregation --- + + fn assert_bad_leaf(tree: Value, reason: &str) { + let error = rqe_error(&engine(), tree.clone(), QUERY_TIME_S); + assert!( + matches!(&error, RqeQueryError::BadRequest(_)), + "{tree}: {error:?}" + ); + assert!(error.to_string().contains(reason), "{tree}: {error}"); + } + + #[test] + fn refuses_a_leaf_on_another_metric() { + assert_bad_leaf( + leaf("sum", "other", "sum", TWO_WINDOWS_MS, &["instance", "job"]), + r#"reads "other", but its aggregation 1 is on "http_requests""#, + ); + } + + #[test] + fn refuses_a_capability_its_aggregation_cannot_answer() { + let mut quantile = series_sum(); + quantile["rqe"]["capability"] = json!("quantile"); + quantile["rqe"]["phi"] = json!(0.5); + assert_bad_leaf(quantile, "can't be answered by its aggregation 1"); + let mut distinct = quantile_leaf(); + distinct["rqe"]["capability"] = json!("cardinality"); + distinct["rqe"].as_object_mut().unwrap().remove("phi"); + assert_bad_leaf(distinct, "can't be answered by its aggregation 2"); + } + + #[test] + fn refuses_a_lookback_its_windows_do_not_tile() { + assert_bad_leaf( + sum_leaf("sum", WINDOW_MS + WINDOW_MS / 2, &["instance", "job"]), + "1500 ms lookback is not a whole number of its aggregation 1's 1000 ms windows", + ); + // Any whole number of windows is served. + let three_windows = sum_leaf("sum", 3 * WINDOW_MS, &["instance", "job"]); + engine() + .handle_rqe_query(&expr(three_windows), QUERY_TIME_S) + .unwrap(); + } + + #[test] + fn refuses_a_lookback_longer_than_its_plan_retains() { + let series = ["instance", "job"]; + assert_bad_leaf( + sum_leaf("sum_retains_two", 3 * WINDOW_MS, &series), + "its 3000 ms lookback reads 3 windows of its aggregation 1, but its plan retains 2", + ); + // The retained windows are served. + let engine = engine(); + assert_eq!( + rqe( + &engine, + sum_leaf("sum_retains_two", TWO_WINDOWS_MS, &series) + ), + rqe(&engine, series_sum()) + ); + } + + #[test] + fn refuses_a_topk_k_its_heap_cannot_hold() { + let mut too_many = topk_leaf(); + too_many["rqe"]["k"] = json!(HEAP_SIZE + 1); + assert_bad_leaf( + too_many, + &format!( + "asks for the top 33, but its aggregation {TOPK_AGG} keeps a heap of {HEAP_SIZE}" + ), + ); + // A whole heap is served. + let mut whole_heap = topk_leaf(); + whole_heap["rqe"]["k"] = json!(HEAP_SIZE); + let (_, rows) = rqe(&engine(), whole_heap); + assert_eq!(rows.len(), 5, "every source fits the heap"); + } + + /// Without a stage, a leaf gets the aggregation's labels; a single- + /// population aggregation can also be merged to a subset of its groups. + #[test] + fn refuses_a_grouping_its_aggregation_does_not_serve() { + assert_bad_leaf( + sum_leaf("sum", TWO_WINDOWS_MS, &["instance", "job", "pod"]), + r#"groups by ["instance", "job", "pod"], but its aggregation 1 serves ["instance", "job"]"#, + ); + // Keyed sums keep their keys: they can't drop `instance` without a stage. + assert_bad_leaf( + leaf("events", "events", "sum", TWO_WINDOWS_MS, &["job"]), + r#"but its aggregation 5 serves ["instance", "job"]"#, + ); + let answer = rqe( + &engine(), + leaf( + "events", + "events", + "sum", + TWO_WINDOWS_MS, + &["instance", "job"], + ), + ); + assert_eq!( + answer.1, + vec![ + row(&["a", "api"], 5.0), + row(&["a", "web"], 8.0), + row(&["b", "api"], 2.0), + ] + ); + // A top-k heap ranks whole series: it is never merged to its buckets. + assert_bad_leaf( + { + let mut top = topk_by_job_leaf(Some(&["job"])); + top["rqe"]["grouping"] = json!(["job"]); + top + }, + "but its aggregation 6 serves", + ); + } + + /// A SQL leaf names its table and value column; the aggregation's + /// metric is `table.value_column`. + #[test] + fn reads_a_sql_leaf_by_table_and_value_column() { + let mut sql = leaf("sql_p90", "http", "quantile", TWO_WINDOWS_MS, &["job"]); + sql["rqe"]["phi"] = json!(0.9); + sql["rqe"]["value"] = json!("latency"); + engine() + .handle_rqe_query(&expr(sql.clone()), QUERY_TIME_S) + .expect("checks pass; the window has no data"); + sql["rqe"].as_object_mut().unwrap().remove("value"); + assert_bad_leaf( + sql, + r#"reads "http", but its aggregation 7 is on "http.latency""#, + ); + } + + /// Repeated labels would mislabel series that are matched by position. + #[test] + fn refuses_duplicate_labels() { + assert_bad_leaf( + json!({"op": "+", "lhs": sum_leaf("sum", TWO_WINDOWS_MS, &["instance", "job", "job"]), + "rhs": series_sum()}), + "more than once", + ); + } + + /// Between ingestion ticks a leaf reads what the PromQL instant path + /// reads at the same time. + #[test] + fn misaligned_time_answers_like_the_promql_path() { + let engine = engine(); + let time = QUERY_TIME_S + 0.4; + let (labels, result) = engine.handle_rqe_query(&expr(series_sum()), time).unwrap(); + let promql_rows = engine + .handle_query_promql("sum_over_time(http_requests[2s])".to_string(), time) + .expect("the PromQL query should execute") + .map(|(labels, result)| sorted(labels, result).1) + .unwrap_or_default(); + assert_eq!(sorted(labels, result).1, promql_rows); + } + + // --- agg --- + + #[test] + fn agg_regroups_its_child_by_each_operator() { + let engine = engine(); + for (op, api, web) in [ + ("avg", 5.0, 11.0), + ("min", 3.0, 11.0), + ("max", 7.0, 11.0), + ("sum", 10.0, 11.0), + ] { + assert_eq!( + rqe( + &engine, + json!({"agg": op, "by": ["job"], "of": series_sum()}) + ), + ( + strings(&["job"]), + vec![row(&["api"], api), row(&["web"], web)] + ), + "{op}" + ); + } + } + + #[test] + fn agg_by_nothing_folds_every_series_into_one() { + assert_eq!( + rqe( + &engine(), + json!({"agg": "avg", "by": [], "of": series_sum()}) + ), + (vec![], vec![row(&[], 7.0)]) + ); + } + + // --- binary --- + + #[test] + fn binary_vector_vector_matches_series_by_labels() { + let engine = engine(); + assert_eq!( + rqe( + &engine, + json!({"op": "-", "lhs": series_sum(), "rhs": latest_sum()}) + ), + ( + strings(&["instance", "job"]), + vec![ + row(&["a", "api"], 1.0), + row(&["a", "web"], 5.0), + row(&["b", "api"], 3.0), + ] + ) + ); + assert_eq!( + rqe( + &engine, + json!({"op": "/", "lhs": series_sum(), "rhs": latest_sum()}) + ) + .1, + vec![ + row(&["a", "api"], 1.5), + row(&["a", "web"], 11.0 / 6.0), + row(&["b", "api"], 7.0 / 4.0), + ] + ); + } + + #[test] + fn binary_vector_scalar_keeps_operand_order() { + let engine = engine(); + assert_eq!( + rqe( + &engine, + json!({"op": "*", "lhs": latest_sum(), "rhs": {"scalar": 2.0}}) + ) + .1, + vec![ + row(&["a", "api"], 4.0), + row(&["a", "web"], 12.0), + row(&["b", "api"], 8.0), + ] + ); + assert_eq!( + rqe( + &engine, + json!({"op": "-", "lhs": {"scalar": 10.0}, "rhs": latest_sum()}) + ), + ( + strings(&["instance", "job"]), + vec![ + row(&["a", "api"], 8.0), + row(&["a", "web"], 4.0), + row(&["b", "api"], 6.0), + ] + ) + ); + } + + #[test] + fn nested_tree_aggregates_a_binary_of_leaves() { + let tree = json!({"agg": "max", "by": ["job"], "of": + {"op": "+", "lhs": series_sum(), "rhs": latest_sum()}}); + assert_eq!( + rqe(&engine(), tree), + ( + strings(&["job"]), + vec![row(&["api"], 11.0), row(&["web"], 17.0)] + ) + ); + } + + // --- refusals and empty windows --- + + #[test] + fn refuses_an_invalid_tree() { + let mut quantile = quantile_leaf(); + quantile["rqe"]["phi"] = json!(1.5); + let error = rqe_error(&engine(), quantile, QUERY_TIME_S); + assert!(matches!(error, RqeQueryError::BadRequest(_)), "{error:?}"); + assert!(error.to_string().contains("phi"), "{error}"); + } + + #[test] + fn refuses_an_unknown_rqe_id_before_reading_anything() { + let tree = json!({"op": "+", "lhs": series_sum(), + "rhs": sum_leaf("nope", WINDOW_MS, &["instance", "job"])}); + let error = rqe_error(&engine(), tree, QUERY_TIME_S); + assert!(matches!(error, RqeQueryError::BadRequest(_)), "{error:?}"); + assert!(error.to_string().contains("\"nope\""), "{error}"); + } + + #[test] + fn refuses_a_bad_time() { + for time in [-1.0, f64::NAN, f64::INFINITY] { + let error = rqe_error(&engine(), series_sum(), time); + assert!(matches!(error, RqeQueryError::BadRequest(_)), "{time}"); + } + // A lookback reaching before the epoch. + let error = rqe_error(&engine(), series_sum(), 1.0); + assert!(matches!(error, RqeQueryError::BadRequest(_)), "{error:?}"); + } + + #[test] + fn refuses_trees_without_a_vector() { + let engine = engine(); + let error = rqe_error(&engine, json!({"scalar": 1.0}), QUERY_TIME_S); + assert!(matches!(error, RqeQueryError::BadRequest(_)), "{error:?}"); + let agg_of_scalar = json!({"op": "+", + "lhs": {"agg": "sum", "by": [], "of": {"scalar": 1.0}}, + "rhs": {"agg": "sum", "by": [], "of": series_sum()}}); + let error = rqe_error(&engine, agg_of_scalar, QUERY_TIME_S); + assert!(matches!(error, RqeQueryError::BadRequest(_)), "{error:?}"); + } + + #[test] + fn a_window_without_data_yet_answers_empty() { + let engine = engine(); + let later = QUERY_TIME_S + 3_600.0; + for tree in [ + series_sum(), + topk_leaf(), + json!({"agg": "sum", "by": ["job"], "of": series_sum()}), + json!({"op": "*", "lhs": series_sum(), "rhs": {"scalar": 2.0}}), + ] { + let parsed = expr(tree.clone()); + let (labels, result) = engine.handle_rqe_query(&parsed, later).unwrap(); + let (labels, rows) = sorted(labels, result); + assert!(rows.is_empty(), "{tree}: {rows:?}"); + assert_eq!(labels, parsed.output_labels(), "{tree}"); + } + } + + // --- HTTP route --- + + async fn server() -> (HttpServer, String) { + server_handling(true).await + } + + async fn server_handling(handle_http_requests: bool) -> (HttpServer, String) { + let engine = Arc::new(engine()); + let store: Arc = Arc::new(SimpleMapStore::new( + Arc::new(StreamingConfig::default()), + CleanupPolicy::NoCleanup, + )); + let config = HttpServerConfig { + port: 0, + handle_http_requests, + // Forwarding on, so a test would see any fallback attempt fail. + adapter_config: AdapterConfig::prometheus_promql( + "http://127.0.0.1:1".to_string(), + true, + 30, + ), + }; + let server = HttpServer::new(config, engine, store, None); + let port = server.start_test_server().await.unwrap(); + ( + server, + format!("http://127.0.0.1:{port}{RQE_QUERY_ENDPOINT}"), + ) + } + + async fn post(url: &str, body: String) -> (u16, Value) { + let response = reqwest::Client::new() + .post(url) + .header("content-type", "application/json") + .body(body) + .send() + .await + .unwrap(); + let status = response.status().as_u16(); + (status, response.json().await.unwrap()) + } + + #[tokio::test] + async fn route_answers_prometheus_vector_json() { + let (_server, url) = server().await; + let body = json!({"expr": {"agg": "sum", "by": ["job"], "of": series_sum()}, + "time": QUERY_TIME_S}); + let (status, response) = post(&url, body.to_string()).await; + assert_eq!(status, 200); + assert_eq!(response["status"], "success"); + assert_eq!(response["data"]["resultType"], "vector"); + let mut result = response["data"]["result"].as_array().unwrap().clone(); + result.sort_by_key(|series| series["metric"]["job"].as_str().unwrap().to_string()); + assert_eq!( + result, + vec![ + json!({"metric": {"job": "api"}, "value": [QUERY_TIME_S, "10"]}), + json!({"metric": {"job": "web"}, "value": [QUERY_TIME_S, "11"]}), + ] + ); + } + + #[tokio::test] + async fn route_answers_an_empty_window_with_an_empty_result() { + let (_server, url) = server().await; + let body = json!({"expr": series_sum(), "time": QUERY_TIME_S + 3_600.0}); + let (status, response) = post(&url, body.to_string()).await; + assert_eq!(status, 200); + assert_eq!( + response, + json!({"status": "success", "data": {"resultType": "vector", "result": []}}) + ); + } + + #[tokio::test] + async fn route_refuses_bad_requests_with_400_bad_data() { + let (_server, url) = server().await; + let mut bad_quantile = quantile_leaf(); + bad_quantile["rqe"]["phi"] = json!(2.0); + for (body, reason) in [ + ("not json".to_string(), "invalid RQE query body"), + ( + json!({"expr": series_sum()}).to_string(), + "missing field `time`", + ), + ( + json!({"expr": {"rqe": {"id": "sum"}}, "time": 1.0}).to_string(), + "invalid RQE query body", + ), + ( + json!({"expr": series_sum(), "time": 1.0, "extra": 1}).to_string(), + "unknown field `extra`", + ), + ( + json!({"expr": bad_quantile, "time": QUERY_TIME_S}).to_string(), + "phi", + ), + ( + json!({"expr": sum_leaf("nope", WINDOW_MS, &["job"]), "time": QUERY_TIME_S}) + .to_string(), + "unknown RQE id", + ), + ( + json!({"expr": topk_by_job_leaf(None), "time": QUERY_TIME_S}).to_string(), + "topk_by", + ), + ( + json!({"expr": leaf("misrolled", "events", "sum", TWO_WINDOWS_MS, &["instance"]), + "time": QUERY_TIME_S}) + .to_string(), + "regroup to", + ), + ] { + let (status, response) = post(&url, body.clone()).await; + assert_eq!(status, 400, "{body}"); + assert_eq!(response["status"], "error", "{body}"); + assert_eq!(response["errorType"], "bad_data", "{body}"); + let error = response["error"].as_str().unwrap(); + assert!(error.contains(reason), "{body}: {error}"); + } + } + + /// The engine's own configs are broken: no fault of the request. + #[tokio::test] + async fn route_reports_an_execution_failure_as_500() { + let (_server, url) = server().await; + let body = json!({"expr": leaf("dangling", "events", "sum", TWO_WINDOWS_MS, &["job"]), + "time": QUERY_TIME_S}); + let (status, response) = post(&url, body.to_string()).await; + assert_eq!(status, 500); + assert_eq!(response["status"], "error"); + assert_eq!(response["errorType"], "internal"); + } + + /// RQE queries have no fallback: disabled handling refuses them. + #[tokio::test] + async fn route_refuses_when_handling_is_disabled() { + let (_server, url) = server_handling(false).await; + let body = json!({"expr": series_sum(), "time": QUERY_TIME_S}); + let (status, response) = post(&url, body.to_string()).await; + assert_eq!(status, 500); + assert_eq!(response["status"], "error"); + assert_eq!(response["errorType"], "internal"); + assert_eq!(response["error"], "RQE query handling is disabled"); + } +}