diff --git a/Cargo.lock b/Cargo.lock index 303579a0..5920ce46 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1043,6 +1043,7 @@ dependencies = [ "axum", "clap", "control_plane", + "form_urlencoded", "futures", "http-body-util", "prost", diff --git a/README.md b/README.md index 4794bf52..88909a1d 100644 --- a/README.md +++ b/README.md @@ -35,7 +35,7 @@ The backend is being rebuilt on the Planner's physical DAG: | Data plane: load and validate `plan.json`, accept Remote Write | done | | Stored summaries (SDS): one definition per computation, panes by end time | done | | Precompute engine: build ingestion-time outputs once per tumbling pane | done | -| Query engine: run query-time nodes over stored panes and raw data | planned | +| Query engine: answer planned instant queries from stored panes and raw samples | done | The data plane stores the plan's `stored_outputs`. Outputs that differ only in which pane they read share one stored summary, kept for the largest @@ -53,9 +53,13 @@ are dropped. Remote Write gets 503 when admitting it would leave more than stopped, which also makes `/-/ready` fail. `/api/v1/status/asap` reports the engine's sample and pane counters as JSON. -Until the query engine lands, the data plane forwards every query to -Prometheus, with its end-to-end headers such as `Authorization` and -`X-Scope-OrgID`, and marks the response `x-asap-execution: forwarded`. SQL and +An instant query whose text is one of the plan's queries is answered by the +plan (`x-asap-execution: plan`). Its stored inputs read the panes ending at +`T - offset`, so `T` must be on the pane grid and those panes stored; its +query-time bound raw-sample windows are read from Prometheus. Anything else is forwarded +to Prometheus with its end-to-end headers, such as `Authorization` and +`X-Scope-OrgID` (`x-asap-execution: forwarded`); a planned query that could not +be answered also says why in `x-asap-forward-reason`. Range queries, SQL and ClickHouse are not supported yet. ## Running @@ -71,8 +75,8 @@ cargo run -p data_plane -- --plan plan.json --prometheus-url http://localhost:90 `--deployment-capabilities caps.json` replaces the default `DeploymentCapabilities`: every summary the Planner can build, raw data kept -in Prometheus, and no query-time result kept across evaluations -(`MatAtQueryKeptAcrossEvaluations` is not implemented). For example +in Prometheus, and no stored outputs built in a `QueryRun` +(`MatAtQueryKeptAcrossEvaluations`). For example `{"raw_data_retained": false}` makes the Planner build every input at ingestion time. `--measured-operation-costs costs.json` gives Stage 3 the deployment's measured sketch operation costs; otherwise it prices every node diff --git a/control_plane/src/lib.rs b/control_plane/src/lib.rs index 7767e805..ba1ee7a5 100644 --- a/control_plane/src/lib.rs +++ b/control_plane/src/lib.rs @@ -83,7 +83,7 @@ pub fn load_plan(path: &Path) -> Result { /// What this backend's data plane can execute and keep: every summary the /// Planner executor builds, raw samples kept in Prometheus, and no -/// query-time result kept across evaluations: the SDS describes such +/// stored outputs built in a `QueryRun` (`MatAtQueryKeptAcrossEvaluations`): the SDS describes such /// outputs, but nothing builds them until the query engine does. pub fn backend_capabilities() -> DeploymentCapabilities { DeploymentCapabilities { diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 5914ecd0..8ce0da62 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -20,8 +20,10 @@ reqwest.workspace = true prost = "0.13" snap = "1.1" futures = "0.3" +form_urlencoded = "1" [dev-dependencies] control_plane.workspace = true tower = { version = "0.5", features = ["util"] } http-body-util = "0.1" +serde_json.workspace = true diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 21f6d082..df882c5a 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -2,13 +2,17 @@ //! it. Planner IR holds `Rc`s, so the plan cannot be shared across the HTTP //! server's threads; handlers send the engine commands instead. +use std::collections::BTreeMap; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc::{channel, Receiver, Sender}; use std::sync::{Arc, Mutex}; use asap_planner::InstallablePlan; +use tokio::sync::oneshot; + use crate::precompute::{Precompute, PrecomputeOptions, PrecomputeStats, RawSample}; +use crate::query::{BoundRawSampleWindow, QueryEngine, ResultSample, Unanswered}; use crate::sds::Sds; /// How the engine thread ingests samples. @@ -31,8 +35,21 @@ impl Default for EngineOptions { } } +type Reply = oneshot::Sender>; + enum Command { Ingest(Vec), + BoundRawSampleWindows { + query: String, + time_ms: i64, + reply: Reply<(usize, Vec)>, + }, + Evaluate { + root: usize, + time_ms: i64, + raw: BTreeMap>, + reply: Reply>, + }, Stop, } @@ -135,6 +152,44 @@ impl EngineHandle { }) } + /// The plan root answering `query` at `time_ms` and the raw samples it + /// needs; see [`QueryEngine::bound_raw_sample_windows`]. + pub async fn bound_raw_sample_windows( + &self, + query: String, + time_ms: i64, + ) -> Result<(usize, Vec), Unanswered> { + self.ask(|reply| Command::BoundRawSampleWindows { + query, + time_ms, + reply, + }) + .await + } + + /// Runs a root over stored panes and `raw`; see [`QueryEngine::evaluate`]. + pub async fn evaluate( + &self, + root: usize, + time_ms: i64, + raw: BTreeMap>, + ) -> Result, Unanswered> { + self.ask(|reply| Command::Evaluate { + root, + time_ms, + raw, + reply, + }) + .await + } + + async fn ask(&self, command: impl FnOnce(Reply) -> Command) -> Result { + let stopped = || Unanswered::Execution("the engine has stopped".into()); + let (reply, answer) = oneshot::channel(); + self.commands.send(command(reply)).map_err(|_| stopped())?; + answer.await.map_err(|_| stopped())? + } + /// Whether the engine thread is still running. pub fn is_running(&self) -> bool { self.running.load(Ordering::SeqCst) @@ -181,6 +236,7 @@ struct Shared { struct Engine { precompute: Precompute, + query: QueryEngine, } impl Engine { @@ -193,9 +249,23 @@ impl Engine { outputs = sds.outputs.len(), "stored summaries" ); + let query_run = sds + .definitions + .values() + .filter(|d| d.built_at == asap_planner::StoredOutputBuild::QueryRun) + .count(); + if query_run > 0 { + tracing::warn!( + definitions = query_run, + "the plan has stored outputs built in a QueryRun, which this data plane does not build yet; \ + queries that read them are forwarded" + ); + } let queries = plan.queries.iter().map(|q| q.query.clone()).collect(); + let query = QueryEngine::new(&plan, &sds); let engine = Self { precompute: Precompute::new(plan.dag.dag, &sds, options)?, + query, }; Ok((engine, queries)) } @@ -210,6 +280,30 @@ impl Engine { .queued_samples .fetch_sub(samples.len() as u64, Ordering::SeqCst); } + Command::BoundRawSampleWindows { + query, + time_ms, + reply, + } => { + let _ = reply.send(self.query.bound_raw_sample_windows( + &query, + time_ms, + &self.precompute.store, + )); + } + Command::Evaluate { + root, + time_ms, + raw, + reply, + } => { + let _ = reply.send(self.query.evaluate( + root, + time_ms, + &self.precompute.store, + &raw, + )); + } Command::Stop => break, } } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index fe037eb7..b65ac699 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -2,11 +2,14 @@ //! the Prometheus query API for an installed plan. //! //! The installed plan is ASAPPlanner's selected `PhysicalASAPDAG`. The -//! engine thread builds its ingestion-time outputs pane by pane; queries are -//! still forwarded to Prometheus. +//! engine thread builds its ingestion-time outputs pane by pane and answers +//! the plan's instant queries from them; every other query, and any planned +//! query it cannot answer, is forwarded to Prometheus. pub mod engine; pub mod precompute; +pub mod prometheus; +pub mod query; pub mod remote_write; pub mod sds; @@ -23,9 +26,13 @@ use axum::routing::{get, post}; use axum::Router; use engine::{EngineHandle, EngineOptions, IngestError}; use precompute::RawSample; +use prometheus::InstantQuery; +use query::Unanswered; -/// Response header naming what answered a query. +/// Response header naming what answered a query: `plan` or `forwarded`. pub const EXECUTION_HEADER: &str = "x-asap-execution"; +/// Response header of a forwarded planned query: why the plan did not answer. +pub const FORWARD_REASON_HEADER: &str = "x-asap-forward-reason"; /// How long a forwarded query may take by default, matching Prometheus's /// default `--query.timeout` of two minutes. @@ -177,6 +184,23 @@ async fn status(State(state): State>) -> Response { .into_response() } +/// `text` as a header value: characters a header cannot carry (controls and +/// non-ASCII) become `?`. +fn header_value(text: &str) -> HeaderValue { + let visible: String = text + .chars() + .map(|c| { + if c == ' ' || c.is_ascii_graphic() { + c + } else { + '?' + } + }) + .collect(); + HeaderValue::from_str(&visible).expect("visible ASCII is a valid header value") +} + +/// Answers a planned instant query from the plan; forwards anything else. async fn query( state: State>, method: Method, @@ -184,7 +208,67 @@ async fn query( headers: HeaderMap, body: Bytes, ) -> Response { - forward(state, "query", method, query, headers, body).await + let instant = match InstantQuery::parse(query.0.as_deref(), &body) { + Ok(Some(instant)) => instant, + Ok(None) => return forward(state, "query", method, query, headers, body).await, + Err(bad) => { + let body = axum::Json(prometheus::bad_data(&bad.to_string())); + return (StatusCode::BAD_REQUEST, body).into_response(); + } + }; + let reason = match answer(&state, instant, &headers).await { + Ok(response) => { + let mut response = axum::Json(response).into_response(); + response + .headers_mut() + .insert(EXECUTION_HEADER, HeaderValue::from_static("plan")); + return response; + } + Err(Unanswered::NotPlanned(_)) => None, + Err(reason) => { + tracing::debug!(%reason, "forwarding a planned query"); + Some(header_value(&reason.to_string())) + } + }; + let mut response = forward(state, "query", method, query, headers, body).await; + if let Some(reason) = reason { + response.headers_mut().insert(FORWARD_REASON_HEADER, reason); + } + response +} + +/// Runs a planned instant query. Bound raw-sample windows are read from Prometheus with +/// the caller's end-to-end `headers`, as a forwarded query would be. +async fn answer( + state: &DataPlane, + query: InstantQuery, + headers: &HeaderMap, +) -> Result { + let time_ms = query.time_ms.unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| d.as_millis() as i64) + }); + let (root, ranges) = state + .engine + .bound_raw_sample_windows(query.query, time_ms) + .await?; + let mut fetch_headers = forwarded_headers(headers); + // The raw read is a GET without the caller's body. + fetch_headers.remove(header::CONTENT_TYPE); + // The bound raw-sample windows are fetched concurrently. + let fetches = ranges.iter().map(|range| async { + prometheus::fetch(&state.client, &state.prometheus_url, &fetch_headers, range) + .await + .map(|samples| (range.node, samples)) + }); + let raw: BTreeMap> = futures::future::try_join_all(fetches) + .await + .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))? + .into_iter() + .collect(); + let samples = state.engine.evaluate(root, time_ms, raw).await?; + Ok(prometheus::vector_response(time_ms, &samples)) } async fn query_range( @@ -288,3 +372,19 @@ async fn forward( .insert(EXECUTION_HEADER, HeaderValue::from_static("forwarded")); response } + +#[cfg(test)] +mod tests { + use super::*; + + /// A forward reason with characters a header cannot carry is still sent, + /// with those characters replaced. + #[test] + fn sanitizes_header_values() { + assert_eq!( + header_value("pane ending at 10 is missing"), + "pane ending at 10 is missing" + ); + assert_eq!(header_value("bad\nlabel é"), "bad?label ?"); + } +} diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index 9658a4fe..3cdedce2 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -141,7 +141,7 @@ impl Precompute { options: PrecomputeOptions, ) -> Result { let mut builders = Vec::new(); - // Ingestion-run definitions in dependency order; query-run ones are + // `IngestionRun` definitions in dependency order; `QueryRun` ones are // the query engine's. let order = sds .ingestion_build_order() @@ -353,17 +353,9 @@ fn build_pane( window_end_ms: end, revision: 1, }; - // The selector takes owned label maps, so each sample's labels are - // copied once per pane build. - let samples: Vec = samples - .iter() - .map(|s| Sample { - labels: (*s.labels).clone(), - timestamp_ms: s.timestamp_ms, - value: s.value, - }) - .collect(); - let selected = builder.selector.window_batch(end, &samples)?; + let selected = builder + .selector + .window_batch(end, &executor_samples(samples))?; let input: Source<'_> = Box::new(Operator::source( builder.raw_schema.clone(), vec![selected], @@ -376,6 +368,19 @@ fn build_pane( Ok(Batch::try_new(schema, rows)?) } +/// `samples` as the executor's [`Sample`]s. The selector takes owned label +/// maps, so each sample's labels are copied. +pub fn executor_samples(samples: &[RawSample]) -> Vec { + samples + .iter() + .map(|s| Sample { + labels: (*s.labels).clone(), + timestamp_ms: s.timestamp_ms, + value: s.value, + }) + .collect() +} + /// Executes `root` of `dag` to completion and returns its rows. pub fn run_root( dag: &asap_executor::plan::PhysicalDAG<'_, Batch, SchemaRef>, diff --git a/data_plane/src/prometheus.rs b/data_plane/src/prometheus.rs new file mode 100644 index 00000000..2bbd90dd --- /dev/null +++ b/data_plane/src/prometheus.rs @@ -0,0 +1,267 @@ +//! The Prometheus HTTP API as the data plane speaks it: instant-query +//! parameters and results, and fetching raw samples for query-time ranges. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use serde_json::{json, Value}; + +use crate::precompute::RawSample; +use crate::query::{BoundRawSampleWindow, ResultSample}; + +/// Parameters of an instant query, from its URL query and form body. +pub struct InstantQuery { + pub query: String, + /// `None` when the request names no time: evaluate now. + pub time_ms: Option, +} + +/// The largest evaluation time, in milliseconds either side of the epoch, the +/// data plane answers; beyond it a request is refused as Prometheus refuses +/// a time it cannot represent. Leaves room for any window offset. +pub const MAX_TIME_MS: i64 = i64::MAX / 4; + +/// Why an instant query's parameters are refused (400, `bad_data`). +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +#[error("invalid parameter \"time\": {0}")] +pub struct BadTime(pub String); + +impl InstantQuery { + /// Reads `query` and `time` (Unix seconds). `Ok(None)` if either is + /// missing or `time` is in another format, such as RFC 3339 (the query + /// is then forwarded); `Err` if `time` is a number that is not finite or + /// beyond [`MAX_TIME_MS`]. + pub fn parse(url_query: Option<&str>, form: &[u8]) -> Result, BadTime> { + let params: BTreeMap = + form_urlencoded::parse(url_query.unwrap_or_default().as_bytes()) + .chain(form_urlencoded::parse(form)) + .map(|(k, v)| (k.into_owned(), v.into_owned())) + .collect(); + let time_ms = match params.get("time") { + None => None, + Some(time) => match time.parse::() { + Err(_) => return Ok(None), + Ok(seconds) => Some(checked_seconds_to_ms(seconds).ok_or_else(|| { + BadTime(format!( + "{time} is not a time the data plane can evaluate at" + )) + })?), + }, + }; + let Some(query) = params.get("query") else { + return Ok(None); + }; + Ok(Some(Self { + query: query.clone(), + time_ms, + })) + } +} + +/// `seconds` in milliseconds, if finite and within [`MAX_TIME_MS`]. +fn checked_seconds_to_ms(seconds: f64) -> Option { + let ms = (seconds * 1000.0).round(); + (ms.is_finite() && ms.abs() <= MAX_TIME_MS as f64).then_some(ms as i64) +} + +/// A Prometheus `bad_data` error response body. +pub fn bad_data(error: &str) -> Value { + json!({"status": "error", "errorType": "bad_data", "error": error}) +} + +fn seconds_to_ms(seconds: f64) -> i64 { + (seconds * 1000.0).round() as i64 +} + +fn ms_to_seconds(ms: i64) -> String { + format!("{:.3}", ms as f64 / 1000.0) +} + +/// A Prometheus instant-query response for `samples` at `time_ms`. +pub fn vector_response(time_ms: i64, samples: &[ResultSample]) -> Value { + let time = time_ms as f64 / 1000.0; + let result: Vec = samples + .iter() + .map(|s| json!({"metric": s.labels, "value": [time, format_value(s.value)]})) + .collect(); + json!({"status": "success", "data": {"resultType": "vector", "result": result}}) +} + +/// Prometheus' sample value format. +fn format_value(value: f64) -> String { + match value { + v if v.is_nan() => "NaN".into(), + v if v == f64::INFINITY => "+Inf".into(), + v if v == f64::NEG_INFINITY => "-Inf".into(), + v => v.to_string(), + } +} + +/// The PromQL range selector of `range`: its metric and label matchers, every +/// value escaped as a PromQL string, over `end - start` milliseconds. +pub fn range_selector(range: &BoundRawSampleWindow) -> String { + let mut matchers = vec![format!("__name__=\"{}\"", escape(&range.metric))]; + matchers.extend( + range + .matchers + .iter() + .map(|m| format!("{}{}\"{}\"", m.name, m.op.as_str(), escape(&m.value))), + ); + format!( + "{{{}}}[{}ms]", + matchers.join(","), + range.end_ms.saturating_sub(range.start_ms) + ) +} + +/// `text` as the inside of a double-quoted PromQL string. +fn escape(text: &str) -> String { + let mut escaped = String::with_capacity(text.len()); + for c in text.chars() { + match c { + '\\' => escaped.push_str("\\\\"), + '"' => escaped.push_str("\\\""), + '\n' => escaped.push_str("\\n"), + c => escaped.push(c), + } + } + escaped +} + +/// The samples of `range`'s metric in `(start, end]`, read from Prometheus +/// with a range selector evaluated at the range's end. `headers` are sent +/// with the request: the caller's end-to-end headers, such as +/// `Authorization` and `X-Scope-OrgID`, so Prometheus authorizes and scopes +/// the read as it would the caller's own query. +pub async fn fetch( + client: &reqwest::Client, + prometheus_url: &str, + headers: &reqwest::header::HeaderMap, + range: &BoundRawSampleWindow, +) -> anyhow::Result> { + let selector = range_selector(range); + let response: Value = client + .get(format!("{prometheus_url}/api/v1/query")) + .headers(headers.clone()) + .query(&[ + ("query", selector.as_str()), + ("time", &ms_to_seconds(range.end_ms)), + ]) + .send() + .await? + .error_for_status()? + .json() + .await?; + let malformed = || anyhow::anyhow!("unexpected Prometheus response: {response}"); + let series = response["data"]["result"] + .as_array() + .ok_or_else(malformed)?; + let mut samples = Vec::new(); + for series in series { + let labels: BTreeMap = serde_json::from_value(series["metric"].clone())?; + let labels = Arc::new(labels); + for point in series["values"].as_array().ok_or_else(malformed)? { + let seconds = point[0].as_f64().ok_or_else(malformed)?; + let value = point[1].as_str().ok_or_else(malformed)?; + let timestamp_ms = seconds_to_ms(seconds); + // Prometheus 2 ranges include their start; the plan's do not. + if timestamp_ms <= range.start_ms { + continue; + } + samples.push(RawSample { + labels: labels.clone(), + timestamp_ms, + value: parse_value(value).ok_or_else(malformed)?, + }); + } + } + Ok(samples) +} + +fn parse_value(value: &str) -> Option { + match value { + "+Inf" | "Inf" => Some(f64::INFINITY), + "-Inf" => Some(f64::NEG_INFINITY), + "NaN" => Some(f64::NAN), + v => v.parse().ok(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Parameters come from the URL or the form body; times are Unix seconds. + #[test] + fn parses_instant_query_parameters() { + let q = InstantQuery::parse(Some("query=up&time=1.5"), b"") + .unwrap() + .unwrap(); + assert_eq!((q.query.as_str(), q.time_ms), ("up", Some(1500))); + let q = InstantQuery::parse(None, b"query=sum%28x%29") + .unwrap() + .unwrap(); + assert_eq!((q.query.as_str(), q.time_ms), ("sum(x)", None)); + assert!(InstantQuery::parse(Some("time=1"), b"").unwrap().is_none()); + assert!( + InstantQuery::parse(Some("query=up&time=2026-01-01T00:00:00Z"), b"") + .unwrap() + .is_none() + ); + for time in ["-1e300", "1e300", "NaN", "inf"] { + let url = format!("query=up&time={time}"); + assert!(InstantQuery::parse(Some(&url), b"").is_err(), "{time}"); + } + } + + /// A range selector carries the metric and the window's label matchers, + /// each value escaped as a PromQL string. + #[test] + fn builds_escaped_range_selectors() { + let range = BoundRawSampleWindow { + node: 1, + metric: "m".into(), + matchers: vec![ + crate::query::LabelMatcher { + name: "job".into(), + op: crate::query::MatchOp::Equal, + value: "a\"b\\c".into(), + }, + crate::query::LabelMatcher { + name: "env".into(), + op: crate::query::MatchOp::NotRegex, + value: "dev|test".into(), + }, + ], + start_ms: 0, + end_ms: 60_000, + }; + assert_eq!( + range_selector(&range), + r#"{__name__="m",job="a\"b\\c",env!~"dev|test"}[60000ms]"# + ); + } + + /// Results use Prometheus' vector shape and value formatting. + #[test] + fn formats_vector_results() { + let samples = [ + ResultSample { + labels: BTreeMap::from([("job".into(), "a".into())]), + value: 2.0, + }, + ResultSample { + labels: BTreeMap::new(), + value: f64::INFINITY, + }, + ]; + assert_eq!( + vector_response(1500, &samples), + json!({"status": "success", "data": {"resultType": "vector", "result": [ + {"metric": {"job": "a"}, "value": [1.5, "2"]}, + {"metric": {}, "value": [1.5, "+Inf"]}, + ]}}) + ); + assert_eq!(ms_to_seconds(4_200_001), "4200.001"); + } +} diff --git a/data_plane/src/query.rs b/data_plane/src/query.rs new file mode 100644 index 00000000..d73012a8 --- /dev/null +++ b/data_plane/src/query.rs @@ -0,0 +1,409 @@ +//! Query engine: answers a planned query at evaluation time `T` by running +//! its root of the plan's DAG over stored panes and raw samples. +//! +//! Each root is compiled once with two kinds of inputs: +//! - stored outputs (the SDS frontier): bound to the pane of their definition +//! ending at `T - offset`, which must be stored; +//! - query-time bound raw-sample windows, `TimeRange(range) <- TimeShift(offset)* <- Scan`: +//! bound to the samples in `(T - offset - range, T - offset]`, which the +//! caller fetches, filtered by the plan's `Scan`. +//! +//! A query the engine cannot answer returns [`Unanswered`]; the caller then +//! forwards it to Prometheus. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; + +use asap_executor::operators::Operator; +use asap_executor::physical_planner::{compile, CompiledPhysicalDAG, InputContract, Source}; +use asap_executor::planner::ir::physical_export::{PhysicalASAPDAG, RawSampleWindow}; +use asap_executor::planner::ir::scalar::{CompareOpKind, ScalarValue}; +use asap_executor::planner::ir::schema::{DataType, FieldDataType, Schema}; +use asap_executor::planner::ir::{NonASAPOp, Operator as IrOperator, ScalarExpr}; +use asap_executor::runtime::Scope; +use asap_executor::values::{Batch, SchemaRef, Value}; +use asap_executor::{series_labels_of_output_row, RawSampleSelector}; +use asap_planner::InstallablePlan; + +use crate::precompute::{executor_samples, run_root, RawSample}; +use crate::sds::{DefinitionId, Sds, Selection, SummaryStore}; + +/// A bound raw-sample window: a `RawSampleWindow` of the plan bound to one +/// evaluation time, i.e. `metric`'s raw samples in `(start_ms, end_ms]` +/// that `matchers` keep, which the caller fetches. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BoundRawSampleWindow { + pub node: u64, + pub metric: String, + /// The window's `Scan` label matchers that PromQL can express, so the + /// fetch reads only the series the plan keeps. The selector applies + /// every matcher again, so one left out only widens the fetch. + pub matchers: Vec, + pub start_ms: i64, + pub end_ms: i64, +} + +/// A PromQL label matcher: `name op "value"`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LabelMatcher { + pub name: String, + pub op: MatchOp, + pub value: String, +} + +/// A PromQL label matching operator. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MatchOp { + /// `=` + Equal, + /// `!=` + NotEqual, + /// `=~` + Regex, + /// `!~` + NotRegex, +} + +impl MatchOp { + pub fn as_str(self) -> &'static str { + match self { + MatchOp::Equal => "=", + MatchOp::NotEqual => "!=", + MatchOp::Regex => "=~", + MatchOp::NotRegex => "!~", + } + } +} + +/// One series of an instant-query result. +#[derive(Debug, Clone, PartialEq)] +pub struct ResultSample { + pub labels: BTreeMap, + pub value: f64, +} + +/// Why the engine does not answer a query; it is then forwarded. +#[derive(Debug, Clone, PartialEq, thiserror::Error)] +pub enum Unanswered { + #[error("the installed plan has no query {0:?}")] + NotPlanned(String), + #[error("the data plane cannot run this query: {0}")] + Unsupported(String), + #[error("{0}")] + MissingPane(String), + #[error("execution failed: {0}")] + Execution(String), +} + +struct StoredInput { + node: u64, + definition: DefinitionId, + offset_ms: i64, + pane_ms: i64, +} + +struct RawInput { + node: u64, + matchers: Vec, + /// The plan's `Scan` label matchers over the range's samples. + selector: RawSampleSelector, + schema: SchemaRef, + range_ms: i64, + offset_ms: i64, +} + +struct Program { + compiled: CompiledPhysicalDAG, + stored: Vec, + raw: Vec, + output: Schema, +} + +pub struct QueryEngine { + /// Each query, with its compiled root or why it cannot run. + roots: Vec<(String, Result)>, +} + +impl QueryEngine { + pub fn new(plan: &InstallablePlan, sds: &Sds) -> Self { + let dag = &plan.dag.dag; + let windows: HashMap = plan + .query_time_raw_windows + .iter() + .map(|w| (w.time_range_node_id as u64, w)) + .collect(); + let roots = plan + .queries + .iter() + .map(|q| { + let program = program(dag, sds, &windows, q.root_node_id as u64); + (q.query.trim().to_string(), program) + }) + .collect(); + Self { roots } + } + + /// The root answering `query` at `time_ms`, and the raw samples it needs. + /// Fails if the query is not planned or a stored pane it reads is missing. + pub fn bound_raw_sample_windows( + &self, + query: &str, + time_ms: i64, + store: &SummaryStore, + ) -> Result<(usize, Vec), Unanswered> { + let (root, program) = self.program(query)?; + for input in &program.stored { + pane(input, time_ms, store)?; + } + let ranges = program + .raw + .iter() + .map(|input| { + let (start_ms, end_ms) = window(time_ms, input.offset_ms, input.range_ms)?; + Ok(BoundRawSampleWindow { + node: input.node, + metric: input.selector.metric().to_string(), + matchers: input.matchers.clone(), + start_ms, + end_ms, + }) + }) + .collect::>()?; + Ok((root, ranges)) + } + + /// Runs `root` at `time_ms` over stored panes and `raw`, the samples of + /// each of its [`BoundRawSampleWindow`]s by node. + pub fn evaluate( + &self, + root: usize, + time_ms: i64, + store: &SummaryStore, + raw: &BTreeMap>, + ) -> Result, Unanswered> { + let program = self.roots[root] + .1 + .as_ref() + .map_err(|reason| Unanswered::Unsupported(reason.clone()))?; + let failed = |e: asap_executor::Error| Unanswered::Execution(e.to_string()); + let scope = Scope::Query { + evaluation_time_ms: time_ms, + revision: 1, + }; + let mut sources: BTreeMap> = BTreeMap::new(); + for input in &program.stored { + let pane = pane(input, time_ms, store)?; + let source = + Operator::source(pane.schema().clone(), vec![pane.clone()]).map_err(failed)?; + sources.insert(input.node, Box::new(source)); + } + for input in &program.raw { + let samples = raw.get(&input.node).map(Vec::as_slice).unwrap_or_default(); + let rows = input + .selector + .window_batch( + window(time_ms, input.offset_ms, input.range_ms)?.1, + &executor_samples(samples), + ) + .map_err(failed)?; + let source = Operator::source(input.schema.clone(), vec![rows]).map_err(failed)?; + sources.insert(input.node, Box::new(source)); + } + let dag = program.compiled.instantiate(sources).map_err(failed)?; + let rows = run_root(&dag, program.compiled.roots()[0], scope).map_err(failed)?; + result_samples(&program.output, rows) + } + + fn program(&self, query: &str) -> Result<(usize, &Program), Unanswered> { + let query = query.trim(); + let root = self + .roots + .iter() + .position(|(planned, _)| planned == query) + .ok_or_else(|| Unanswered::NotPlanned(query.to_string()))?; + match &self.roots[root].1 { + Ok(program) => Ok((root, program)), + Err(reason) => Err(Unanswered::Unsupported(reason.clone())), + } + } +} + +/// `(start, end]` of a window `width_ms` wide ending `offset_ms` before +/// `time_ms`, or why it does not fit in `i64` milliseconds. +fn window(time_ms: i64, offset_ms: i64, width_ms: i64) -> Result<(i64, i64), Unanswered> { + let out_of_range = + || Unanswered::Unsupported(format!("evaluation time {time_ms} ms is out of range")); + let end = time_ms.checked_sub(offset_ms).ok_or_else(out_of_range)?; + let start = end.checked_sub(width_ms).ok_or_else(out_of_range)?; + Ok((start, end)) +} + +/// The stored pane `input` reads at `time_ms`. +fn pane<'a>( + input: &StoredInput, + time_ms: i64, + store: &'a SummaryStore, +) -> Result<&'a Batch, Unanswered> { + let (_, end) = window(time_ms, input.offset_ms, input.pane_ms)?; + if end.rem_euclid(input.pane_ms) != 0 { + return Err(Unanswered::MissingPane(format!( + "{time_ms} is not on the {} ms pane grid", + input.pane_ms + ))); + } + store + .get(&input.definition, Selection::ending_at(end)) + .ok_or_else(|| { + Unanswered::MissingPane(format!( + "pane ending at {end} of {:?} is not stored", + input.definition + )) + }) +} + +/// Compiles `root` with its stored and raw inputs, or says why it cannot. +fn program( + dag: &PhysicalASAPDAG, + sds: &Sds, + windows: &HashMap, + root: u64, +) -> Result { + let nodes: HashMap = dag.nodes.iter().map(|n| (n.id as u64, n)).collect(); + let mut producers: HashMap> = HashMap::new(); + for edge in &dag.edges { + producers + .entry(edge.consumer as u64) + .or_default() + .push(edge.producer as u64); + } + let (mut stored, mut raw) = (Vec::new(), Vec::new()); + let mut pending = vec![root]; + let mut seen = BTreeSet::new(); + while let Some(id) = pending.pop() { + if !seen.insert(id) { + continue; + } + if let Some((definition, binding)) = sds.outputs.get(&id) { + stored.push(StoredInput { + node: id, + definition: definition.clone(), + offset_ms: binding.offset_ms, + pane_ms: sds.definitions[definition].pane_ms, + }); + continue; + } + match &nodes[&id].payload { + IrOperator::NonASAP(NonASAPOp::TimeRange { .. }) => { + let window = windows.get(&id).ok_or_else(|| { + format!("node {id}: a TimeRange the plan lists no window for") + })?; + let selector = + RawSampleSelector::new(dag, window).map_err(|e| format!("node {id}: {e}"))?; + raw.push(RawInput { + node: id, + matchers: label_matchers(dag, window.scan_node_id as u64), + selector, + schema: Arc::new(nodes[&id].output_schema.clone()), + range_ms: window.width_ms, + offset_ms: window.end_offset_ms, + }); + continue; + } + IrOperator::NonASAP(NonASAPOp::Scan { .. }) => { + return Err(format!("node {id}: a Scan read without a TimeRange")) + } + _ => pending.extend(producers.get(&id).into_iter().flatten().copied()), + } + } + let mut inputs = BTreeMap::new(); + for input in &stored { + let schema = Arc::new(nodes[&input.node].output_schema.clone()); + inputs.insert(input.node, InputContract::bounded(schema)); + } + for input in &raw { + inputs.insert(input.node, InputContract::bounded(input.schema.clone())); + } + let compiled = compile(dag, inputs, &[root]).map_err(|e| e.to_string())?; + Ok(Program { + compiled, + stored, + raw, + output: nodes[&root].output_schema.clone(), + }) +} + +/// The label matchers of the `Scan` at `scan` that PromQL can express: a +/// label column compared (`=`, `!=`, `=~`, `!~`) with a text literal. +/// Other predicates are left out of the fetch; the selector applies them. +fn label_matchers(dag: &PhysicalASAPDAG, scan: u64) -> Vec { + let Some(node) = dag.nodes.iter().find(|n| n.id as u64 == scan) else { + return Vec::new(); + }; + let IrOperator::NonASAP(NonASAPOp::Scan { + predicates, schema, .. + }) = &node.payload + else { + return Vec::new(); + }; + predicates + .iter() + .filter_map(|predicate| { + let ScalarExpr::Compare { + left, op, right, .. + } = &predicate.0 + else { + return None; + }; + let (ScalarExpr::Column(column), ScalarExpr::Literal(ScalarValue::Utf8(value))) = + (left.as_ref(), right.as_ref()) + else { + return None; + }; + let op = match op { + CompareOpKind::Eq => MatchOp::Equal, + CompareOpKind::Ne => MatchOp::NotEqual, + CompareOpKind::Regex => MatchOp::Regex, + CompareOpKind::NotRegex => MatchOp::NotRegex, + _ => return None, + }; + let name = &schema.fields.get(*column)?.name; + // Only Prometheus label names; the series identity column and + // other derived columns are not labels. + let is_label = !name.is_empty() + && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') + && !name.starts_with(|c: char| c.is_ascii_digit()); + is_label.then(|| LabelMatcher { + name: name.clone(), + op, + value: value.clone(), + }) + }) + .collect() +} + +/// PromQL vector samples from result rows: labels from the series identity +/// and the other text columns, the value from the one `Float64` column. +fn result_samples(schema: &Schema, rows: Vec>) -> Result, Unanswered> { + let unsupported = |reason: &str| Unanswered::Unsupported(format!("result rows: {reason}")); + let values: Vec = schema + .fields + .iter() + .enumerate() + .filter(|(_, f)| f.dtype == FieldDataType::Plain(DataType::Float64)) + .map(|(i, _)| i) + .collect(); + let [value] = values[..] else { + return Err(unsupported("expected one Float64 value column")); + }; + rows.into_iter() + .map(|row| { + let labels = series_labels_of_output_row(schema, &row) + .map_err(|e| Unanswered::Execution(e.to_string()))?; + let value = match &row[value] { + Value::Float64(v) => *v, + other => return Err(Unanswered::Execution(format!("value {other:?}"))), + }; + Ok(ResultSample { labels, value }) + }) + .collect() +} diff --git a/data_plane/src/sds.rs b/data_plane/src/sds.rs index b5354310..184f5f2c 100644 --- a/data_plane/src/sds.rs +++ b/data_plane/src/sds.rs @@ -43,9 +43,10 @@ pub struct DefinitionId(pub String); #[derive(Debug, Clone, PartialEq)] pub struct SummaryDefinition { pub id: DefinitionId, - /// Which run builds its summaries: an ingestion run, or the query run of - /// the first evaluation that needs one (kept once its coverage is - /// closed). Ingestion when any output of the definition says so. + /// Which run builds its stored outputs (`StoredOutputBuild`): an + /// `IngestionRun`, or the `QueryRun` of the first evaluation that needs + /// one (stored once its coverage is closed). `IngestionRun` when any + /// output of the definition says so. pub built_at: StoredOutputBuild, /// Coverage width: each stored summary covers `(end - pane_ms, end]`. pub pane_ms: i64, @@ -105,9 +106,9 @@ pub enum SdsError { /// Stored outputs read each other in a cycle. #[error("stored outputs read each other in a cycle through {0:?}")] Cycle(DefinitionId), - /// An ingestion-run definition reads a query-run one, which is built + /// An `IngestionRun` definition reads a `QueryRun` one, which is built /// only when an evaluation needs it, so never in time for ingestion. - #[error("ingestion-run summary {definition:?} reads query-run summary {input:?}")] + #[error("stored output built in an IngestionRun ({definition:?}) reads one built in a QueryRun ({input:?})")] IngestionReadsQueryRun { definition: DefinitionId, input: DefinitionId, @@ -115,8 +116,8 @@ pub enum SdsError { } impl Sds { - /// The definitions behind `plan`'s stored outputs, built in an ingestion - /// run or in a query run and kept. + /// The definitions behind `plan`'s stored outputs, built in an + /// `IngestionRun` or in a `QueryRun`. pub fn from_plan(plan: &InstallablePlan) -> Result { let dag = &plan.dag.dag; let nodes: HashMap = @@ -194,7 +195,7 @@ impl Sds { /// The definitions built in ingestion runs, each after every definition /// it reads (`stored_inputs`). Fails on a cycle, and when an - /// ingestion-run definition reads a query-run one. + /// `IngestionRun` definition reads a `QueryRun` one. pub fn ingestion_build_order(&self) -> Result, SdsError> { let mut order = Vec::new(); let mut state: HashMap<&DefinitionId, bool> = HashMap::new(); // false: visiting diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs new file mode 100644 index 00000000..62232375 --- /dev/null +++ b/data_plane/tests/answers.rs @@ -0,0 +1,393 @@ +//! The data plane answers planned instant queries over HTTP, from stored +//! panes and from raw samples it fetches, and forwards what it cannot answer. + +mod common; + +use std::collections::BTreeMap; +use std::sync::Arc; + +use axum::body::Body; +use axum::extract::Query; +use axum::http::{HeaderMap, Request, StatusCode}; +use axum::routing::get; +use axum::{Json, Router}; +use data_plane::remote_write::{encode, Label, Sample, TimeSeries, WriteRequest}; +use data_plane::{router, DataPlane, EXECUTION_HEADER, FORWARD_REASON_HEADER}; +use http_body_util::BodyExt; +use serde_json::{json, Value}; +use tower::ServiceExt; + +const MINUTE: i64 = 60_000; + +/// `(labels, timestamp ms, value)`, labels including `__name__`. +type Point = (BTreeMap, i64, f64); + +fn series(metric: &str, labels: &[(&str, &str)]) -> BTreeMap { + let mut map = BTreeMap::from([("__name__".to_string(), metric.to_string())]); + map.extend(labels.iter().map(|(k, v)| (k.to_string(), v.to_string()))); + map +} + +/// A stand-in Prometheus holding `points`. It answers +/// `{__name__="m"}[Nms]` at `time` with the samples in `[time - N, time]`, +/// both ends included as Prometheus 2 does, and anything else with +/// `{"forwarded": query}`. +async fn fake_prometheus(points: Vec) -> String { + fake_prometheus_requiring(points, &[]).await +} + +/// [`fake_prometheus`] that answers 401 to a request missing any of +/// `required` (header name, value). +async fn fake_prometheus_requiring(points: Vec, required: &[(&str, &str)]) -> String { + let points = Arc::new(points); + let required: Arc> = Arc::new( + required + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + ); + let handler = move |headers: HeaderMap, Query(params): Query>| { + let points = points.clone(); + let required = required.clone(); + async move { + let authorized = required.iter().all(|(name, value)| { + headers + .get(name.as_str()) + .is_some_and(|v| v == value.as_str()) + }); + if !authorized { + return (StatusCode::UNAUTHORIZED, Json(json!({"status": "error"}))); + } + let query = ¶ms["query"]; + let parsed = query + .strip_prefix("{__name__=\"") + .and_then(|rest| rest.split_once("\"}[")) + .and_then(|(metric, rest)| { + Some((metric, rest.strip_suffix("ms]")?.parse::().ok()?)) + }); + let Some((metric, range)) = parsed else { + return (StatusCode::OK, Json(json!({ "forwarded": query }))); + }; + let time = (params["time"].parse::().unwrap() * 1000.0).round() as i64; + let mut by_series: BTreeMap<&BTreeMap, Vec> = BTreeMap::new(); + for (labels, t, v) in points.iter() { + if labels["__name__"] == metric && time - range <= *t && *t <= time { + by_series + .entry(labels) + .or_default() + .push(json!([*t as f64 / 1000.0, v.to_string()])); + } + } + let result: Vec = by_series + .into_iter() + .map(|(labels, values)| json!({"metric": labels, "values": values})) + .collect(); + ( + StatusCode::OK, + Json( + json!({"status": "success", "data": {"resultType": "matrix", "result": result}}), + ), + ) + } + }; + let app = Router::new().route("/api/v1/query", get(handler)); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let url = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + url +} + +struct Response { + execution: Option, + reason: Option, + body: Value, +} + +async fn instant(app: &Router, query: &str, time_ms: i64) -> Response { + let uri = format!( + "/api/v1/query?{}", + form_urlencoded::Serializer::new(String::new()) + .append_pair("query", query) + .append_pair("time", &(time_ms as f64 / 1000.0).to_string()) + .finish() + ); + let response = app + .clone() + .oneshot(Request::get(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let header = |name| { + response + .headers() + .get(name) + .map(|v: &axum::http::HeaderValue| v.to_str().unwrap().to_string()) + }; + let (execution, reason) = (header(EXECUTION_HEADER), header(FORWARD_REASON_HEADER)); + let body = response.into_body().collect().await.unwrap().to_bytes(); + Response { + execution, + reason, + body: serde_json::from_slice(&body).unwrap(), + } +} + +/// `(job, value)` of each result series, by job. +fn by_job(body: &Value) -> Vec<(String, f64)> { + let mut rows: Vec<(String, f64)> = body["data"]["result"] + .as_array() + .unwrap() + .iter() + .map(|s| { + ( + s["metric"]["job"].as_str().unwrap().to_string(), + s["value"][1].as_str().unwrap().parse().unwrap(), + ) + }) + .collect(); + rows.sort_by(|a, b| a.0.cmp(&b.0)); + rows +} + +/// Example 1 runs at query time: the data plane fetches the last minute of +/// raw samples from Prometheus and runs the plan over them. A sample at the +/// range's start, which Prometheus 2 returns, is excluded as in PromQL. +#[tokio::test] +async fn answers_query_time_plans_from_raw_samples() { + let mut points = Vec::new(); + for (job, instance) in [("a", "1"), ("a", "2"), ("b", "1")] { + let labels = series( + "http_requests_total", + &[("job", job), ("instance", instance)], + ); + for t in (0..10 * MINUTE).step_by(15_000) { + points.push((labels.clone(), t, t as f64 / 1000.0)); + } + } + let state = DataPlane::new( + common::plan_json("example1", true), + fake_prometheus(points).await, + 0, + ) + .unwrap(); + let app = router(Arc::new(state)); + let response = instant( + &app, + "sum by (job) (rate(http_requests_total[1m]))", + 5 * MINUTE, + ) + .await; + assert_eq!( + response.execution.as_deref(), + Some("plan"), + "{:?}", + response.body + ); + let result = by_job(&response.body); + let jobs: Vec<&str> = result.iter().map(|(job, _)| job.as_str()).collect(); + assert_eq!(jobs, ["a", "b"]); + for ((_, value), expected) in result.iter().zip([2.0, 1.0]) { + assert!((value - expected).abs() < 1e-9, "{result:?}"); + } + // (4 min, 5 min] holds the samples at 4:15 .. 5:00: 255 + 270 + 285 + 300. + // Including the one at 4:00 would add 240. + let response = instant( + &app, + "topk by (job) (10, sum_over_time(http_requests_total[1m]))", + 5 * MINUTE, + ) + .await; + assert_eq!( + response.execution.as_deref(), + Some("plan"), + "{:?}", + response.body + ); + let result = by_job(&response.body); + assert_eq!(result.len(), 3, "{result:?}"); + assert!( + result.iter().all(|(_, value)| *value == 1110.0), + "{result:?}" + ); + + let response = instant(&app, "up", 5 * MINUTE).await; + assert_eq!(response.execution.as_deref(), Some("forwarded")); + assert_eq!( + response.reason, None, + "unplanned queries are simply forwarded" + ); + assert_eq!(response.body, json!({"forwarded": "up"})); +} + +/// Example 4b reads six stored 10-minute KLL panes built from Remote Write; +/// at a time off the pane grid it forwards and says why. +#[tokio::test] +async fn answers_from_panes_built_from_remote_write() { + let state = Arc::new( + DataPlane::new( + common::plan_json("example4b", false), + fake_prometheus(vec![]).await, + 0, + ) + .unwrap(), + ); + let app = router(state.clone()); + let timeseries = ["a", "b"] + .into_iter() + .enumerate() + .map(|(i, job)| TimeSeries { + labels: vec![ + Label { + name: "__name__".into(), + value: "latency_ms".into(), + }, + Label { + name: "job".into(), + value: job.into(), + }, + ], + samples: (MINUTE..75 * MINUTE) + .step_by(15_000) + .map(|t| Sample { + timestamp: t, + value: (i * 1000) as f64 + ((t / 15_000) % 240) as f64, + }) + .collect(), + }) + .collect(); + let write = app + .clone() + .oneshot( + Request::post("/api/v1/write") + .body(Body::from(encode(&WriteRequest { timeseries }))) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(write.status(), StatusCode::NO_CONTENT); + // Panes ending at 20..70 min; the first, (0, 10 min], is skipped. + for _ in 0..500 { + if state.engine.queued_samples() == 0 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + assert_eq!(state.engine.stats().panes_built, 6); + + let query = "quantile_over_time(0.99, latency_ms[60m])"; + let response = instant(&app, query, 70 * MINUTE).await; + assert_eq!( + response.execution.as_deref(), + Some("plan"), + "{:?}", + response.body + ); + let result = by_job(&response.body); + assert_eq!(result.len(), 2); + for ((_, value), base) in result.iter().zip([0.0, 1000.0]) { + assert!( + *value >= base + 236.0 && *value <= base + 239.0, + "{result:?}" + ); + } + + let response = instant(&app, query, 70 * MINUTE + 1_000).await; + assert_eq!(response.execution.as_deref(), Some("forwarded")); + assert!(response.reason.unwrap().contains("pane grid")); +} + +/// A planned query's raw-range reads carry the caller's auth and tenancy +/// headers, as a forwarded query does: a Prometheus that requires them +/// serves the read, and the plan answers. +#[tokio::test] +async fn raw_range_reads_forward_auth_and_tenancy_headers() { + let labels = series("http_requests_total", &[("job", "a"), ("instance", "1")]); + let points = (0..10 * MINUTE) + .step_by(15_000) + .map(|t| (labels.clone(), t, t as f64 / 1000.0)) + .collect(); + let required = [ + ("authorization", "Bearer secret"), + ("x-scope-orgid", "tenant-a"), + ]; + let state = DataPlane::new( + common::plan_json("example1", true), + fake_prometheus_requiring(points, &required).await, + 0, + ) + .unwrap(); + let app = router(Arc::new(state)); + let uri = format!( + "/api/v1/query?{}", + form_urlencoded::Serializer::new(String::new()) + .append_pair("query", "sum by (job) (rate(http_requests_total[1m]))") + .append_pair("time", "300") + .finish() + ); + let request = |with_headers: bool| { + let mut request = Request::get(uri.clone()); + if with_headers { + request = request + .header("authorization", "Bearer secret") + .header("x-scope-orgid", "tenant-a"); + } + request.body(Body::empty()).unwrap() + }; + let response = app.clone().oneshot(request(true)).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers().get(EXECUTION_HEADER).unwrap(), + "plan", + "the raw read was authorized" + ); + // Without them the read is refused, so the query is forwarded, and + // Prometheus refuses that too. + let response = app.clone().oneshot(request(false)).await.unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + assert_eq!( + response.headers().get(EXECUTION_HEADER).unwrap(), + "forwarded" + ); + assert!(response + .headers() + .get(FORWARD_REASON_HEADER) + .unwrap() + .to_str() + .unwrap() + .contains("fetching raw samples")); +} + +/// A planned query at a time Prometheus rejects (not finite, or beyond the +/// millisecond range) gets 400 like Prometheus, and the engine keeps +/// running. +#[tokio::test] +async fn rejects_evaluation_times_out_of_range() { + let state = Arc::new( + DataPlane::new( + common::plan_json("example4b", true), + fake_prometheus(vec![]).await, + 0, + ) + .unwrap(), + ); + let app = router(state.clone()); + for time in ["-1e300", "1e300", "NaN", "inf"] { + let uri = format!( + "/api/v1/query?{}", + form_urlencoded::Serializer::new(String::new()) + .append_pair("query", "quantile_over_time(0.99, latency_ms[60m])") + .append_pair("time", time) + .finish() + ); + let response = app + .clone() + .oneshot(Request::get(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::BAD_REQUEST, "time={time}"); + let body = response.into_body().collect().await.unwrap().to_bytes(); + let body: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["errorType"], "bad_data", "time={time}"); + } + assert!(state.engine.is_running()); +} diff --git a/data_plane/tests/precompute.rs b/data_plane/tests/precompute.rs index 34c1622d..72cb9269 100644 --- a/data_plane/tests/precompute.rs +++ b/data_plane/tests/precompute.rs @@ -424,7 +424,7 @@ fn far_future_unrouted_sample_counts_as_unrouted() { /// Summaries built in a query run are not built at ingestion. #[test] -fn does_not_build_query_run_summaries() { +fn does_not_build_query_run_outputs() { let mut plan = common::plan("example4b", false); for output in &mut plan.stored_outputs { output.built_at = asap_planner::StoredOutputBuild::QueryRun; diff --git a/data_plane/tests/query.rs b/data_plane/tests/query.rs new file mode 100644 index 00000000..ae9b698e --- /dev/null +++ b/data_plane/tests/query.rs @@ -0,0 +1,320 @@ +mod common; + +use std::collections::BTreeMap; +use std::sync::Arc; + +use asap_planner::InstallablePlan; +use data_plane::precompute::{Precompute, PrecomputeOptions, RawSample}; +use data_plane::query::{ + BoundRawSampleWindow, LabelMatcher, MatchOp, QueryEngine, ResultSample, Unanswered, +}; +use data_plane::sds::Sds; + +const MINUTE: i64 = 60_000; + +struct Fixture { + precompute: Precompute, + query: QueryEngine, +} + +fn fixture(plan: InstallablePlan) -> Fixture { + let sds = Sds::from_plan(&plan).unwrap(); + let options = PrecomputeOptions { + allowed_lateness_ms: 0, + ..PrecomputeOptions::default() + }; + Fixture { + precompute: Precompute::new(plan.dag.dag.clone(), &sds, options).unwrap(), + query: QueryEngine::new(&plan, &sds), + } +} + +/// Samples every 15 s in `[from, to)` for each `(metric, labels)`, with +/// value `f(series index, t)`. +fn samples( + series: &[(&str, &[(&str, &str)])], + from: i64, + to: i64, + f: impl Fn(usize, i64) -> f64, +) -> Vec { + let series: Vec<_> = series + .iter() + .map(|(metric, labels)| { + let mut map = BTreeMap::from([("__name__".to_string(), metric.to_string())]); + map.extend(labels.iter().map(|(k, v)| (k.to_string(), v.to_string()))); + Arc::new(map) + }) + .collect(); + (from..to) + .step_by(15_000) + .flat_map(|t| { + series + .iter() + .enumerate() + .map(move |(i, labels)| (i, labels.clone(), t)) + }) + .map(|(i, labels, timestamp_ms)| RawSample { + labels, + timestamp_ms, + value: f(i, timestamp_ms), + }) + .collect() +} + +/// What a deployment fetches for each bound raw-sample window: its metric's samples in +/// `(start, end]`. +fn fetch(ranges: &[BoundRawSampleWindow], all: &[RawSample]) -> BTreeMap> { + ranges + .iter() + .map(|range| { + let samples = all + .iter() + .filter(|s| { + s.labels["__name__"] == range.metric + && range.start_ms < s.timestamp_ms + && s.timestamp_ms <= range.end_ms + }) + .cloned() + .collect(); + (range.node, samples) + }) + .collect() +} + +fn answer(f: &Fixture, query: &str, time_ms: i64, all: &[RawSample]) -> Vec { + let (root, ranges) = f + .query + .bound_raw_sample_windows(query, time_ms, &f.precompute.store) + .unwrap(); + let mut result = f + .query + .evaluate(root, time_ms, &f.precompute.store, &fetch(&ranges, all)) + .unwrap(); + result.sort_by(|a, b| a.labels.cmp(&b.labels)); + result +} + +const P99: &str = "quantile_over_time(0.99, latency_ms[60m])"; + +/// Example 4b's p99 merged from six stored 10-minute KLL panes equals one +/// query-time build over the same hour of raw samples (KLL is exact below k). +#[test] +fn stored_panes_answer_like_a_single_build() { + let series: &[(&str, &[(&str, &str)])] = &[ + ("latency_ms", &[("job", "a")]), + ("latency_ms", &[("job", "b")]), + ]; + // Values cycle through 0..240 per series, offset by 1000 for series b. + let all = samples(series, MINUTE, 75 * MINUTE, |i, t| { + (i * 1000) as f64 + ((t / 15_000) % 240) as f64 + }); + let mut panes = fixture(common::plan("example4b", false)); + assert!( + panes + .query + .bound_raw_sample_windows(P99, 70 * MINUTE, &panes.precompute.store) + .is_err(), + "no panes before ingestion" + ); + panes.precompute.ingest(&all); + let (_, ranges) = panes + .query + .bound_raw_sample_windows(P99, 70 * MINUTE, &panes.precompute.store) + .unwrap(); + assert!(ranges.is_empty(), "every input is a stored pane"); + let from_panes = answer(&panes, P99, 70 * MINUTE, &all); + + let single = fixture(common::plan("example4b", true)); + let (_, ranges) = single + .query + .bound_raw_sample_windows(P99, 70 * MINUTE, &single.precompute.store) + .unwrap(); + assert_eq!( + ranges + .iter() + .map(|r| (r.start_ms, r.end_ms)) + .collect::>(), + vec![(10 * MINUTE, 70 * MINUTE)], + "one raw hour, left-open" + ); + let from_raw = answer(&single, P99, 70 * MINUTE, &all); + + assert_eq!(from_panes, from_raw); + assert_eq!(from_panes.len(), 2); + for (sample, base) in from_panes.iter().zip([0.0, 1000.0]) { + assert!( + sample.value >= base + 236.0 && sample.value <= base + 239.0, + "{sample:?}" + ); + } +} + +/// Example 1 runs entirely at query time over raw samples: each counter grows +/// by 1 per second, so `sum by (job) (rate(...[1m]))` is the series count. +#[test] +fn query_time_plans_read_raw_ranges() { + let series: &[(&str, &[(&str, &str)])] = &[ + ("http_requests_total", &[("job", "a"), ("instance", "1")]), + ("http_requests_total", &[("job", "a"), ("instance", "2")]), + ("http_requests_total", &[("job", "b"), ("instance", "1")]), + ]; + let all = samples(series, 0, 10 * MINUTE, |_, t| t as f64 / 1000.0); + let f = fixture(common::plan("example1", true)); + let result = answer( + &f, + "sum by (job) (rate(http_requests_total[1m]))", + 5 * MINUTE, + &all, + ); + let by_job: Vec<(String, f64)> = result + .iter() + .map(|s| (s.labels["job"].clone(), s.value)) + .collect(); + assert_eq!(by_job.len(), 2); + assert_eq!(by_job[0].0, "a"); + assert!((by_job[0].1 - 2.0).abs() < 1e-9, "{by_job:?}"); + assert_eq!(by_job[1].0, "b"); + assert!((by_job[1].1 - 1.0).abs() < 1e-9, "{by_job:?}"); + + let topk = answer( + &f, + "topk by (job) (10, sum_over_time(http_requests_total[1m]))", + 5 * MINUTE, + &all, + ); + assert_eq!( + topk.len(), + 3, + "every series is in its job's top 10: {topk:?}" + ); +} + +/// Queries outside the plan, and stored panes that are missing or off the +/// pane grid, are not answered, so the caller forwards them. +#[test] +fn reports_what_it_cannot_answer() { + let mut f = fixture(common::plan("example4b", false)); + let store = &f.precompute.store; + assert!(matches!( + f.query.bound_raw_sample_windows("up", 70 * MINUTE, store), + Err(Unanswered::NotPlanned(_)) + )); + assert!(matches!( + f.query.bound_raw_sample_windows(P99, 70 * MINUTE, store), + Err(Unanswered::MissingPane(_)) + )); + let all = samples(&[("latency_ms", &[])], MINUTE, 75 * MINUTE, |_, t| t as f64); + f.precompute.ingest(&all); + let store = &f.precompute.store; + assert!(f + .query + .bound_raw_sample_windows(P99, 70 * MINUTE, store) + .is_ok()); + assert!(matches!( + f.query + .bound_raw_sample_windows(P99, 70 * MINUTE + 1, store), + Err(Unanswered::MissingPane(_)) + )); + // The pane ending at 10 min was skipped as possibly partial. + assert!(matches!( + f.query.bound_raw_sample_windows(P99, 60 * MINUTE, store), + Err(Unanswered::MissingPane(_)) + )); +} + +/// Evaluation times at the ends of `i64` are not answered and do not panic: +/// binding a window or pane to them would overflow. +#[test] +fn extreme_evaluation_times_are_not_answered() { + let f = fixture(common::plan("example4b", true)); + for time_ms in [i64::MIN, i64::MIN + 1] { + assert!(matches!( + f.query + .bound_raw_sample_windows(P99, time_ms, &f.precompute.store), + Err(Unanswered::Unsupported(_)) + )); + } + // The window ending at `i64::MAX` fits; binding it does not panic. + assert!(f + .query + .bound_raw_sample_windows(P99, i64::MAX, &f.precompute.store) + .is_ok()); + let panes = fixture(common::plan("example4b", false)); + assert!(matches!( + panes + .query + .bound_raw_sample_windows(P99, i64::MIN, &panes.precompute.store), + Err(Unanswered::Unsupported(_)) + )); +} + +/// Every query of examples 1, 3b and 4b, with and without raw data, answers +/// exactly like the Planner's query-time reference for the same plan +/// (nothing materialized) over the same samples. +#[test] +fn examples_answer_like_the_query_time_reference() { + let series: &[(&str, &[(&str, &str)])] = &[ + ("latency_ms", &[("job", "a")]), + ("latency_ms", &[("job", "b")]), + ("http_requests_total", &[("job", "a"), ("instance", "1")]), + ("http_requests_total", &[("job", "a"), ("instance", "2")]), + ("http_requests_total", &[("job", "b"), ("instance", "1")]), + ]; + // Samples through 70:45: the newest closed pane of every example ends + // at 70 min, the evaluation time. + let all = samples(series, MINUTE, 71 * MINUTE, |i, t| { + (i * 1000) as f64 + ((t / 15_000) % 240) as f64 + }); + for (example, raw) in [ + ("example1", true), + ("example3b", true), + ("example3b", false), + ("example4b", true), + ("example4b", false), + ] { + let plan = common::plan(example, raw); + let reference = plan.query_time_reference().unwrap(); + let mut installed = fixture(plan.clone()); + installed.precompute.ingest(&all); + let by_reference = fixture(reference); + for query in plan.queries.iter().map(|q| q.query.as_str()) { + let expected = answer(&by_reference, query, 70 * MINUTE, &all); + assert!(!expected.is_empty(), "{example} raw={raw} {query}"); + assert_eq!( + answer(&installed, query, 70 * MINUTE, &all), + expected, + "{example} raw={raw} {query}" + ); + } + } +} + +/// A raw window's label matchers go into the fetch, so Prometheus returns +/// only the series the plan keeps. +#[test] +fn bound_raw_sample_windows_carry_the_scans_label_matchers() { + let mut workload = common::workload("example4b"); + workload.query_workload.repeating_queries.as_mut().unwrap()[0].query = + asap_executor::planner::workload::Query( + "quantile_over_time(0.99, latency_ms{job=\"a\"}[60m])".into(), + ); + let plan = control_plane::plan_workload(&workload, &common::inputs(true)).unwrap(); + let f = fixture(plan); + let (_, ranges) = f + .query + .bound_raw_sample_windows( + "quantile_over_time(0.99, latency_ms{job=\"a\"}[60m])", + 70 * MINUTE, + &f.precompute.store, + ) + .unwrap(); + assert_eq!(ranges.len(), 1); + assert_eq!( + ranges[0].matchers, + vec![LabelMatcher { + name: "job".into(), + op: MatchOp::Equal, + value: "a".into(), + }] + ); +} diff --git a/data_plane/tests/sds.rs b/data_plane/tests/sds.rs index 1e27b416..53019fd4 100644 --- a/data_plane/tests/sds.rs +++ b/data_plane/tests/sds.rs @@ -119,7 +119,7 @@ fn identity_depends_on_the_source() { /// Outputs built in a query run and kept are stored summaries too: their /// definition says so, and no ingestion run builds them. #[test] -fn query_run_outputs_are_query_run_definitions() { +fn query_run_outputs_have_query_run_definitions() { let mut plan = plan("example4b", false); for output in &mut plan.stored_outputs { output.built_at = StoredOutputBuild::QueryRun; @@ -132,7 +132,7 @@ fn query_run_outputs_are_query_run_definitions() { assert!(sds.ingestion_build_order().unwrap().is_empty()); } -/// A definition shared by an ingestion-run and a query-run output is built +/// A definition shared by an `IngestionRun` and a `QueryRun` output is built /// at ingestion. #[test] fn ingestion_wins_when_outputs_share_a_definition() { @@ -231,7 +231,7 @@ fn build_order_puts_inputs_first() { assert_eq!(order, ["c", "b", "a"]); } -/// An ingestion-run definition cannot read a query-run one: that summary +/// An `IngestionRun` definition cannot read a `QueryRun` one: that summary /// is only built when an evaluation needs it. #[test] fn build_order_refuses_ingestion_reading_query_run() {