From 6f7b53c1dc0a7fb1c116c20b10039c047269d4dc Mon Sep 17 00:00:00 2001 From: zzylol Date: Wed, 7 Oct 2026 03:18:31 +0000 Subject: [PATCH] feat(query): answer planned instant queries by running the plan's roots Each plan root is compiled once with two kinds of inputs: stored outputs, bound at evaluation time T to their definition's pane ending at T - offset, and query-time raw ranges (TimeRange <- TimeShift* <- Scan), bound to the metric's samples in (T - offset - range, T - offset] fetched from Prometheus and filtered by the plan's Scan. An instant query whose text is a planned query is answered with a Prometheus vector and x-asap-execution: plan. Unplanned queries are forwarded; planned ones that cannot be answered (pane missing or off the pane grid, fetch or execution failure) are forwarded with x-asap-forward-reason. Co-Authored-By: Claude Opus 5.5 --- Cargo.lock | 1 + README.md | 10 +- data_plane/Cargo.toml | 2 + data_plane/src/engine.rs | 91 ++++++++- data_plane/src/lib.rs | 55 +++++- data_plane/src/precompute.rs | 77 ++++---- data_plane/src/prometheus.rs | 162 +++++++++++++++++ data_plane/src/query.rs | 344 +++++++++++++++++++++++++++++++++++ data_plane/tests/answers.rs | 265 +++++++++++++++++++++++++++ data_plane/tests/query.rs | 215 ++++++++++++++++++++++ 10 files changed, 1179 insertions(+), 43 deletions(-) create mode 100644 data_plane/src/prometheus.rs create mode 100644 data_plane/src/query.rs create mode 100644 data_plane/tests/answers.rs create mode 100644 data_plane/tests/query.rs diff --git a/Cargo.lock b/Cargo.lock index 33d20ff4..00c78698 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -490,6 +490,7 @@ dependencies = [ "axum", "clap", "control_plane", + "form_urlencoded", "futures", "http-body-util", "prost", diff --git a/README.md b/README.md index 744af865..c42eb009 100644 --- a/README.md +++ b/README.md @@ -33,14 +33,18 @@ 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 | Panes close `--allowed-lateness-ms` (default 15 s) after the newest sample passes their end; later samples for them are dropped. The first pane after startup is skipped because it may be partial. -Until the query engine lands, the data plane forwards every query to -Prometheus 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 raw ranges are read from Prometheus. Anything else is forwarded +to Prometheus (`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 diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 3c6a4655..1d7eca3f 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -21,7 +21,9 @@ reqwest.workspace = true prost = "0.13" snap = "1.1" futures = "0.3" +form_urlencoded = "1" [dev-dependencies] 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 b5b643a9..e2e220a5 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -2,6 +2,7 @@ //! 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::mpsc::{sync_channel, Receiver, SyncSender, TrySendError}; use asap_executor::physical_planner::frontier_from_timing; @@ -9,14 +10,28 @@ use control_plane::InstalledPlan; use tokio::sync::oneshot; use crate::precompute::{Precompute, PrecomputeStats, RawSample}; +use crate::query::{QueryEngine, RawRange, ResultSample, Unanswered}; use crate::sds::Sds; /// Commands queued for the engine before ingestion pushes back. const QUEUE: usize = 1024; +type Reply = oneshot::Sender>; + pub enum Command { Ingest(Vec), Stats(oneshot::Sender), + RawRanges { + query: String, + time_ms: i64, + reply: Reply<(usize, Vec)>, + }, + Evaluate { + root: usize, + time_ms: i64, + raw: BTreeMap>, + reply: Reply>, + }, } #[derive(Debug, thiserror::Error)] @@ -58,32 +73,74 @@ impl EngineHandle { } } + /// The plan root answering `query` at `time_ms` and the raw samples it + /// needs; see [`QueryEngine::raw_ranges`]. + pub async fn raw_ranges( + &self, + query: String, + time_ms: i64, + ) -> Result<(usize, Vec), Unanswered> { + self.ask(|reply| Command::RawRanges { + 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 busy = || Unanswered::Execution("the engine is busy".into()); + let (reply, answer) = oneshot::channel(); + self.commands.try_send(command(reply)).map_err(|_| busy())?; + answer.await.map_err(|_| busy())? + } + /// Precompute counters, once every command queued before has run. pub async fn stats(&self) -> PrecomputeStats { let (reply, stats) = oneshot::channel(); - let _ = self.commands.send(Command::Stats(reply)); + if self.commands.try_send(Command::Stats(reply)).is_err() { + return PrecomputeStats::default(); + } stats.await.unwrap_or_default() } } struct Engine { precompute: Precompute, + query: QueryEngine, } impl Engine { fn new(plan_json: &str, allowed_lateness_ms: i64) -> anyhow::Result { let plan: InstalledPlan = serde_json::from_str(plan_json)?; plan.validate()?; - let dag = plan.dag.dag; - let frontier = frontier_from_timing(&dag)?; - let sds = Sds::from_dag(&dag, &frontier)?; + let frontier = frontier_from_timing(&plan.dag.dag)?; + let sds = Sds::from_dag(&plan.dag.dag, &frontier)?; tracing::info!( definitions = sds.definitions.len(), outputs = sds.outputs.len(), "stored summaries" ); + let query = QueryEngine::new(&plan, &sds); Ok(Self { - precompute: Precompute::new(dag, &sds, allowed_lateness_ms)?, + precompute: Precompute::new(plan.dag.dag, &sds, allowed_lateness_ms)?, + query, }) } @@ -94,6 +151,30 @@ impl Engine { Command::Stats(reply) => { let _ = reply.send(self.precompute.stats); } + Command::RawRanges { + query, + time_ms, + reply, + } => { + let _ = reply.send(self.query.raw_ranges( + &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, + )); + } } } } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 9de90c76..0223c40a 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,12 +26,16 @@ use axum::Router; use control_plane::InstalledPlan; use engine::EngineHandle; use precompute::RawSample; +use prometheus::InstantQuery; +use query::Unanswered; /// Prometheus' staleness marker: a NaN that ends a series, not a sample. const STALE_NAN_BITS: u64 = 0x7ff0_0000_0000_0002; -/// 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"; pub struct DataPlane { /// The installed plan's queries, in root order. @@ -110,6 +117,7 @@ async fn write(State(state): State>, body: Bytes) -> Response { } } +/// Answers a planned instant query from the plan; forwards anything else. async fn query( state: State>, method: Method, @@ -117,7 +125,46 @@ async fn query( headers: HeaderMap, body: Bytes, ) -> Response { - forward(state, "query", method, query, headers, body).await + let Some(instant) = InstantQuery::parse(query.0.as_deref(), &body) else { + return forward(state, "query", method, query, headers, body).await; + }; + let reason = match answer(&state, instant).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"); + HeaderValue::from_str(&reason.to_string()).ok() + } + }; + 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 +} + +async fn answer(state: &DataPlane, query: InstantQuery) -> 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.raw_ranges(query.query, time_ms).await?; + let mut raw = BTreeMap::new(); + for range in &ranges { + let samples = prometheus::fetch(&state.client, &state.prometheus_url, range) + .await + .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))?; + raw.insert(range.node, samples); + } + let samples = state.engine.evaluate(root, time_ms, raw).await?; + Ok(prometheus::vector_response(time_ms, &samples)) } async fn query_range( diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index a0aabe86..eb836431 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -62,11 +62,17 @@ pub struct PrecomputeStats { pub panes_failed: u64, } +/// A time-series `Scan` of the plan: its node, metric and row schema. +#[derive(Debug, Clone)] +pub struct ScanInput { + pub node: u64, + pub metric: String, + pub schema: SchemaRef, +} + struct Builder { definition: SummaryDefinition, - metric: String, - scan: u64, - scan_schema: SchemaRef, + scan: ScanInput, build: CompiledPhysicalDAG, /// Oldest pane offset any output reads, plus one pane: panes ending /// before `watermark - retention_ms` are no longer read. @@ -98,10 +104,10 @@ impl Precompute { for definition in sds.definitions.values() { let error = |message: String| PrecomputeError::Definition(definition.id.clone(), message); - let (scan, metric, scan_schema) = scan_below(&dag, definition.raw_input) + let scan = scan_below(&dag, definition.raw_input) .ok_or_else(|| error("no time-series Scan below its TimeRange".into()))?; let raw_schema = Arc::new(output_schema(&dag, definition.raw_input).clone()); - if raw_schema != scan_schema { + if raw_schema != scan.schema { return Err(error("TimeRange rows differ from its Scan's".into())); } let build = compile( @@ -118,9 +124,7 @@ impl Precompute { .unwrap_or(0); builders.push(Builder { definition: definition.clone(), - metric, scan, - scan_schema, build, retention_ms: max_offset + 2 * definition.pane_ms, open: BTreeMap::new(), @@ -150,7 +154,7 @@ impl Precompute { let metric = sample.labels.get("__name__").map(String::as_str); let mut accepted = false; for builder in &mut self.builders { - if metric != Some(builder.metric.as_str()) { + if metric != Some(builder.scan.metric.as_str()) { continue; } let end = pane_end(sample.timestamp_ms, builder.definition.pane_ms); @@ -225,23 +229,9 @@ fn build_pane( window_end_ms: end, revision: 1, }; - let rows = samples - .iter() - .map(|s| series_row(&builder.scan_schema, &s.labels, s.timestamp_ms, s.value)) - .collect::, _>>()?; - let raw = Batch::try_new(builder.scan_schema.clone(), rows)?; - let mut sources = DataSources::default(); - sources.register( - ScanSource::TimeSeries { - metric: builder.metric.clone(), - }, - Arc::new(MemorySource::new(builder.scan_schema.clone(), vec![raw])?), - )?; - let scan = bind_with_data_sources(dag, BTreeMap::new(), &[builder.scan], &sources)?; - let selected = run(&scan, &[builder.scan], scope.clone())?; - let selected = Batch::try_new(builder.scan_schema.clone(), selected)?; + let selected = select(dag, &builder.scan, samples, scope.clone())?; let input: Source<'_> = Box::new(Operator::source( - builder.scan_schema.clone(), + builder.scan.schema.clone(), vec![selected], )?); let build = builder @@ -252,6 +242,31 @@ fn build_pane( Ok(Batch::try_new(schema, rows)?) } +/// The rows of `samples` the plan's `scan` selects: its label matchers +/// applied by the executor's `Scan`. +pub fn select( + dag: &PhysicalASAPDAG, + scan: &ScanInput, + samples: &[RawSample], + scope: Scope, +) -> Result { + let rows = samples + .iter() + .map(|s| series_row(&scan.schema, &s.labels, s.timestamp_ms, s.value)) + .collect::, _>>()?; + let raw = Batch::try_new(scan.schema.clone(), rows)?; + let mut sources = DataSources::default(); + sources.register( + ScanSource::TimeSeries { + metric: scan.metric.clone(), + }, + Arc::new(MemorySource::new(scan.schema.clone(), vec![raw])?), + )?; + let dag = bind_with_data_sources(dag, BTreeMap::new(), &[scan.node], &sources)?; + let selected = run(&dag, &[scan.node], scope)?; + Batch::try_new(scan.schema.clone(), selected) +} + /// Executes `dag`'s single root to completion and returns its rows. pub fn run( dag: &asap_executor::plan::PhysicalDAG<'_, Batch, SchemaRef>, @@ -277,8 +292,8 @@ fn output_schema(dag: &PhysicalASAPDAG, id: u64) -> &planner_types::ir::schema:: .output_schema } -/// The time-series `Scan` that `id` reads: its node id, metric and row schema. -fn scan_below(dag: &PhysicalASAPDAG, id: u64) -> Option<(u64, String, SchemaRef)> { +/// The time-series `Scan` that `id` reads through single-input operators. +pub fn scan_below(dag: &PhysicalASAPDAG, id: u64) -> Option { let inputs: HashMap = dag .edges .iter() @@ -295,11 +310,11 @@ fn scan_below(dag: &PhysicalASAPDAG, id: u64) -> Option<(u64, String, SchemaRef) }, } = &node.payload { - return Some(( - current, - metric.clone(), - Arc::new(node.output_schema.clone()), - )); + return Some(ScanInput { + node: current, + metric: metric.clone(), + schema: Arc::new(node.output_schema.clone()), + }); } current = *inputs.get(¤t)?; } diff --git a/data_plane/src/prometheus.rs b/data_plane/src/prometheus.rs new file mode 100644 index 00000000..08e753d8 --- /dev/null +++ b/data_plane/src/prometheus.rs @@ -0,0 +1,162 @@ +//! 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::{RawRange, 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, +} + +impl InstantQuery { + /// Reads `query` and `time` (Unix seconds); `None` if either is missing + /// or `time` is in another format, such as RFC 3339. + pub fn parse(url_query: Option<&str>, form: &[u8]) -> Option { + 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) => Some(seconds_to_ms(time.parse().ok()?)), + }; + Some(Self { + query: params.get("query")?.clone(), + time_ms, + }) + } +} + +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 samples of `range`'s metric in `(start, end]`, read from Prometheus +/// with a range selector evaluated at the range's end. +pub async fn fetch( + client: &reqwest::Client, + prometheus_url: &str, + range: &RawRange, +) -> anyhow::Result> { + let selector = format!( + "{{__name__=\"{}\"}}[{}ms]", + range.metric, + range.end_ms - range.start_ms + ); + let response: Value = client + .get(format!("{prometheus_url}/api/v1/query")) + .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(); + assert_eq!((q.query.as_str(), q.time_ms), ("up", Some(1500))); + let q = InstantQuery::parse(None, b"query=sum%28x%29").unwrap(); + assert_eq!((q.query.as_str(), q.time_ms), ("sum(x)", None)); + assert!(InstantQuery::parse(Some("time=1"), b"").is_none()); + assert!(InstantQuery::parse(Some("query=up&time=2026-01-01T00:00:00Z"), b"").is_none()); + } + + /// 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..e3dc40b0 --- /dev/null +++ b/data_plane/src/query.rs @@ -0,0 +1,344 @@ +//! 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 raw ranges, `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::promql_rows::{ + decode_series_identity, SERIES_IDENTITY_COLUMN, +}; +use asap_executor::physical_planner::{compile, CompiledPhysicalDAG, InputContract, Source}; +use asap_executor::runtime::Scope; +use asap_executor::values::{Batch, Value}; +use control_plane::InstalledPlan; +use planner_types::ir::export::{ + LogicalASAPOperatorPayload as Payload, NonASAPOpKind, PhysicalASAPDAG, +}; +use planner_types::ir::schema::{DataType, FieldDataType, Schema}; + +use crate::precompute::{run, scan_below, select, RawSample, ScanInput}; +use crate::sds::{DefinitionId, Sds, SummaryStore}; + +/// Raw samples a query needs: `metric`'s samples in `(start_ms, end_ms]`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RawRange { + pub node: u64, + pub metric: String, + pub start_ms: i64, + pub end_ms: i64, +} + +/// 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, + scan: ScanInput, + range_ms: i64, + offset_ms: i64, +} + +struct Program { + compiled: CompiledPhysicalDAG, + stored: Vec, + raw: Vec, + output: Schema, +} + +pub struct QueryEngine { + dag: PhysicalASAPDAG, + /// Each query, with its compiled root or why it cannot run. + roots: Vec<(String, Result)>, +} + +impl QueryEngine { + pub fn new(plan: &InstalledPlan, sds: &Sds) -> Self { + let dag = plan.dag.dag.clone(); + let roots = plan + .queries + .iter() + .zip(&dag.roots) + .map(|(query, root)| { + ( + query.trim().to_string(), + program(&dag, sds, u64::from(root.0)), + ) + }) + .collect(); + Self { dag, 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 raw_ranges( + &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| RawRange { + node: input.node, + metric: input.scan.metric.clone(), + start_ms: time_ms - input.offset_ms - input.range_ms, + end_ms: time_ms - input.offset_ms, + }) + .collect(); + Ok((root, ranges)) + } + + /// Runs `root` at `time_ms` over stored panes and `raw`, the samples of + /// each of its [`RawRange`]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 = select(&self.dag, &input.scan, samples, scope.clone()).map_err(failed)?; + let source = Operator::source(input.scan.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(&dag, program.compiled.roots(), 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())), + } + } +} + +/// 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 = time_ms - input.offset_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, 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, root: u64) -> Result { + let nodes: HashMap = dag.nodes.iter().map(|n| (u64::from(n.id.0), n)).collect(); + let mut producers: HashMap> = HashMap::new(); + for edge in &dag.edges { + producers + .entry(u64::from(edge.consumer.0)) + .or_default() + .push(u64::from(edge.producer.0)); + } + 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 { + Payload::Relational { + operator: NonASAPOpKind::TimeRange { range, .. }, + } => { + raw.push(raw_input( + dag, + &nodes, + &producers, + id, + range.as_millis() as i64, + )?); + continue; + } + Payload::Relational { + operator: NonASAPOpKind::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.scan.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(), + }) +} + +/// A query-time `TimeRange`: the `Scan` below it and the `TimeShift` +/// offsets between them. +fn raw_input( + dag: &PhysicalASAPDAG, + nodes: &HashMap, + producers: &HashMap>, + time_range: u64, + range_ms: i64, +) -> Result { + let unsupported = |reason: &str| format!("node {time_range}: {reason}"); + let mut offset_ms = 0; + let mut current = time_range; + loop { + current = match producers.get(¤t).map(Vec::as_slice) { + Some([input]) => *input, + _ => return Err(unsupported("a TimeRange must read one input")), + }; + match &nodes[¤t].payload { + Payload::Relational { + operator: NonASAPOpKind::TimeShift { shift }, + } if shift.at.is_none() => offset_ms += shift.offset_ms, + Payload::Relational { + operator: NonASAPOpKind::Scan { .. }, + } => break, + _ => { + return Err(unsupported( + "only offset TimeShifts between TimeRange and Scan", + )) + } + } + } + let scan = scan_below(dag, time_range).ok_or_else(|| unsupported("not a time-series Scan"))?; + if Arc::new(nodes[&time_range].output_schema.clone()) != scan.schema { + return Err(unsupported("TimeRange rows differ from its Scan's")); + } + Ok(RawInput { + node: time_range, + scan, + range_ms, + offset_ms, + }) +} + +/// 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 mut labels = BTreeMap::new(); + for (field, cell) in schema.fields.iter().zip(&row) { + let Value::Utf8(text) = cell else { continue }; + if field.name == SERIES_IDENTITY_COLUMN { + let identity = decode_series_identity(text) + .map_err(|e| Unanswered::Execution(e.to_string()))?; + labels.extend(identity); + } else if !text.is_empty() { + labels.insert(field.name.clone(), text.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/tests/answers.rs b/data_plane/tests/answers.rs new file mode 100644 index 00000000..c08ec7ba --- /dev/null +++ b/data_plane/tests/answers.rs @@ -0,0 +1,265 @@ +//! 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::{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 { + let points = Arc::new(points); + let handler = move |Query(params): Query>| { + let points = points.clone(); + async move { + 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 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(); + 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. + assert_eq!(state.engine.stats().await.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")); +} diff --git a/data_plane/tests/query.rs b/data_plane/tests/query.rs new file mode 100644 index 00000000..04af1d49 --- /dev/null +++ b/data_plane/tests/query.rs @@ -0,0 +1,215 @@ +mod common; + +use std::collections::BTreeMap; +use std::sync::Arc; + +use asap_executor::physical_planner::frontier_from_timing; +use control_plane::InstalledPlan; +use data_plane::precompute::{Precompute, RawSample}; +use data_plane::query::{QueryEngine, RawRange, ResultSample, Unanswered}; +use data_plane::sds::Sds; + +const MINUTE: i64 = 60_000; + +struct Fixture { + precompute: Precompute, + query: QueryEngine, +} + +fn fixture(plan: InstalledPlan) -> Fixture { + let dag = &plan.dag.dag; + let sds = Sds::from_dag(dag, &frontier_from_timing(dag).unwrap()).unwrap(); + Fixture { + precompute: Precompute::new(dag.clone(), &sds, 0).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 raw range: its metric's samples in +/// `(start, end]`. +fn fetch(ranges: &[RawRange], 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 + .raw_ranges(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 + .raw_ranges(P99, 70 * MINUTE, &panes.precompute.store) + .is_err(), + "no panes before ingestion" + ); + panes.precompute.ingest(&all); + let (_, ranges) = panes + .query + .raw_ranges(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 + .raw_ranges(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.raw_ranges("up", 70 * MINUTE, store), + Err(Unanswered::NotPlanned(_)) + )); + assert!(matches!( + f.query.raw_ranges(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.raw_ranges(P99, 70 * MINUTE, store).is_ok()); + assert!(matches!( + f.query.raw_ranges(P99, 70 * MINUTE + 1, store), + Err(Unanswered::MissingPane(_)) + )); + // The pane ending at 10 min was skipped as possibly partial. + assert!(matches!( + f.query.raw_ranges(P99, 60 * MINUTE, store), + Err(Unanswered::MissingPane(_)) + )); +}