From 48450b9256a6cb7e0eb0179972f4b8f358315568 Mon Sep 17 00:00:00 2001 From: zzylol Date: Sun, 18 Oct 2026 15:36:25 +0000 Subject: [PATCH 1/6] feat(query): answer planned instant queries by running the plan's roots Each plan root (InstallablePlan.queries) is compiled once with two kinds of inputs: stored outputs, bound at evaluation time T to their definition's summary ending at T - offset, and the plan's query-time raw windows (query_time_raw_windows), bound to the metric's samples in (T - offset - width, T - offset] fetched from Prometheus and selected by the window's RawSampleSelector. 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 (summary missing or off the coverage grid, fetch or execution failure) are forwarded with x-asap-forward-reason. Raw-window reads carry the caller's end-to-end headers (Authorization, X-Scope-OrgID and the rest query forwarding keeps), so Prometheus authorizes and scopes them as it would the caller's own query. Answer labels come from series_labels_of_output_row. --- Cargo.lock | 1 + README.md | 12 +- data_plane/Cargo.toml | 2 + data_plane/src/engine.rs | 82 ++++++++ data_plane/src/lib.rs | 65 ++++++- data_plane/src/precompute.rs | 27 +-- data_plane/src/prometheus.rs | 167 ++++++++++++++++ data_plane/src/query.rs | 302 +++++++++++++++++++++++++++++ data_plane/tests/answers.rs | 358 +++++++++++++++++++++++++++++++++++ data_plane/tests/query.rs | 217 +++++++++++++++++++++ 10 files changed, 1214 insertions(+), 19 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 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..39213724 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 raw ranges 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 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..858cfe2c 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::{QueryEngine, RawRange, 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), + RawRanges { + 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::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 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 { @@ -194,8 +250,10 @@ impl Engine { "stored summaries" ); 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 +268,30 @@ impl Engine { .queued_samples .fetch_sub(samples.len() as u64, Ordering::SeqCst); } + 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, + )); + } Command::Stop => break, } } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index fe037eb7..b545af3a 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,7 @@ async fn status(State(state): State>) -> Response { .into_response() } +/// Answers a planned instant query from the plan; forwards anything else. async fn query( state: State>, method: Method, @@ -184,7 +192,56 @@ 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, &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"); + 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 +} + +/// Runs a planned instant query. Raw ranges 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.raw_ranges(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); + let mut raw = BTreeMap::new(); + for range in &ranges { + let samples = + prometheus::fetch(&state.client, &state.prometheus_url, &fetch_headers, 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 9658a4fe..1a0b1b31 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -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..22f4f032 --- /dev/null +++ b/data_plane/src/prometheus.rs @@ -0,0 +1,167 @@ +//! 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. `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: &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")) + .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(); + 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..775c66bf --- /dev/null +++ b/data_plane/src/query.rs @@ -0,0 +1,302 @@ +//! 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::{compile, CompiledPhysicalDAG, InputContract, Source}; +use asap_executor::planner::ir::physical_export::{PhysicalASAPDAG, RawSampleWindow}; +use asap_executor::planner::ir::schema::{DataType, FieldDataType, Schema}; +use asap_executor::planner::ir::{NonASAPOp, Operator as IrOperator}; +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}; + +/// 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, + /// 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 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.selector.metric().to_string(), + 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 = input + .selector + .window_batch(time_ms - input.offset_ms, &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())), + } + } +} + +/// 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, 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, + 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(), + }) +} + +/// 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/tests/answers.rs b/data_plane/tests/answers.rs new file mode 100644 index 00000000..9589f5dd --- /dev/null +++ b/data_plane/tests/answers.rs @@ -0,0 +1,358 @@ +//! 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")); +} diff --git a/data_plane/tests/query.rs b/data_plane/tests/query.rs new file mode 100644 index 00000000..14b44a43 --- /dev/null +++ b/data_plane/tests/query.rs @@ -0,0 +1,217 @@ +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::{QueryEngine, RawRange, 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 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(_)) + )); +} From 6be30eb2e91fb4dad6981f0f636c94b6ab809db3 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 20 Oct 2026 23:36:51 +0000 Subject: [PATCH 2/6] fix(query): bound evaluation times, push label matchers into raw fetches - A query whose time is not finite or beyond MAX_TIME_MS gets 400 bad_data, as Prometheus answers; window and pane binding use checked arithmetic, so no evaluation time can overflow (time=-1e300 panicked the engine thread in debug builds). - Raw-window fetches select only the series the plan keeps: the Scan's label matchers PromQL can express go into the range selector, with metric and values escaped. The selector still applies every matcher. - Raw windows are fetched concurrently. - A forward reason with characters a header cannot carry is sanitized instead of dropped. --- data_plane/src/lib.rs | 62 +++++++++++++--- data_plane/src/prometheus.rs | 132 ++++++++++++++++++++++++++++++----- data_plane/src/query.rs | 126 ++++++++++++++++++++++++++++++--- data_plane/tests/answers.rs | 35 ++++++++++ data_plane/tests/query.rs | 98 +++++++++++++++++++++++++- 5 files changed, 415 insertions(+), 38 deletions(-) diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index b545af3a..0c26989d 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -184,6 +184,22 @@ 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>, @@ -192,8 +208,13 @@ async fn query( headers: HeaderMap, body: Bytes, ) -> Response { - let Some(instant) = InstantQuery::parse(query.0.as_deref(), &body) else { - return 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) => { @@ -206,7 +227,7 @@ async fn query( Err(Unanswered::NotPlanned(_)) => None, Err(reason) => { tracing::debug!(%reason, "forwarding a planned query"); - HeaderValue::from_str(&reason.to_string()).ok() + Some(header_value(&reason.to_string())) } }; let mut response = forward(state, "query", method, query, headers, body).await; @@ -232,14 +253,17 @@ async fn answer( let mut fetch_headers = forwarded_headers(headers); // The raw read is a GET without the caller's body. fetch_headers.remove(header::CONTENT_TYPE); - let mut raw = BTreeMap::new(); - for range in &ranges { - let samples = - prometheus::fetch(&state.client, &state.prometheus_url, &fetch_headers, range) - .await - .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))?; - raw.insert(range.node, samples); - } + // The raw 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)) } @@ -345,3 +369,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/prometheus.rs b/data_plane/src/prometheus.rs index 22f4f032..2e18a871 100644 --- a/data_plane/src/prometheus.rs +++ b/data_plane/src/prometheus.rs @@ -16,10 +16,22 @@ pub struct InstantQuery { 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); `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 { + /// 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)) @@ -27,15 +39,36 @@ impl InstantQuery { .collect(); let time_ms = match params.get("time") { None => None, - Some(time) => Some(seconds_to_ms(time.parse().ok()?)), + 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" + )) + })?), + }, }; - Some(Self { - query: params.get("query")?.clone(), + 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 } @@ -64,6 +97,37 @@ fn format_value(value: f64) -> 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: &RawRange) -> 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 @@ -75,11 +139,7 @@ pub async fn fetch( headers: &reqwest::header::HeaderMap, range: &RawRange, ) -> anyhow::Result> { - let selector = format!( - "{{__name__=\"{}\"}}[{}ms]", - range.metric, - range.end_ms - range.start_ms - ); + let selector = range_selector(range); let response: Value = client .get(format!("{prometheus_url}/api/v1/query")) .headers(headers.clone()) @@ -134,12 +194,52 @@ mod tests { /// 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(); + 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(); + 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"").is_none()); - assert!(InstantQuery::parse(Some("query=up&time=2026-01-01T00:00:00Z"), b"").is_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 = RawRange { + 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. diff --git a/data_plane/src/query.rs b/data_plane/src/query.rs index 775c66bf..bca9bffb 100644 --- a/data_plane/src/query.rs +++ b/data_plane/src/query.rs @@ -17,8 +17,9 @@ 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}; +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}; @@ -27,15 +28,52 @@ use asap_planner::InstallablePlan; use crate::precompute::{executor_samples, run_root, RawSample}; use crate::sds::{DefinitionId, Sds, Selection, SummaryStore}; -/// Raw samples a query needs: `metric`'s samples in `(start_ms, end_ms]`. +/// Raw samples a query needs: `metric`'s samples in `(start_ms, end_ms]` +/// that `matchers` keep. #[derive(Debug, Clone, PartialEq, Eq)] pub struct RawRange { 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 { @@ -65,6 +103,7 @@ struct StoredInput { struct RawInput { node: u64, + matchers: Vec, /// The plan's `Scan` label matchers over the range's samples. selector: RawSampleSelector, schema: SchemaRef, @@ -118,13 +157,17 @@ impl QueryEngine { let ranges = program .raw .iter() - .map(|input| RawRange { - node: input.node, - metric: input.selector.metric().to_string(), - start_ms: time_ms - input.offset_ms - input.range_ms, - end_ms: time_ms - input.offset_ms, + .map(|input| { + let (start_ms, end_ms) = window(time_ms, input.offset_ms, input.range_ms)?; + Ok(RawRange { + node: input.node, + metric: input.selector.metric().to_string(), + matchers: input.matchers.clone(), + start_ms, + end_ms, + }) }) - .collect(); + .collect::>()?; Ok((root, ranges)) } @@ -157,7 +200,10 @@ impl QueryEngine { let samples = raw.get(&input.node).map(Vec::as_slice).unwrap_or_default(); let rows = input .selector - .window_batch(time_ms - input.offset_ms, &executor_samples(samples)) + .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)); @@ -181,13 +227,23 @@ impl QueryEngine { } } +/// `(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 = time_ms - input.offset_ms; + 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", @@ -244,6 +300,7 @@ fn program( 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, @@ -274,6 +331,55 @@ fn program( }) } +/// 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> { diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs index 9589f5dd..62232375 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -356,3 +356,38 @@ async fn raw_range_reads_forward_auth_and_tenancy_headers() { .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/query.rs b/data_plane/tests/query.rs index 14b44a43..a6775ef8 100644 --- a/data_plane/tests/query.rs +++ b/data_plane/tests/query.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use asap_planner::InstallablePlan; use data_plane::precompute::{Precompute, PrecomputeOptions, RawSample}; -use data_plane::query::{QueryEngine, RawRange, ResultSample, Unanswered}; +use data_plane::query::{LabelMatcher, MatchOp, QueryEngine, RawRange, ResultSample, Unanswered}; use data_plane::sds::Sds; const MINUTE: i64 = 60_000; @@ -215,3 +215,99 @@ fn reports_what_it_cannot_answer() { 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.raw_ranges(P99, time_ms, &f.precompute.store), + Err(Unanswered::Unsupported(_)) + )); + } + // The window ending at `i64::MAX` fits; binding it does not panic. + assert!(f + .query + .raw_ranges(P99, i64::MAX, &f.precompute.store) + .is_ok()); + let panes = fixture(common::plan("example4b", false)); + assert!(matches!( + panes + .query + .raw_ranges(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 raw_ranges_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 + .raw_ranges( + "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(), + }] + ); +} From afa5bc6db5e0aae9150d0d6474b05dc3d90238d7 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 20 Oct 2026 23:44:28 +0000 Subject: [PATCH 3/6] feat(engine): warn at startup about query-run summaries not built yet A plan that keeps query-run summaries installs, but this data plane does not build them yet, so queries reading them are forwarded. Say so in the log at startup. --- data_plane/src/engine.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 858cfe2c..6c962658 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -249,6 +249,18 @@ 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 keeps query-run summaries, 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 { From 9bfa31fbdfbe47efe31c02850a0b532047509114 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 20 Oct 2026 19:59:45 +0000 Subject: [PATCH 4/6] feat(query): build query-run summaries on first use and keep them once closed A stored output the plan builds in a query run and keeps (StoredOutputBuild::QueryRun, #509 Example 4's B3) is now executed. The query engine builds its definition's summary from the coverage's raw samples the first time an evaluation needs it, and stores it in the SDS store under the same SummaryInstanceKey as ingestion-built summaries once its coverage is closed. Later evaluations reuse it. - Closure is judged by the ingestion progress snapshotted when the evaluation chose its fetch (raw_ranges returns it with the ranges and evaluate takes it back), not when the evaluation runs, so a partial summary built from an earlier fetch is never kept. - Eviction is by the plan's retention_ms plus one pane, measured from the newest closed coverage at or before the watermark, as for ingestion-built summaries. A far-future evaluation no longer stops keeping. - One fetch per query-run definition covers all its missing summaries. - KeptResultStats (built, kept, kept_while_metric_behind, reused) is on /api/v1/status/asap. kept_while_metric_behind counts summaries kept while their own metric's newest sample had not reached their end: the watermark is global. - A query-run definition whose selector does not build no longer fails the install: its queries are forwarded. SummaryBuild, shared by the precompute and query engines, compiles a definition's build and runs it over one coverage's samples. Samples of a metric only query-run summaries read move the watermark (counted as samples_query_run_only) without being buffered. The backend now declares query_time_results_retained: true. examples/planner-layering-example4-b3.json is a workload for which the Planner selects B3 under the default cost model. --- README.md | 14 +- control_plane/src/lib.rs | 8 +- control_plane/tests/plan_workload.rs | 17 + data_plane/src/engine.rs | 58 ++-- data_plane/src/lib.rs | 6 +- data_plane/src/precompute.rs | 219 ++++++++----- data_plane/src/query.rs | 284 ++++++++++++++-- data_plane/src/sds.rs | 4 + data_plane/tests/kept.rs | 356 +++++++++++++++++++++ data_plane/tests/precompute.rs | 30 +- data_plane/tests/query.rs | 94 ++++-- data_plane/tests/sds.rs | 24 ++ examples/planner-layering-example4-b3.json | 44 +++ 13 files changed, 986 insertions(+), 172 deletions(-) create mode 100644 data_plane/tests/kept.rs create mode 100644 examples/planner-layering-example4-b3.json diff --git a/README.md b/README.md index 39213724..1940572e 100644 --- a/README.md +++ b/README.md @@ -56,7 +56,16 @@ engine's sample and pane counters as JSON. 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 +query-time raw ranges are read from Prometheus. A summary the plan builds in +a query run and keeps (`MatAtQueryKeptAcrossEvaluations`) is built from its +coverage's raw samples the first time an evaluation needs it, stored in the +SDS under the same key as ingestion-built summaries once its coverage is +closed (its end at or before the watermark when the evaluation chose its +fetch), reused by later evaluations, and evicted like ingestion-built +summaries: the plan's `retention_ms` plus one pane, measured from the newest +closed coverage at or before the watermark. `/api/v1/status/asap` reports +`query_run` counters, including summaries kept while their own metric lagged +behind the global watermark. 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 @@ -75,8 +84,7 @@ 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 query-time results kept across evaluations. 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..32ecbbe6 100644 --- a/control_plane/src/lib.rs +++ b/control_plane/src/lib.rs @@ -82,14 +82,14 @@ 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 -/// outputs, but nothing builds them until the query engine does. +/// Planner executor builds, raw samples kept in Prometheus, and query-time +/// results kept across evaluations (the query engine builds query-run +/// summaries on first use and keeps them once their coverage is closed). pub fn backend_capabilities() -> DeploymentCapabilities { DeploymentCapabilities { supported_summaries: None, raw_data_retained: true, - query_time_results_retained: false, + query_time_results_retained: true, } } diff --git a/control_plane/tests/plan_workload.rs b/control_plane/tests/plan_workload.rs index f4855511..e2680f64 100644 --- a/control_plane/tests/plan_workload.rs +++ b/control_plane/tests/plan_workload.rs @@ -88,3 +88,20 @@ fn load_rejects_a_plan_without_a_query_per_root() { )); std::fs::remove_file(path).unwrap(); } + +/// The backend keeps query-time results across evaluations, so for a few +/// series sampled every millisecond the Planner selects #509 Example 4's +/// B3: five 1-minute KLL summaries built in a query run and kept. +#[test] +fn selects_summaries_kept_across_evaluations() { + let inputs = DeploymentInputs::default(); + assert!(inputs.deployment_capabilities.query_time_results_retained); + let plan = plan_workload(&example("example4-b3"), &inputs).unwrap(); + let kept: Vec<_> = plan + .stored_outputs + .iter() + .filter(|o| o.built_at == asap_planner::StoredOutputBuild::QueryRun) + .collect(); + assert_eq!(kept.len(), 5, "five one-minute summaries"); + assert!(kept.iter().all(|o| o.pane.width_ms == 60_000)); +} diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 6c962658..17aa8653 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -12,7 +12,9 @@ use asap_planner::InstallablePlan; use tokio::sync::oneshot; use crate::precompute::{Precompute, PrecomputeOptions, PrecomputeStats, RawSample}; -use crate::query::{QueryEngine, RawRange, ResultSample, Unanswered}; +use crate::query::{ + IngestionProgress, KeptResultStats, Prepared, QueryEngine, ResultSample, Unanswered, +}; use crate::sds::Sds; /// How the engine thread ingests samples. @@ -42,11 +44,12 @@ enum Command { RawRanges { query: String, time_ms: i64, - reply: Reply<(usize, Vec)>, + reply: Reply, }, Evaluate { root: usize, time_ms: i64, + progress: IngestionProgress, raw: BTreeMap>, reply: Reply>, }, @@ -75,6 +78,8 @@ pub struct EngineHandle { running: Arc, /// The counters as of the last command the engine ran. stats: Arc>, + /// The query-run summary counters as of the last evaluation. + kept_stats: Arc>, } /// Marks the engine stopped when its thread exits, panics included. @@ -96,9 +101,11 @@ impl EngineHandle { let running = Arc::new(AtomicBool::new(true)); let queued_samples = Arc::new(AtomicU64::new(0)); let stats = Arc::new(Mutex::new(PrecomputeStats::default())); + let kept_stats = Arc::new(Mutex::new(KeptResultStats::default())); let shared = Shared { queued_samples: queued_samples.clone(), stats: stats.clone(), + kept_stats: kept_stats.clone(), }; let guard = RunningGuard(running.clone()); std::thread::Builder::new() @@ -122,6 +129,7 @@ impl EngineHandle { max_queued_samples: options.max_queued_samples, running, stats, + kept_stats, }; Ok((handle, queries)) } @@ -154,11 +162,7 @@ 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> { + pub async fn raw_ranges(&self, query: String, time_ms: i64) -> Result { self.ask(|reply| Command::RawRanges { query, time_ms, @@ -167,22 +171,29 @@ impl EngineHandle { .await } - /// Runs a root over stored panes and `raw`; see [`QueryEngine::evaluate`]. + /// Runs a prepared root over stored summaries and `raw`; see + /// [`QueryEngine::evaluate`]. pub async fn evaluate( &self, - root: usize, + prepared: Prepared, time_ms: i64, raw: BTreeMap>, ) -> Result, Unanswered> { self.ask(|reply| Command::Evaluate { - root, + root: prepared.root, time_ms, + progress: prepared.progress, raw, reply, }) .await } + /// The query-run summary counters as of the last evaluation. + pub fn kept_stats(&self) -> KeptResultStats { + *self.kept_stats.lock().unwrap_or_else(|e| e.into_inner()) + } + async fn ask(&self, command: impl FnOnce(Reply) -> Command) -> Result { let stopped = || Unanswered::Execution("the engine has stopped".into()); let (reply, answer) = oneshot::channel(); @@ -232,6 +243,7 @@ fn admit(queued: u64, count: u64, max_queued: u64) -> Result { struct Shared { queued_samples: Arc, stats: Arc>, + kept_stats: Arc>, } struct Engine { @@ -249,18 +261,6 @@ 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 keeps query-run summaries, 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 { @@ -289,20 +289,28 @@ impl Engine { &query, time_ms, &self.precompute.store, + self.precompute.progress(), )); } Command::Evaluate { root, time_ms, + progress, raw, reply, } => { - let _ = reply.send(self.query.evaluate( + let answer = self.query.evaluate( root, time_ms, - &self.precompute.store, + &mut self.precompute.store, + &progress, &raw, - )); + ); + // Published before the reply, so a caller that awaited + // the evaluation reads its counters. + *shared.kept_stats.lock().unwrap_or_else(|e| e.into_inner()) = + self.query.kept_stats; + let _ = reply.send(answer); } Command::Stop => break, } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 0c26989d..448a547b 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -176,6 +176,7 @@ async fn status(State(state): State>) -> Response { "engine_running": state.engine.is_running(), "samples_queued": state.engine.queued_samples(), "precompute": state.engine.stats(), + "query_run": state.engine.kept_stats(), }); ( [(header::CONTENT_TYPE, "application/json")], @@ -249,7 +250,8 @@ async fn answer( .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 prepared = state.engine.raw_ranges(query.query, time_ms).await?; + let ranges = &prepared.ranges; let mut fetch_headers = forwarded_headers(headers); // The raw read is a GET without the caller's body. fetch_headers.remove(header::CONTENT_TYPE); @@ -264,7 +266,7 @@ async fn answer( .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))? .into_iter() .collect(); - let samples = state.engine.evaluate(root, time_ms, raw).await?; + let samples = state.engine.evaluate(prepared, time_ms, raw).await?; Ok(prometheus::vector_response(time_ms, &samples)) } diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index 1a0b1b31..3775fdae 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -29,6 +29,8 @@ use asap_executor::{RawSampleSelector, Sample}; use futures::executor::block_on; use futures::StreamExt; +use asap_planner::StoredOutputBuild; + use crate::sds::{ DefinitionId, Sds, Selection, SummaryDefinition, SummaryInstanceKey, SummaryStore, }; @@ -68,6 +70,10 @@ pub struct PrecomputeStats { pub samples_dropped_future: u64, /// Samples of a metric no stored summary reads. pub samples_unrouted: u64, + /// Samples of a metric only query-run summaries read: they move the + /// watermark, which says when such a summary is closed, but are not + /// buffered (the query engine reads them from Prometheus). + pub samples_query_run_only: u64, /// Panes built and stored. pub panes_built: u64, /// Panes whose build failed; they are skipped, not retried. @@ -95,13 +101,84 @@ impl Default for PrecomputeOptions { } } -struct Builder { - definition: SummaryDefinition, - /// Turns a pane's samples into the build's input: the `Scan`'s label - /// matchers, then the pane's `TimeRange` rows. +/// How one definition's summaries are built from raw samples. The +/// precompute engine builds ingestion-run summaries with it, and the query +/// engine query-run ones. +pub struct SummaryBuild { + pub definition: SummaryDefinition, + /// Turns a coverage's samples into the build's input: the `Scan`'s label + /// matchers, then the coverage's `TimeRange` rows. selector: RawSampleSelector, raw_schema: SchemaRef, build: CompiledPhysicalDAG, +} + +impl SummaryBuild { + /// Compiles `definition`'s build from its `TimeRange` up to its builder. + /// A definition built from other stored outputs is refused until the + /// engine can build it once its inputs are sealed. + pub fn new( + dag: &PhysicalASAPDAG, + definition: &SummaryDefinition, + ) -> Result { + let error = |message: String| PrecomputeError::Definition(definition.id.clone(), message); + if !definition.stored_inputs.is_empty() { + return Err(error( + "built from other stored outputs, which the data plane cannot build yet".into(), + )); + } + if definition.pane_ms <= 0 { + return Err(error(format!("pane width {} ms", definition.pane_ms))); + } + let raw_schema = Arc::new(output_schema(dag, definition.raw_input()).clone()); + // `window_batch` is called with each coverage's end, so the window's + // offset from evaluation time is not used. + let selector = RawSampleSelector::new(dag, &definition.window)?; + let build = compile( + dag, + BTreeMap::from([( + definition.raw_input(), + InputContract::bounded(raw_schema.clone()), + )]), + &[definition.builder], + )?; + Ok(Self { + definition: definition.clone(), + selector, + raw_schema, + build, + }) + } + + /// The metric whose samples the build reads. + pub fn metric(&self) -> &str { + self.selector.metric() + } + + /// Selects the coverage ending at `end_ms` from `samples` with the plan's + /// `Scan` matchers, then runs the build over those rows under `scope`. + pub fn run( + &self, + end_ms: i64, + samples: &[RawSample], + scope: Scope, + ) -> Result { + let selected = self + .selector + .window_batch(end_ms, &executor_samples(samples))?; + let input: Source<'_> = + Box::new(Operator::source(self.raw_schema.clone(), vec![selected])?); + let build = self + .build + .instantiate(BTreeMap::from([(self.definition.raw_input(), input)]))?; + let rows = run_root(&build, self.definition.builder, scope)?; + let schema = Arc::new(self.definition.output_schema.clone()); + Ok(Batch::try_new(schema, rows)?) + } +} + +struct Builder { + summary: SummaryBuild, /// Samples of open panes, by pane end. open: BTreeMap>, /// End of the first pane seen, which is skipped. `None` until then. @@ -121,6 +198,9 @@ enum Outcome { pub struct Precompute { builders: Vec, + /// Metrics read by query-run summaries, with each one's newest routed + /// sample time. + query_run_metrics: BTreeMap>, allowed_lateness_ms: i64, max_future_ms: i64, newest_sample_ms: Option, @@ -147,43 +227,35 @@ impl Precompute { .ingestion_build_order() .map_err(|e| PrecomputeError::Plan(e.to_string()))?; for id in &order { - let definition = &sds.definitions[id]; - let error = - |message: String| PrecomputeError::Definition(definition.id.clone(), message); - if !definition.stored_inputs.is_empty() { - return Err(error( - "built from other stored outputs, which the precompute engine cannot build yet" - .into(), - )); - } - if definition.pane_ms <= 0 { - return Err(error(format!("pane width {} ms", definition.pane_ms))); - } - let raw_schema = Arc::new(output_schema(&dag, definition.raw_input()).clone()); - // `window_batch` is called with each pane's end, so the window's - // offset from evaluation time is not used. - let selector = RawSampleSelector::new(&dag, &definition.window)?; - let build = compile( - &dag, - BTreeMap::from([( - definition.raw_input(), - InputContract::bounded(raw_schema.clone()), - )]), - &[definition.builder], - )?; builders.push(Builder { - definition: definition.clone(), - selector, - raw_schema, - build, + summary: SummaryBuild::new(&dag, &sds.definitions[id])?, open: BTreeMap::new(), first_pane_end: None, built_through: None, }); } + // Samples of a metric only query-run summaries read move the + // watermark, which says when such a summary's coverage is closed. + // A query-run definition whose selector does not build is not + // tracked here; the query engine forwards the queries that read it. + let mut query_run_metrics = BTreeMap::new(); + for definition in sds.definitions.values() { + if definition.built_at == StoredOutputBuild::QueryRun { + match RawSampleSelector::new(&dag, &definition.window) { + Ok(selector) => { + query_run_metrics.insert(selector.metric().to_string(), None); + } + Err(error) => tracing::warn!( + definition = ?definition.id, %error, + "query-run summary cannot be built; its queries are forwarded" + ), + } + } + } let millis = |value: u64| i64::try_from(value).unwrap_or(i64::MAX); Ok(Self { builders, + query_run_metrics, allowed_lateness_ms: millis(options.allowed_lateness_ms), max_future_ms: millis(options.max_future_ms), newest_sample_ms: None, @@ -193,6 +265,19 @@ impl Precompute { }) } + /// What ingestion has seen, for the query engine's keep decisions: the + /// watermark and each query-run metric's newest sample. + pub fn progress(&self) -> crate::query::IngestionProgress { + crate::query::IngestionProgress { + watermark_ms: self.watermark_ms(), + newest_sample_ms: self + .query_run_metrics + .iter() + .filter_map(|(metric, newest)| Some((metric.clone(), (*newest)?))) + .collect(), + } + } + /// The watermark: no sample older than this is expected any more. It is /// the newest routed sample, capped at the wall clock, minus the allowed /// lateness. @@ -220,10 +305,12 @@ impl Precompute { let latest_allowed_ms = now_ms.saturating_add(self.max_future_ms); for sample in samples { let metric = sample.labels.get("__name__").map(String::as_str); - if !self - .builders - .iter() - .any(|b| metric == Some(b.selector.metric())) + let query_run = metric.is_some_and(|m| self.query_run_metrics.contains_key(m)); + if !query_run + && !self + .builders + .iter() + .any(|b| metric == Some(b.summary.metric())) { self.stats.samples_unrouted += 1; continue; @@ -234,10 +321,10 @@ impl Precompute { } let mut outcome: Option = None; for builder in &mut self.builders { - if metric != Some(builder.selector.metric()) { + if metric != Some(builder.summary.metric()) { continue; } - let this = match pane_end(sample.timestamp_ms, builder.definition.pane_ms) { + let this = match pane_end(sample.timestamp_ms, builder.summary.definition.pane_ms) { None => Outcome::Late, Some(end) => { let first = *builder.first_pane_end.get_or_insert(end); @@ -255,15 +342,14 @@ impl Precompute { outcome = outcome.max(Some(this)); } match outcome { - None => self.stats.samples_unrouted += 1, - Some(outcome) => { - match outcome { - Outcome::Buffered => self.stats.samples_accepted += 1, - Outcome::FirstPane => self.stats.samples_skipped_first_pane += 1, - Outcome::Late => self.stats.samples_dropped_late += 1, - } - self.newest_sample_ms = self.newest_sample_ms.max(Some(sample.timestamp_ms)); - } + None => self.stats.samples_query_run_only += 1, + Some(Outcome::Buffered) => self.stats.samples_accepted += 1, + Some(Outcome::FirstPane) => self.stats.samples_skipped_first_pane += 1, + Some(Outcome::Late) => self.stats.samples_dropped_late += 1, + } + self.newest_sample_ms = self.newest_sample_ms.max(Some(sample.timestamp_ms)); + if let Some(newest) = metric.and_then(|m| self.query_run_metrics.get_mut(m)) { + *newest = (*newest).max(Some(sample.timestamp_ms)); } } self.close_panes(); @@ -277,8 +363,8 @@ impl Precompute { let Some(built) = builder.built_through else { continue; }; - let pane_ms = builder.definition.pane_ms; - let id = &builder.definition.id; + let pane_ms = builder.summary.definition.pane_ms; + let id = &builder.summary.definition.id; // The newest evaluation stored panes can serve is at the end of // the newest closed pane. Panes are kept for one extra pane // width beyond the plan's retention, so an evaluation one pane @@ -286,6 +372,7 @@ impl Precompute { // skip them. let newest_closed_end = watermark.div_euclid(pane_ms).saturating_mul(pane_ms); let oldest_kept_end = builder + .summary .definition .oldest_kept_end_ms(newest_closed_end.saturating_sub(pane_ms)); let Some(mut end) = built.checked_add(pane_ms) else { @@ -304,7 +391,12 @@ impl Precompute { // a pane without samples is stored as empty, not missing. while end <= watermark { let samples = builder.open.remove(&end).unwrap_or_default(); - match build_pane(builder, end, &samples) { + let scope = Scope::Ingestion { + window_start_ms: end - pane_ms, + window_end_ms: end, + revision: 1, + }; + match builder.summary.run(end, &samples, scope) { Ok(state) => { self.store.insert( SummaryInstanceKey::new(id.clone(), Selection::ending_at(end)), @@ -341,33 +433,6 @@ pub fn pane_end(timestamp_ms: i64, pane_ms: i64) -> Option { .checked_mul(pane_ms) } -/// Selects the pane's rows from `samples` with the plan's `Scan` matchers, -/// then runs the definition's build over them, for the pane ending at `end`. -fn build_pane( - builder: &Builder, - end: i64, - samples: &[RawSample], -) -> Result { - let scope = Scope::Ingestion { - window_start_ms: end - builder.definition.pane_ms, - window_end_ms: end, - revision: 1, - }; - let selected = builder - .selector - .window_batch(end, &executor_samples(samples))?; - let input: Source<'_> = Box::new(Operator::source( - builder.raw_schema.clone(), - vec![selected], - )?); - let build = builder - .build - .instantiate(BTreeMap::from([(builder.definition.raw_input(), input)]))?; - let rows = run_root(&build, builder.definition.builder, scope)?; - let schema = Arc::new(builder.definition.output_schema.clone()); - 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 { diff --git a/data_plane/src/query.rs b/data_plane/src/query.rs index bca9bffb..f0999097 100644 --- a/data_plane/src/query.rs +++ b/data_plane/src/query.rs @@ -1,12 +1,18 @@ //! 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`. +//! Each root is compiled once with three kinds of inputs: +//! - stored outputs built in an ingestion run: bound to their definition's +//! summary ending at `T - offset`, which must be stored; +//! - stored outputs built in a query run (`StoredOutputBuild::QueryRun`): +//! bound to their definition's summary ending at `T - offset`. A missing +//! one is built from its coverage's raw samples, which the caller fetches, +//! and kept for later evaluations once its coverage is closed (its end at +//! or before the precompute engine's watermark); kept summaries are +//! evicted by the plan's retention, measured from the newest evaluation; +//! - the plan's query-time raw windows: bound to the samples in +//! `(T - offset - width, T - offset]`, which the caller fetches, selected +//! by the window's `RawSampleSelector`. //! //! A query the engine cannot answer returns [`Unanswered`]; the caller then //! forwards it to Prometheus. @@ -23,10 +29,10 @@ 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 asap_planner::{InstallablePlan, StoredOutputBuild}; -use crate::precompute::{executor_samples, run_root, RawSample}; -use crate::sds::{DefinitionId, Sds, Selection, SummaryStore}; +use crate::precompute::{executor_samples, run_root, RawSample, SummaryBuild}; +use crate::sds::{DefinitionId, Sds, Selection, SummaryInstanceKey, SummaryStore}; /// Raw samples a query needs: `metric`'s samples in `(start_ms, end_ms]` /// that `matchers` keep. @@ -111,14 +117,73 @@ struct RawInput { offset_ms: i64, } +/// A stored output built in a query run and kept. +struct KeptInput { + node: u64, + definition: DefinitionId, + offset_ms: i64, + /// The output's own `TimeRange`, which its coverage's samples feed. + raw_input: u64, +} + struct Program { compiled: CompiledPhysicalDAG, stored: Vec, + kept: Vec, + /// For each query-run definition the root reads, the node its one + /// fetch is keyed by: the lowest `TimeRange` among its outputs. + kept_fetch_node: BTreeMap, raw: Vec, output: Schema, } +/// What happened to the query-run summaries evaluations read. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, serde::Serialize)] +pub struct KeptResultStats { + /// Summaries built from raw samples because none was kept. + pub built: u64, + /// Of those, the ones kept: their coverage was closed. + pub kept: u64, + /// Of those kept, the ones whose own metric's newest sample had not + /// passed the coverage's end: the watermark is global, so a metric + /// lagging behind others can have a summary kept before all its samples + /// arrived. A per-metric watermark is the follow-up. + pub kept_while_metric_behind: u64, + /// Summaries read from an earlier evaluation instead of rebuilt. + pub reused: u64, +} + +/// What ingestion had seen when an evaluation decided what to fetch. The +/// keep decision uses this snapshot, not the state when the evaluation +/// runs: samples ingested in between were not in the fetched data. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct IngestionProgress { + /// The precompute engine's watermark; `None` before any sample. + pub watermark_ms: Option, + /// The newest sample time of each metric query-run summaries read. + pub newest_sample_ms: BTreeMap, +} + +/// One raw read an evaluation needs, prepared by [`QueryEngine::raw_ranges`]. +#[derive(Debug, Clone)] +pub struct Prepared { + /// The root answering the query. + pub root: usize, + /// The raw samples to fetch. + pub ranges: Vec, + /// What ingestion had seen when `ranges` were chosen. + pub progress: IngestionProgress, +} + pub struct QueryEngine { + /// How many query-run summaries evaluations built, kept and reused. + pub kept_stats: KeptResultStats, + /// How each query-run definition's summaries are built, or why they + /// cannot be. + kept_builds: BTreeMap>, + /// The label matchers of each query-run definition's `Scan`, pushed into + /// its fetch. + kept_matchers: BTreeMap>, /// Each query, with its compiled root or why it cannot run. roots: Vec<(String, Result)>, } @@ -131,30 +196,56 @@ impl QueryEngine { .iter() .map(|w| (w.time_range_node_id as u64, w)) .collect(); + let kept_builds: BTreeMap<_, _> = sds + .definitions + .values() + .filter(|d| d.built_at == StoredOutputBuild::QueryRun) + .map(|d| { + let build = SummaryBuild::new(dag, d).map_err(|e| e.to_string()); + (d.id.clone(), build) + }) + .collect(); let roots = plan .queries .iter() .map(|q| { - let program = program(dag, sds, &windows, q.root_node_id as u64); + let program = program(dag, sds, &kept_builds, &windows, q.root_node_id as u64); (q.query.trim().to_string(), program) }) .collect(); - Self { roots } + let kept_matchers = kept_builds + .keys() + .map(|id| { + let scan = sds.definitions[id].window.scan_node_id as u64; + (id.clone(), label_matchers(dag, scan)) + }) + .collect(); + Self { + kept_stats: KeptResultStats::default(), + kept_builds, + kept_matchers, + 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. + /// The root answering `query` at `time_ms`, and the raw samples it needs: + /// its query-time windows, and for each query-run definition one range + /// covering every summary it reads that is not stored. Fails if the + /// query is not planned, a time does not fit, or an ingestion-built + /// summary it reads is missing. `progress` is what ingestion has seen + /// now; pass the returned copy to [`Self::evaluate`]. pub fn raw_ranges( &self, query: &str, time_ms: i64, store: &SummaryStore, - ) -> Result<(usize, Vec), Unanswered> { + progress: IngestionProgress, + ) -> Result { let (root, program) = self.program(query)?; for input in &program.stored { pane(input, time_ms, store)?; } - let ranges = program + let mut ranges: Vec = program .raw .iter() .map(|input| { @@ -168,16 +259,53 @@ impl QueryEngine { }) }) .collect::>()?; - Ok((root, ranges)) + // One fetch per query-run definition, covering all its missing + // summaries. + let mut missing: BTreeMap<&DefinitionId, (i64, i64)> = BTreeMap::new(); + for input in &program.kept { + let build = self.kept_build(&input.definition); + let (start_ms, end_ms) = window(time_ms, input.offset_ms, build.definition.pane_ms)?; + if store + .get(&input.definition, Selection::ending_at(end_ms)) + .is_none() + { + let span = missing + .entry(&input.definition) + .or_insert((start_ms, end_ms)); + *span = (span.0.min(start_ms), span.1.max(end_ms)); + } + } + for (definition, (start_ms, end_ms)) in missing { + let build = self.kept_build(definition); + ranges.push(RawRange { + node: program.kept_fetch_node[definition], + metric: build.metric().to_string(), + matchers: self.kept_matchers[definition].clone(), + start_ms, + end_ms, + }); + } + Ok(Prepared { + root, + ranges, + progress, + }) } - /// Runs `root` at `time_ms` over stored panes and `raw`, the samples of - /// each of its [`RawRange`]s by node. + /// Runs `root` at `time_ms` over stored summaries and `raw`, the samples + /// of each of its [`RawRange`]s by node. A query-run summary not stored + /// is built from its raw samples and stored in `store`, under the same + /// key as ingestion-built ones, if its coverage was closed in + /// `progress` (the snapshot [`Self::raw_ranges`] chose the fetch by). + /// Then query-run summaries older than their definition's retention, + /// measured from the newest closed pane at or before the watermark, are + /// evicted. pub fn evaluate( - &self, + &mut self, root: usize, time_ms: i64, - store: &SummaryStore, + store: &mut SummaryStore, + progress: &IngestionProgress, raw: &BTreeMap>, ) -> Result, Unanswered> { let program = self.roots[root] @@ -196,6 +324,70 @@ impl QueryEngine { Operator::source(pane.schema().clone(), vec![pane.clone()]).map_err(failed)?; sources.insert(input.node, Box::new(source)); } + for input in &program.kept { + let build = self.kept_builds[&input.definition] + .as_ref() + .expect("compiled roots read only query-run definitions that build"); + let (_, end_ms) = window(time_ms, input.offset_ms, build.definition.pane_ms)?; + let selection = Selection::ending_at(end_ms); + let summary = match store.get(&input.definition, selection) { + Some(summary) => { + self.kept_stats.reused += 1; + summary.clone() + } + None => { + // `raw_ranges` asked for these samples unless the + // summary was stored then and evicted since. + let samples = raw + .get(&program.kept_fetch_node[&input.definition]) + .ok_or_else(|| { + Unanswered::MissingPane(format!( + "query-run summary ending at {end_ms} of {:?} was evicted", + input.definition + )) + })?; + // Built as at its coverage's end, so every evaluation + // that reads it reads the same summary. + let build_scope = Scope::Query { + evaluation_time_ms: end_ms, + revision: 1, + }; + let summary = build + .run(end_ms, samples, build_scope) + .map_err(|e| Unanswered::Execution(e.to_string()))?; + self.kept_stats.built += 1; + // Kept only once closed, as the precompute engine closes + // panes: a sample may still arrive for an open one. + if progress.watermark_ms.is_some_and(|w| end_ms <= w) { + self.kept_stats.kept += 1; + let metric_newest = progress.newest_sample_ms.get(build.metric()); + if metric_newest.is_none_or(|newest| *newest < end_ms) { + self.kept_stats.kept_while_metric_behind += 1; + } + store.insert( + SummaryInstanceKey::new(input.definition.clone(), selection), + summary.clone(), + ); + } + summary + } + }; + let source = + Operator::source(summary.schema().clone(), vec![summary]).map_err(failed)?; + sources.insert(input.node, Box::new(source)); + } + if let Some(watermark) = progress.watermark_ms { + for build in self.kept_builds.values().flatten() { + let definition = &build.definition; + let pane_ms = definition.pane_ms; + // As for ingestion-built summaries: measured from the newest + // closed coverage, one pane of margin. + let newest_closed_end = watermark.div_euclid(pane_ms).saturating_mul(pane_ms); + let oldest = + definition.oldest_kept_end_ms(newest_closed_end.saturating_sub(pane_ms)); + store.evict_before(&definition.id, Selection::ending_at(oldest)); + } + } for input in &program.raw { let samples = raw.get(&input.node).map(Vec::as_slice).unwrap_or_default(); let rows = input @@ -213,6 +405,13 @@ impl QueryEngine { result_samples(&program.output, rows) } + /// The build of a query-run definition a compiled root reads. + fn kept_build(&self, definition: &DefinitionId) -> &SummaryBuild { + self.kept_builds[definition] + .as_ref() + .expect("compiled roots read only query-run definitions that build") + } + fn program(&self, query: &str) -> Result<(usize, &Program), Unanswered> { let query = query.trim(); let root = self @@ -264,6 +463,7 @@ fn pane<'a>( fn program( dag: &PhysicalASAPDAG, sds: &Sds, + kept_builds: &BTreeMap>, windows: &HashMap, root: u64, ) -> Result { @@ -275,7 +475,7 @@ fn program( .or_default() .push(edge.producer as u64); } - let (mut stored, mut raw) = (Vec::new(), Vec::new()); + let (mut stored, mut kept, mut raw) = (Vec::new(), Vec::new(), Vec::new()); let mut pending = vec![root]; let mut seen = BTreeSet::new(); while let Some(id) = pending.pop() { @@ -283,12 +483,23 @@ fn program( 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, - }); + match sds.definitions[definition].built_at { + StoredOutputBuild::IngestionRun => stored.push(StoredInput { + node: id, + definition: definition.clone(), + offset_ms: binding.offset_ms, + pane_ms: sds.definitions[definition].pane_ms, + }), + StoredOutputBuild::QueryRun => { + kept_builds[definition].as_ref().map_err(Clone::clone)?; + kept.push(KeptInput { + node: id, + definition: definition.clone(), + offset_ms: binding.offset_ms, + raw_input: binding.raw_input, + }); + } + } continue; } match &nodes[&id].payload { @@ -315,17 +526,30 @@ fn program( } } 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 node in stored + .iter() + .map(|i| i.node) + .chain(kept.iter().map(|i| i.node)) + { + let schema = Arc::new(nodes[&node].output_schema.clone()); + inputs.insert(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())?; + let mut kept_fetch_node: BTreeMap = BTreeMap::new(); + for input in &kept { + let node = kept_fetch_node + .entry(input.definition.clone()) + .or_insert(input.raw_input); + *node = (*node).min(input.raw_input); + } Ok(Program { compiled, stored, + kept, + kept_fetch_node, raw, output: nodes[&root].output_schema.clone(), }) diff --git a/data_plane/src/sds.rs b/data_plane/src/sds.rs index b5354310..f45b37f2 100644 --- a/data_plane/src/sds.rs +++ b/data_plane/src/sds.rs @@ -85,6 +85,9 @@ impl SummaryDefinition { #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct PaneBinding { pub offset_ms: i64, + /// The output's own `TimeRange` node: where its coverage's raw samples + /// feed when a query-run summary is built for it. + pub raw_input: u64, } /// The stored summaries of one plan. @@ -153,6 +156,7 @@ impl Sds { definition.retention_ms = definition.retention_ms.max(retention_ms); let binding = PaneBinding { offset_ms: output.pane.end_offset_ms, + raw_input: output.pane.time_range_node_id as u64, }; outputs.insert(node, (id, binding)); } diff --git a/data_plane/tests/kept.rs b/data_plane/tests/kept.rs new file mode 100644 index 00000000..47496865 --- /dev/null +++ b/data_plane/tests/kept.rs @@ -0,0 +1,356 @@ +//! Query-run summaries kept across evaluations +//! (`MatAtQueryKeptAcrossEvaluations`): #509 Example 4's B3, five 1-minute +//! KLL summaries built at query time and stored for the next evaluations of +//! `quantile_over_time(0.99, latency_ms[5m])` every minute. + +mod common; + +use std::collections::BTreeMap; +use std::sync::Arc; + +use asap_executor::values::Batch; +use asap_planner::{InstallablePlan, StoredOutputBuild}; +use data_plane::engine::{EngineHandle, EngineOptions}; +use data_plane::precompute::RawSample; +use data_plane::query::{IngestionProgress, KeptResultStats, QueryEngine, RawRange, ResultSample}; +use data_plane::sds::{DefinitionId, Sds, Selection, SummaryStore}; + +const MINUTE: i64 = 60_000; +const P99: &str = "quantile_over_time(0.99, latency_ms[5m])"; + +/// A query engine and the SDS store its query-run summaries go to. +struct Kept { + engine: QueryEngine, + store: SummaryStore, + definition: DefinitionId, +} + +fn kept(plan: &InstallablePlan) -> Kept { + let sds = Sds::from_plan(plan).unwrap(); + assert!( + sds.ingestion_build_order().unwrap().is_empty(), + "nothing is built at ingestion" + ); + let definition = sds + .definitions + .values() + .find(|d| d.built_at == StoredOutputBuild::QueryRun) + .map(|d| d.id.clone()) + .unwrap_or_else(|| DefinitionId(String::new())); + Kept { + engine: QueryEngine::new(plan, &sds), + store: SummaryStore::new(), + definition, + } +} + +impl Kept { + /// The ends, in minutes, of the stored query-run summaries. + fn ends(&self) -> Vec { + (0..=30) + .map(|m| m * MINUTE) + .filter(|end| { + self.store + .get(&self.definition, Selection::ending_at(*end)) + .is_some() + }) + .map(|end| end / MINUTE) + .collect() + } +} + +/// Two series every 15 s in `[0, 30 min]`, the `n`th sample `25n mod 101` +/// (plus 100 for series b): each 5-minute window's p99 falls in an older +/// pane and changes every minute. KLL is exact below k. +fn samples() -> Vec { + let series: Vec<_> = ["a", "b"] + .iter() + .map(|job| { + Arc::new(BTreeMap::from([ + ("__name__".to_string(), "latency_ms".to_string()), + ("job".to_string(), job.to_string()), + ])) + }) + .collect(); + (0..=30 * MINUTE) + .step_by(15_000) + .flat_map(|t| { + series.iter().enumerate().map(move |(i, labels)| RawSample { + labels: labels.clone(), + timestamp_ms: t, + value: (i * 100) as f64 + (t / 15_000 * 25 % 101) as f64, + }) + }) + .collect() +} + +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() +} + +/// Answers `P99` at `time_ms` with ingestion's watermark at `watermark_ms` +/// (every metric's newest sample there too), returning the answer and the +/// raw ranges it read. +fn answer_at_watermark( + k: &mut Kept, + time_ms: i64, + watermark_ms: Option, + all: &[RawSample], +) -> (Vec, Vec) { + let progress = IngestionProgress { + watermark_ms, + newest_sample_ms: watermark_ms + .map(|w| BTreeMap::from([("latency_ms".to_string(), w)])) + .unwrap_or_default(), + }; + let prepared = k + .engine + .raw_ranges(P99, time_ms, &k.store, progress) + .unwrap(); + let ranges = prepared.ranges.clone(); + let mut result = k + .engine + .evaluate( + prepared.root, + time_ms, + &mut k.store, + &prepared.progress, + &fetch(&ranges, all), + ) + .unwrap(); + result.sort_by(|a, b| a.labels.cmp(&b.labels)); + (result, ranges) +} + +/// Answers `P99` at `time_ms` once every sample through `time_ms` has +/// arrived: the watermark is `time_ms`. +fn answer(k: &mut Kept, time_ms: i64, all: &[RawSample]) -> (Vec, Vec) { + answer_at_watermark(k, time_ms, Some(time_ms), all) +} + +fn spans(ranges: &[RawRange]) -> Vec<(i64, i64)> { + ranges + .iter() + .map(|r| (r.start_ms / MINUTE, r.end_ms / MINUTE)) + .collect() +} + +/// The installed B3 plan and the Planner's reference for it: the same +/// selection with nothing materialized. +fn plans() -> (InstallablePlan, InstallablePlan) { + let plan = common::plan("example4-b3", true); + let reference = plan.query_time_reference().unwrap(); + (plan, reference) +} + +/// The first evaluation builds all five summaries from one fetch and +/// answers like the reference; the next one, a minute later, fetches only +/// the newest minute, builds one summary, reuses four, and still answers +/// like the reference. +#[test] +fn kept_summaries_answer_like_the_reference_and_are_reused() { + let (plan, reference_plan) = plans(); + let all = samples(); + let mut k = kept(&plan); + let mut reference = kept(&reference_plan); + + let (first, ranges) = answer(&mut k, 10 * MINUTE, &all); + assert_eq!( + spans(&ranges), + vec![(5, 10)], + "one fetch for five summaries" + ); + assert_eq!(first, answer(&mut reference, 10 * MINUTE, &all).0); + assert_eq!(first.len(), 2); + assert_eq!( + k.engine.kept_stats, + KeptResultStats { + built: 5, + kept: 5, + kept_while_metric_behind: 0, + reused: 0 + } + ); + + let (second, ranges) = answer(&mut k, 11 * MINUTE, &all); + assert_eq!(spans(&ranges), vec![(10, 11)], "only the newest summary"); + assert_eq!(second, answer(&mut reference, 11 * MINUTE, &all).0); + assert_ne!(second, first, "the answer moves with the window"); + assert_eq!( + (k.engine.kept_stats.built, k.engine.kept_stats.reused), + (6, 4) + ); + + // The same evaluation again builds nothing. + let (again, ranges) = answer(&mut k, 11 * MINUTE, &all); + assert!(ranges.is_empty()); + assert_eq!(again, second); + assert_eq!( + (k.engine.kept_stats.built, k.engine.kept_stats.reused), + (6, 9) + ); + assert_eq!( + reference.engine.kept_stats, + KeptResultStats::default(), + "the reference keeps nothing" + ); +} + +/// Stored query-run summaries are evicted by the plan's retention (5 min, +/// the oldest summary a 5-minute window reads), plus one pane, measured +/// from the newest closed coverage at or before the watermark, as for +/// ingestion-built summaries. A far-future evaluation does not move that +/// point: its summaries are built but not kept, and keeping resumes. +#[test] +fn kept_summaries_are_evicted_by_the_plans_retention() { + let (plan, _) = plans(); + let retention = plan.stored_outputs.iter().map(|o| o.retention_ms).max(); + assert_eq!(retention, Some(5 * MINUTE as u64)); + let all = samples(); + let mut k = kept(&plan); + answer(&mut k, 10 * MINUTE, &all); + assert_eq!(k.ends(), vec![6, 7, 8, 9, 10]); + answer(&mut k, 11 * MINUTE, &all); + assert_eq!(k.ends(), vec![6, 7, 8, 9, 10, 11]); + answer(&mut k, 12 * MINUTE, &all); + assert_eq!(k.ends(), vec![7, 8, 9, 10, 11, 12]); + // An evaluation far ahead of the watermark is answered, but its + // summaries' coverages are not closed, so nothing is kept or evicted. + let (far, ranges) = answer_at_watermark(&mut k, 20 * MINUTE, Some(12 * MINUTE), &all); + assert_eq!(far.len(), 2); + assert_eq!(spans(&ranges), vec![(15, 20)]); + assert_eq!(k.ends(), vec![7, 8, 9, 10, 11, 12]); + // Keeping resumes with the watermark. + answer(&mut k, 13 * MINUTE, &all); + assert_eq!(k.ends(), vec![8, 9, 10, 11, 12, 13]); +} + +/// A summary is kept only once its coverage is closed (its end at or +/// before the watermark). One read before then is built for that +/// evaluation but not kept, so a later evaluation builds it again and +/// includes a sample that arrived late. +#[test] +fn summaries_are_kept_only_once_closed() { + let (plan, reference_plan) = plans(); + let mut all = samples(); + let mut k = kept(&plan); + let mut reference = kept(&reference_plan); + + // Watermark at 9.5 min: the summary ending at 10 min is still open. + let (first, ranges) = answer_at_watermark(&mut k, 10 * MINUTE, Some(570_000), &all); + assert_eq!(spans(&ranges), vec![(5, 10)]); + assert_eq!(first, answer(&mut reference, 10 * MINUTE, &all).0); + assert_eq!(k.ends(), vec![6, 7, 8, 9], "the open summary is not kept"); + + // A late sample for the open coverage arrives before it closes. + let late = RawSample { + timestamp_ms: 10 * MINUTE - 1_000, + value: 1_000.0, + ..all.iter().find(|s| s.labels["job"] == "a").unwrap().clone() + }; + all.push(late); + let (again, ranges) = answer_at_watermark(&mut k, 10 * MINUTE, Some(585_000), &all); + assert_eq!( + spans(&ranges), + vec![(9, 10)], + "the open summary is built again" + ); + assert_eq!(again, answer(&mut reference, 10 * MINUTE, &all).0); + assert_ne!(again, first, "the late sample is included"); + + // Once the watermark reaches 10 min, the summary is kept. + answer_at_watermark(&mut k, 10 * MINUTE, Some(10 * MINUTE), &all); + assert_eq!(k.ends(), vec![6, 7, 8, 9, 10]); + // Without any sample there is no watermark, and nothing is kept. + let mut fresh = kept(&plan); + answer_at_watermark(&mut fresh, 10 * MINUTE, None, &all); + assert!(fresh.store.is_empty()); +} + +/// The keep decision uses the watermark when the evaluation chose what to +/// fetch, not when it runs: ingestion that closes a coverage after the +/// fetch was chosen does not make the partial summary built from that +/// fetch permanent. +#[tokio::test] +async fn keeps_by_the_watermark_the_fetch_was_chosen_at() { + let (plan, _) = plans(); + let options = EngineOptions { + precompute: data_plane::precompute::PrecomputeOptions { + allowed_lateness_ms: 0, + ..Default::default() + }, + ..EngineOptions::default() + }; + let (engine, _) = EngineHandle::spawn(serde_json::to_string(&plan).unwrap(), options).unwrap(); + let all = samples(); + let through = |t: i64| -> Vec { + all.iter() + .filter(|s| s.timestamp_ms <= t) + .cloned() + .collect() + }; + let settle = |engine: &EngineHandle| { + let engine = engine.clone(); + async move { + while engine.queued_samples() > 0 { + tokio::time::sleep(std::time::Duration::from_millis(5)).await; + } + } + }; + engine.try_ingest(through(570_000)).unwrap(); + settle(&engine).await; + let prepared = engine.raw_ranges(P99.into(), 10 * MINUTE).await.unwrap(); + assert_eq!(prepared.progress.watermark_ms, Some(570_000)); + // The fetch sees only what had arrived; then ingestion closes 10 min. + let raw = fetch(&prepared.ranges, &through(570_000)); + engine.try_ingest(through(11 * MINUTE)).unwrap(); + settle(&engine).await; + engine.evaluate(prepared, 10 * MINUTE, raw).await.unwrap(); + assert_eq!(engine.kept_stats().kept, 4, "6..9 kept, 10 not"); + // So the next evaluation builds the summary ending at 10 min again. + let prepared = engine.raw_ranges(P99.into(), 10 * MINUTE).await.unwrap(); + assert_eq!(spans(&prepared.ranges), vec![(9, 10)]); +} + +/// A summary kept while its own metric's newest sample has not reached its +/// end (the global watermark moved on another metric) is counted. +#[test] +fn counts_summaries_kept_while_their_metric_is_behind() { + let (plan, _) = plans(); + let all = samples(); + let mut k = kept(&plan); + let progress = IngestionProgress { + watermark_ms: Some(10 * MINUTE), + newest_sample_ms: BTreeMap::from([("latency_ms".to_string(), 8 * MINUTE)]), + }; + let prepared = k + .engine + .raw_ranges(P99, 10 * MINUTE, &k.store, progress) + .unwrap(); + let raw = fetch(&prepared.ranges, &all); + k.engine + .evaluate( + prepared.root, + 10 * MINUTE, + &mut k.store, + &prepared.progress, + &raw, + ) + .unwrap(); + assert_eq!(k.engine.kept_stats.kept, 5); + assert_eq!(k.engine.kept_stats.kept_while_metric_behind, 2, "9 and 10"); +} diff --git a/data_plane/tests/precompute.rs b/data_plane/tests/precompute.rs index 34c1622d..c661e181 100644 --- a/data_plane/tests/precompute.rs +++ b/data_plane/tests/precompute.rs @@ -432,7 +432,7 @@ fn does_not_build_query_run_summaries() { let (mut engine, _) = precompute(plan); engine.ingest(&samples("latency_ms", &["a"], 0, 45 * MINUTE)); assert_eq!(engine.stats.panes_built, 0); - assert_eq!(engine.stats.samples_unrouted, 180); + assert_eq!(engine.stats.samples_query_run_only, 180); } /// Retention is measured from the newest evaluation the stored panes can @@ -473,3 +473,31 @@ fn keeps_one_pane_of_margin_beyond_the_retention() { // the one more the evaluation at 280 reads. assert_eq!(engine.store.len(), 7); } + +/// Samples of a metric only query-run summaries read (B3) are not buffered +/// but move the watermark, which says when those summaries are closed. +#[test] +fn query_run_metrics_move_the_watermark() { + let (mut engine, _) = precompute_with(common::plan("example4-b3", true), 0); + assert_eq!(engine.watermark_ms(), None); + engine.ingest(&samples("latency_ms", &["a"], 0, 3 * MINUTE)); + assert_eq!(engine.stats.samples_query_run_only, 12); + assert_eq!(engine.stats.samples_accepted, 0); + assert_eq!(engine.watermark_ms(), Some(3 * MINUTE - 15_000)); +} + +/// Samples of a query-run-only metric get the same future bound and +/// wall-clock cap as routed samples: one far ahead is dropped, and one +/// ahead of the clock but within the bound moves the watermark only to the +/// clock. +#[test] +fn query_run_metrics_are_bounded_like_routed_samples() { + let (mut engine, _) = precompute_with(common::plan("example4-b3", true), 0); + let now = 10 * MINUTE; + engine.ingest_at(&[sample("latency_ms", "a", 30 * DAY, 1.0)], now); + assert_eq!(engine.stats.samples_dropped_future, 1); + assert_eq!(engine.watermark_ms(), None); + engine.ingest_at(&[sample("latency_ms", "a", now + 2 * MINUTE, 1.0)], now); + assert_eq!(engine.stats.samples_query_run_only, 1); + assert_eq!(engine.watermark_ms(), Some(now)); +} diff --git a/data_plane/tests/query.rs b/data_plane/tests/query.rs index a6775ef8..ee1a6d1c 100644 --- a/data_plane/tests/query.rs +++ b/data_plane/tests/query.rs @@ -79,14 +79,21 @@ fn fetch(ranges: &[RawRange], all: &[RawSample]) -> BTreeMap .collect() } -fn answer(f: &Fixture, query: &str, time_ms: i64, all: &[RawSample]) -> Vec { - let (root, ranges) = f +fn answer(f: &mut Fixture, query: &str, time_ms: i64, all: &[RawSample]) -> Vec { + let prepared = f .query - .raw_ranges(query, time_ms, &f.precompute.store) + .raw_ranges(query, time_ms, &f.precompute.store, f.precompute.progress()) .unwrap(); + let raw = fetch(&prepared.ranges, all); let mut result = f .query - .evaluate(root, time_ms, &f.precompute.store, &fetch(&ranges, all)) + .evaluate( + prepared.root, + time_ms, + &mut f.precompute.store, + &prepared.progress, + &raw, + ) .unwrap(); result.sort_by(|a, b| a.labels.cmp(&b.labels)); result @@ -110,23 +117,40 @@ fn stored_panes_answer_like_a_single_build() { assert!( panes .query - .raw_ranges(P99, 70 * MINUTE, &panes.precompute.store) + .raw_ranges( + P99, + 70 * MINUTE, + &panes.precompute.store, + Default::default() + ) .is_err(), "no panes before ingestion" ); panes.precompute.ingest(&all); - let (_, ranges) = panes + let ranges = panes .query - .raw_ranges(P99, 70 * MINUTE, &panes.precompute.store) - .unwrap(); + .raw_ranges( + P99, + 70 * MINUTE, + &panes.precompute.store, + Default::default(), + ) + .unwrap() + .ranges; assert!(ranges.is_empty(), "every input is a stored pane"); - let from_panes = answer(&panes, P99, 70 * MINUTE, &all); + let from_panes = answer(&mut panes, P99, 70 * MINUTE, &all); - let single = fixture(common::plan("example4b", true)); - let (_, ranges) = single + let mut single = fixture(common::plan("example4b", true)); + let ranges = single .query - .raw_ranges(P99, 70 * MINUTE, &single.precompute.store) - .unwrap(); + .raw_ranges( + P99, + 70 * MINUTE, + &single.precompute.store, + Default::default(), + ) + .unwrap() + .ranges; assert_eq!( ranges .iter() @@ -135,7 +159,7 @@ fn stored_panes_answer_like_a_single_build() { vec![(10 * MINUTE, 70 * MINUTE)], "one raw hour, left-open" ); - let from_raw = answer(&single, P99, 70 * MINUTE, &all); + let from_raw = answer(&mut single, P99, 70 * MINUTE, &all); assert_eq!(from_panes, from_raw); assert_eq!(from_panes.len(), 2); @@ -157,9 +181,9 @@ fn query_time_plans_read_raw_ranges() { ("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 mut f = fixture(common::plan("example1", true)); let result = answer( - &f, + &mut f, "sum by (job) (rate(http_requests_total[1m]))", 5 * MINUTE, &all, @@ -175,7 +199,7 @@ fn query_time_plans_read_raw_ranges() { assert!((by_job[1].1 - 1.0).abs() < 1e-9, "{by_job:?}"); let topk = answer( - &f, + &mut f, "topk by (job) (10, sum_over_time(http_requests_total[1m]))", 5 * MINUTE, &all, @@ -194,24 +218,31 @@ 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), + f.query + .raw_ranges("up", 70 * MINUTE, store, Default::default()), Err(Unanswered::NotPlanned(_)) )); assert!(matches!( - f.query.raw_ranges(P99, 70 * MINUTE, store), + f.query + .raw_ranges(P99, 70 * MINUTE, store, Default::default()), 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!(f + .query + .raw_ranges(P99, 70 * MINUTE, store, Default::default()) + .is_ok()); assert!(matches!( - f.query.raw_ranges(P99, 70 * MINUTE + 1, store), + f.query + .raw_ranges(P99, 70 * MINUTE + 1, store, Default::default()), Err(Unanswered::MissingPane(_)) )); // The pane ending at 10 min was skipped as possibly partial. assert!(matches!( - f.query.raw_ranges(P99, 60 * MINUTE, store), + f.query + .raw_ranges(P99, 60 * MINUTE, store, Default::default()), Err(Unanswered::MissingPane(_)) )); } @@ -223,20 +254,21 @@ 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.raw_ranges(P99, time_ms, &f.precompute.store), + f.query + .raw_ranges(P99, time_ms, &f.precompute.store, Default::default()), Err(Unanswered::Unsupported(_)) )); } // The window ending at `i64::MAX` fits; binding it does not panic. assert!(f .query - .raw_ranges(P99, i64::MAX, &f.precompute.store) + .raw_ranges(P99, i64::MAX, &f.precompute.store, Default::default()) .is_ok()); let panes = fixture(common::plan("example4b", false)); assert!(matches!( panes .query - .raw_ranges(P99, i64::MIN, &panes.precompute.store), + .raw_ranges(P99, i64::MIN, &panes.precompute.store, Default::default()), Err(Unanswered::Unsupported(_)) )); } @@ -269,12 +301,12 @@ fn examples_answer_like_the_query_time_reference() { let reference = plan.query_time_reference().unwrap(); let mut installed = fixture(plan.clone()); installed.precompute.ingest(&all); - let by_reference = fixture(reference); + let mut 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); + let expected = answer(&mut by_reference, query, 70 * MINUTE, &all); assert!(!expected.is_empty(), "{example} raw={raw} {query}"); assert_eq!( - answer(&installed, query, 70 * MINUTE, &all), + answer(&mut installed, query, 70 * MINUTE, &all), expected, "{example} raw={raw} {query}" ); @@ -293,14 +325,16 @@ fn raw_ranges_carry_the_scans_label_matchers() { ); let plan = control_plane::plan_workload(&workload, &common::inputs(true)).unwrap(); let f = fixture(plan); - let (_, ranges) = f + let ranges = f .query .raw_ranges( "quantile_over_time(0.99, latency_ms{job=\"a\"}[60m])", 70 * MINUTE, &f.precompute.store, + Default::default(), ) - .unwrap(); + .unwrap() + .ranges; assert_eq!(ranges.len(), 1); assert_eq!( ranges[0].matchers, diff --git a/data_plane/tests/sds.rs b/data_plane/tests/sds.rs index 1e27b416..cdfecf37 100644 --- a/data_plane/tests/sds.rs +++ b/data_plane/tests/sds.rs @@ -251,3 +251,27 @@ fn build_order_refuses_ingestion_reading_query_run() { }) ); } + +/// B3's five query-time KLL summaries kept across evaluations are one +/// query-run definition with a 1-minute coverage, each output reading the +/// summary `i * 1min` before evaluation. Nothing is built at ingestion. +#[test] +fn kept_summaries_share_one_query_run_definition() { + let sds = Sds::from_plan(&plan("example4-b3", true)).unwrap(); + assert!(sds.ingestion_build_order().unwrap().is_empty()); + let mut kept = sds + .definitions + .values() + .filter(|d| d.built_at == StoredOutputBuild::QueryRun); + let definition = kept.next().unwrap(); + assert!(kept.next().is_none()); + assert_eq!(definition.pane_ms, 60_000); + let mut offsets: Vec = sds + .outputs + .values() + .filter(|(id, _)| id == &definition.id) + .map(|(_, binding)| binding.offset_ms) + .collect(); + offsets.sort(); + assert_eq!(offsets, (0..5).map(|i| i * 60_000).collect::>()); +} diff --git a/examples/planner-layering-example4-b3.json b/examples/planner-layering-example4-b3.json new file mode 100644 index 00000000..9ad9481d --- /dev/null +++ b/examples/planner-layering-example4-b3.json @@ -0,0 +1,44 @@ +{ + "query_workload": { + "language": "promql", + "query_batch": null, + "repeating_queries": [ + { + "query": "quantile_over_time(0.99, latency_ms[5m])", + "demand": { + "fixed_interval": 60000 + }, + "requirements": { + "accuracy": { + "accuracy_metric": "rank", + "max_error": 0.01, + "failure_probability": 0.01, + "accuracy_scope": "whole_result" + }, + "max_latency_ms": 200.0 + }, + "predictability": { + "predictable": { + "known_at": null + } + } + } + ] + }, + "data_workload": { + "arrival": { + "ingesting": { + "rows_per_second": 10000.0 + } + }, + "sources": { + "latency_ms": { + "series_count": 10, + "metric_type": null, + "sample_interval": 1, + "stream_norms": null, + "max_distinct_per_evaluation": null + } + } + } +} From 7fbd02a679326347f26057eb271e95e6af44b8c2 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 20 Oct 2026 19:40:49 +0000 Subject: [PATCH 5/6] feat(precompute): event-time barrier for Remote Write producers Ports #814 to the plan-executing engine. POST /api/v1/precompute/watermark?time_ms=T: the producer states that every sample at or before T has been written. The request queues an AdvanceWatermark command behind every write already acknowledged and answers {"barrier_ms": B}, the barrier in force, once it has run: every pane with end <= T is closed, keeping the left-open (end - width, end] convention, so the last pane of a series that stopped sending is built and its window served from the plan without a later sample. - T is capped at the wall clock, like the sample watermark. - A repeated or lower T does nothing; the answer is the barrier in force. - The closing reuses the expired-pane skip and checked arithmetic, so T = i64::MAX terminates. - Input stays open; a sample arriving later for a pane the barrier closed counts as late. - A missing or non-integer time_ms gets 400 bad_data with the reason. - With several producers, the barrier must follow all of their writes through T; the README says so. --- README.md | 12 ++++ data_plane/src/engine.rs | 24 ++++++++ data_plane/src/lib.rs | 49 ++++++++++++++++ data_plane/src/precompute.rs | 48 +++++++++++++--- data_plane/tests/answers.rs | 102 +++++++++++++++++++++++++++++++++ data_plane/tests/precompute.rs | 99 ++++++++++++++++++++++++++++++++ 6 files changed, 326 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 1940572e..569f2cad 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,18 @@ 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. +A producer can close panes without waiting for later samples: +`POST /api/v1/precompute/watermark?time_ms=T` states that every sample at or +before `T` has been written. After the writes acknowledged before it have +been ingested, every pane ending at or before `T` (capped at the wall clock) +is closed, so the last pane of a series that stopped sending is built and +served. A repeated or lower `T` does nothing and input stays open. A sample +that arrives later for a pane the barrier closed (one ending at or before +`T`) counts as late; a later sample in a pane still open is accepted, even if +it is at or before `T`. The response is `{"barrier_ms": B}`, the barrier in +force: `T` capped at the wall clock, or an earlier higher barrier. With several producers, +send the barrier only after all of them have written through `T`. + 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 diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 17aa8653..638d517c 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -53,6 +53,10 @@ enum Command { raw: BTreeMap>, reply: Reply>, }, + AdvanceWatermark { + barrier_ms: i64, + reply: tokio::sync::oneshot::Sender, + }, Stop, } @@ -216,6 +220,21 @@ impl EngineHandle { *self.stats.lock().unwrap_or_else(|e| e.into_inner()) } + /// Queues an event-time barrier behind every write already acknowledged + /// and waits until it has run: every pane ending at or before + /// `barrier_ms` (capped at the wall clock) is then closed. Returns the + /// barrier in force; see [`Precompute::advance_watermark_at`]. + pub async fn advance_watermark(&self, barrier_ms: i64) -> Result { + if !self.is_running() { + return Err(IngestError::Stopped); + } + let (reply, done) = tokio::sync::oneshot::channel(); + self.commands + .send(Command::AdvanceWatermark { barrier_ms, reply }) + .map_err(|_| IngestError::Stopped)?; + done.await.map_err(|_| IngestError::Stopped) + } + /// For tests: stops the engine thread after the commands queued before. /// Samples queued between this call and the stop are acknowledged and /// then lost; ingestion after the stop fails with @@ -312,6 +331,11 @@ impl Engine { self.query.kept_stats; let _ = reply.send(answer); } + Command::AdvanceWatermark { barrier_ms, reply } => { + let barrier_ms = self.precompute.advance_watermark(barrier_ms); + *shared.stats.lock().unwrap_or_else(|e| e.into_inner()) = self.precompute.stats; + let _ = reply.send(barrier_ms); + } Command::Stop => break, } } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 448a547b..5b7b0a43 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -106,6 +106,7 @@ pub fn router(state: Arc) -> Router { .route("/api/v1/query", get(query).post(query)) .route("/api/v1/query_range", get(query_range).post(query_range)) .route("/api/v1/status/asap", get(status)) + .route("/api/v1/precompute/watermark", post(watermark)) .route("/-/ready", get(ready)) .with_state(state) } @@ -159,6 +160,54 @@ async fn write(State(state): State>, body: Bytes) -> Response { } } +/// `POST /api/v1/precompute/watermark?time_ms=T`: the producer states that +/// every sample at or before `T` has been written. Once the engine has run +/// every write acknowledged before this request, every pane with +/// `end <= T` is closed, so the last pane of a series that stopped sending +/// is built without a later sample. `T` is capped at the wall clock; a +/// repeated or lower `T` does nothing. A sample that arrives later for a +/// pane the barrier closed counts as late. Input stays open. With several +/// producers, call it only after all of them have written through `T`. +/// Answers `{"barrier_ms": B}`, the barrier in force (`T` capped at the wall +/// clock, or a higher earlier barrier); 400 `bad_data` with the reason when +/// `time_ms` is missing or not an integer. +async fn watermark(State(state): State>, RawQuery(query): RawQuery) -> Response { + let time_ms = form_urlencoded::parse(query.unwrap_or_default().as_bytes()) + .find(|(name, _)| name == "time_ms") + .map(|(_, value)| value.into_owned()); + let time_ms = match time_ms.as_deref().map(str::parse::) { + Some(Ok(time_ms)) => time_ms, + Some(Err(_)) => { + let reason = format!( + "invalid parameter \"time_ms\": {:?} is not an integer number of milliseconds", + time_ms.unwrap_or_default() + ); + return ( + StatusCode::BAD_REQUEST, + axum::Json(prometheus::bad_data(&reason)), + ) + .into_response(); + } + None => { + let reason = "missing parameter \"time_ms\" (integer milliseconds)"; + return ( + StatusCode::BAD_REQUEST, + axum::Json(prometheus::bad_data(reason)), + ) + .into_response(); + } + }; + match state.engine.advance_watermark(time_ms).await { + Ok(barrier_ms) => { + axum::Json(serde_json::json!({ "barrier_ms": barrier_ms })).into_response() + } + Err(error) => { + tracing::error!(%error, "event-time barrier refused"); + (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()).into_response() + } + } +} + /// Ready while the engine thread runs; a stopped engine never ingests /// again, so the data plane must be restarted. async fn ready(State(state): State>) -> StatusCode { diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index 3775fdae..1eba1f2d 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -208,6 +208,9 @@ pub struct Precompute { /// passes it, so a series stamped ahead of the others cannot close /// panes early. newest_clock_ms: Option, + /// The newest event-time barrier a producer declared: every pane ending + /// at or before it is closed. + barrier_ms: Option, pub store: SummaryStore, pub stats: PrecomputeStats, } @@ -260,6 +263,7 @@ impl Precompute { max_future_ms: millis(options.max_future_ms), newest_sample_ms: None, newest_clock_ms: None, + barrier_ms: None, store: SummaryStore::new(), stats: PrecomputeStats::default(), }) @@ -280,21 +284,42 @@ impl Precompute { /// The watermark: no sample older than this is expected any more. It is /// the newest routed sample, capped at the wall clock, minus the allowed - /// lateness. + /// lateness, or the newest barrier if that is later. pub fn watermark_ms(&self) -> Option { let newest = match self.newest_clock_ms { - Some(clock) => self.newest_sample_ms?.min(clock), - None => self.newest_sample_ms?, + Some(clock) => self.newest_sample_ms.map(|newest| newest.min(clock)), + None => self.newest_sample_ms, }; - Some(newest.saturating_sub(self.allowed_lateness_ms)) + let from_samples = newest.map(|newest| newest.saturating_sub(self.allowed_lateness_ms)); + from_samples.max(self.barrier_ms) + } + + /// [`Self::advance_watermark_at`] the current wall-clock time. + pub fn advance_watermark(&mut self, barrier_ms: i64) -> i64 { + self.advance_watermark_at(barrier_ms, wall_clock_ms()) + } + + /// An event-time barrier: the producer states that every sample at or + /// before `barrier_ms` has been written. Closes every pane with + /// `end <= barrier_ms` (panes are `(end - width, end]`), so a pane whose + /// series stopped sending closes without a later sample. The barrier is + /// capped at `now_ms`; a repeated or lower barrier does nothing. A sample + /// that arrives later for a pane the barrier closed counts as late. + /// Returns the barrier in force. + pub fn advance_watermark_at(&mut self, barrier_ms: i64, now_ms: i64) -> i64 { + self.newest_clock_ms = self.newest_clock_ms.max(Some(now_ms)); + let barrier_ms = barrier_ms.min(now_ms); + if self.barrier_ms.is_some_and(|current| current >= barrier_ms) { + return self.barrier_ms.unwrap_or(barrier_ms); + } + self.barrier_ms = Some(barrier_ms); + self.close_panes(); + barrier_ms } /// [`Self::ingest_at`] the current wall-clock time. pub fn ingest(&mut self, samples: &[RawSample]) { - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX)); - self.ingest_at(samples, now_ms); + self.ingest_at(samples, wall_clock_ms()); } /// Buffers `samples` into their panes, then builds every pane the @@ -421,6 +446,13 @@ impl Precompute { } } +/// Milliseconds since the Unix epoch. +fn wall_clock_ms() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX)) +} + /// End of the tumbling pane of width `pane_ms` that contains `timestamp_ms`, /// or `None` when it does not fit in an `i64`. Panes are /// `(end - pane_ms, end]`, like PromQL ranges: a sample at a pane boundary diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs index 62232375..c6ac3a27 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -391,3 +391,105 @@ async fn rejects_evaluation_times_out_of_range() { } assert!(state.engine.is_running()); } + +/// After a producer's event-time barrier, the last pane of a series that +/// stopped sending is closed, and the planned query is served from the +/// plan instead of forwarded. Input stays open. +#[tokio::test] +async fn barrier_serves_the_window_of_a_stopped_series() { + let state = Arc::new( + DataPlane::new( + common::plan_json("example4b", false), + fake_prometheus(vec![]).await, + 15_000, + ) + .unwrap(), + ); + let app = router(state.clone()); + let write = |from: i64, to: i64| { + let request = WriteRequest { + timeseries: vec![TimeSeries { + labels: vec![ + Label { + name: "__name__".into(), + value: "latency_ms".into(), + }, + Label { + name: "job".into(), + value: "a".into(), + }, + ], + samples: (from..to) + .step_by(15_000) + .map(|t| Sample { + timestamp: t, + value: (t / 15_000 % 240) as f64, + }) + .collect(), + }], + }; + Request::post("/api/v1/write") + .body(Body::from(encode(&request))) + .unwrap() + }; + let written = app + .clone() + .oneshot(write(MINUTE, 70 * MINUTE)) + .await + .unwrap(); + assert_eq!(written.status(), StatusCode::NO_CONTENT); + let query = "quantile_over_time(0.99, latency_ms[60m])"; + // The series stopped at 69:45; with 15 s lateness (60, 70 min] is open. + let response = instant(&app, query, 70 * MINUTE).await; + assert_eq!(response.execution.as_deref(), Some("forwarded")); + + let barrier = app + .clone() + .oneshot( + Request::post(format!( + "/api/v1/precompute/watermark?time_ms={}", + 70 * MINUTE + )) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(barrier.status(), StatusCode::OK); + let body = barrier.into_body().collect().await.unwrap().to_bytes(); + let body: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["barrier_ms"], 70 * MINUTE, "{body}"); + let response = instant(&app, query, 70 * MINUTE).await; + assert_eq!( + response.execution.as_deref(), + Some("plan"), + "{:?}", + response.body + ); + + let written = app + .clone() + .oneshot(write(70 * MINUTE + 15_000, 72 * MINUTE)) + .await + .unwrap(); + assert_eq!(written.status(), StatusCode::NO_CONTENT, "input stays open"); + for uri in [ + "/api/v1/precompute/watermark", + "/api/v1/precompute/watermark?time_ms=soon", + "/api/v1/precompute/watermark?time_ms=1.5", + ] { + let bad = app + .clone() + .oneshot(Request::post(uri).body(Body::empty()).unwrap()) + .await + .unwrap(); + assert_eq!(bad.status(), StatusCode::BAD_REQUEST, "{uri}"); + let body = bad.into_body().collect().await.unwrap().to_bytes(); + let body: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["errorType"], "bad_data", "{uri}"); + assert!( + body["error"].as_str().unwrap().contains("time_ms"), + "{uri}: {body}" + ); + } +} diff --git a/data_plane/tests/precompute.rs b/data_plane/tests/precompute.rs index c661e181..ce848e79 100644 --- a/data_plane/tests/precompute.rs +++ b/data_plane/tests/precompute.rs @@ -501,3 +501,102 @@ fn query_run_metrics_are_bounded_like_routed_samples() { assert_eq!(engine.stats.samples_query_run_only, 1); assert_eq!(engine.watermark_ms(), Some(now)); } + +/// An event-time barrier closes the last pane of a series that stopped +/// sending, without a later sample. +#[test] +fn barrier_closes_the_last_pane_of_a_stopped_series() { + let (mut engine, sds) = precompute(common::plan("example4b", false)); + let id = sds.definitions.keys().next().unwrap().clone(); + engine.ingest_at( + &samples("latency_ms", &["a"], MINUTE, 25 * MINUTE), + 31 * MINUTE, + ); + assert!(engine + .store + .get(&id, Selection::ending_at(2 * PANE)) + .is_some()); + assert!( + engine + .store + .get(&id, Selection::ending_at(3 * PANE)) + .is_none(), + "(20, 30 min] open" + ); + assert_eq!(engine.advance_watermark_at(3 * PANE, 31 * MINUTE), 3 * PANE); + assert_eq!( + engine + .store + .get(&id, Selection::ending_at(3 * PANE)) + .unwrap() + .num_rows(), + 1 + ); +} + +/// A repeated or lower barrier does nothing, and a barrier ahead of the +/// wall clock is capped at it. +#[test] +fn barrier_is_monotonic_and_capped_at_the_wall_clock() { + let (mut engine, sds) = precompute(common::plan("example4b", false)); + let id = sds.definitions.keys().next().unwrap().clone(); + engine.ingest_at( + &samples("latency_ms", &["a"], MINUTE, 25 * MINUTE), + 31 * MINUTE, + ); + assert_eq!( + engine.advance_watermark_at(4 * PANE, 31 * MINUTE), + 31 * MINUTE + ); + assert!(engine + .store + .get(&id, Selection::ending_at(3 * PANE)) + .is_some()); + assert!( + engine + .store + .get(&id, Selection::ending_at(4 * PANE)) + .is_none(), + "capped at 31 min" + ); + let stats = engine.stats; + assert_eq!( + engine.advance_watermark_at(31 * MINUTE, 31 * MINUTE), + 31 * MINUTE + ); + assert_eq!( + engine.advance_watermark_at(2 * PANE, 31 * MINUTE), + 31 * MINUTE + ); + assert_eq!(engine.stats, stats); +} + +/// A barrier at `i64::MAX` terminates and builds only retained panes. +#[test] +fn barrier_at_i64_max_terminates() { + let (mut engine, _) = precompute(common::plan("example4b", false)); + engine.ingest_at( + &samples("latency_ms", &["a"], MINUTE, 25 * MINUTE), + i64::MAX, + ); + let built = engine.stats.panes_built; + assert_eq!(engine.advance_watermark_at(i64::MAX, i64::MAX), i64::MAX); + assert!(engine.stats.panes_built - built <= 7); +} + +/// Input stays open after a barrier: a later sample is accepted, and one at +/// or before the barrier counts as late. +#[test] +fn input_stays_open_after_a_barrier() { + let (mut engine, _) = precompute(common::plan("example4b", false)); + engine.ingest_at( + &samples("latency_ms", &["a"], MINUTE, 25 * MINUTE), + 31 * MINUTE, + ); + engine.advance_watermark_at(3 * PANE, 31 * MINUTE); + let accepted = engine.stats.samples_accepted; + engine.ingest_at(&[sample("latency_ms", "a", 35 * MINUTE, 1.0)], 36 * MINUTE); + assert_eq!(engine.stats.samples_accepted, accepted + 1); + engine.ingest_at(&[sample("latency_ms", "a", 29 * MINUTE, 1.0)], 36 * MINUTE); + assert_eq!(engine.stats.samples_dropped_late, 1); +} From a70c2382713b13bcf551c16d3ef06adf211456a2 Mon Sep 17 00:00:00 2001 From: zzylol Date: Tue, 20 Oct 2026 23:50:03 +0000 Subject: [PATCH 6/6] feat(data-plane): sharded ingestion workers, a query thread pool and a shared SDS Series are sharded by label set across ingestion workers that each own their shard's open panes. The global watermark is the minimum over the workers that have seen a sample; the worker or barrier that moves it closes the passed panes from every worker's samples. The event-time barrier closes every worker's panes. Planned queries run on a pool of query threads, one query per thread, over a plan compiled once and shared through Arc, and wait for the writes acknowledged before them. The SDS store holds Arc instances under a RwLock and never rewrites one. jemalloc is the data plane's allocator. --- Cargo.lock | 74 +++-- README.md | 7 + data_plane/Cargo.toml | 1 + data_plane/src/engine.rs | 374 +++++++++++++--------- data_plane/src/lib.rs | 9 +- data_plane/src/main.rs | 11 + data_plane/src/precompute.rs | 566 ++++++++++++++++++++++++--------- data_plane/src/query.rs | 66 ++-- data_plane/src/sds.rs | 56 +++- data_plane/tests/answers.rs | 103 ++++++ data_plane/tests/kept.rs | 16 +- data_plane/tests/precompute.rs | 110 ++++++- data_plane/tests/sds.rs | 68 +++- 13 files changed, 1071 insertions(+), 390 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5920ce46..a218ee69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1053,6 +1053,7 @@ dependencies = [ "sha2 0.10.9", "snap", "thiserror 1.0.69", + "tikv-jemallocator", "tokio", "tower", "tracing", @@ -1098,7 +1099,7 @@ dependencies = [ "flate2", "futures", "indexmap", - "itertools 0.14.0", + "itertools", "liblzma", "log", "object_store", @@ -1130,7 +1131,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", "parking_lot", @@ -1155,7 +1156,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", ] @@ -1174,7 +1175,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "libc", "log", "object_store", @@ -1221,7 +1222,7 @@ dependencies = [ "flate2", "futures", "glob", - "itertools 0.14.0", + "itertools", "liblzma", "log", "object_store", @@ -1252,7 +1253,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "object_store", "tokio", ] @@ -1326,7 +1327,7 @@ dependencies = [ "datafusion-pruning", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", "parking_lot", @@ -1379,7 +1380,7 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr-common", "indexmap", - "itertools 0.14.0", + "itertools", "recursive", "serde_json", "sqlparser", @@ -1394,7 +1395,7 @@ dependencies = [ "arrow", "datafusion-common", "indexmap", - "itertools 0.14.0", + "itertools", ] [[package]] @@ -1418,7 +1419,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hex", - "itertools 0.14.0", + "itertools", "log", "md-5", "memchr", @@ -1481,7 +1482,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hashbrown 0.17.1", - "itertools 0.14.0", + "itertools", "itoa", "log", "memchr", @@ -1554,7 +1555,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools 0.14.0", + "itertools", "log", "recursive", "regex", @@ -1576,7 +1577,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "parking_lot", "petgraph", "recursive", @@ -1595,7 +1596,7 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools 0.14.0", + "itertools", ] [[package]] @@ -1610,7 +1611,7 @@ dependencies = [ "datafusion-expr-common", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "parking_lot", "pin-project", ] @@ -1630,7 +1631,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", - "itertools 0.14.0", + "itertools", "recursive", ] @@ -1659,7 +1660,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "log", "num-traits", "parking_lot", @@ -2340,15 +2341,6 @@ version = "1.70.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" -[[package]] -name = "itertools" -version = "0.10.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" -dependencies = [ - "either", -] - [[package]] name = "itertools" version = "0.14.0" @@ -2789,7 +2781,7 @@ dependencies = [ "futures-util", "http", "humantime", - "itertools 0.14.0", + "itertools", "parking_lot", "percent-encoding", "thiserror 2.0.20", @@ -3070,7 +3062,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools", "proc-macro2", "quote", "syn 2.0.119", @@ -3387,7 +3379,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3787,7 +3779,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3850,6 +3842,26 @@ dependencies = [ "ordered-float", ] +[[package]] +name = "tikv-jemalloc-sys" +version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd8aa5b2ab86a2cefa406d889139c162cbb230092f7d1d7cbc1716405d852a3b" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.55" @@ -4334,7 +4346,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/README.md b/README.md index 569f2cad..a7d4233c 100644 --- a/README.md +++ b/README.md @@ -53,6 +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. +Series are sharded by their label set across `--ingestion-workers` (default +one per core); with several workers, a pane closes once the newest sample of +every worker that has received one has passed it. Planned queries run on +`--query-threads` (default one per core), one query per thread, and each +sees the panes of every write acknowledged before it. The data plane uses +jemalloc. + A producer can close panes without waiting for later samples: `POST /api/v1/precompute/watermark?time_ms=T` states that every sample at or before `T` has been written. After the writes acknowledged before it have diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 8ce0da62..acb25e0a 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -7,6 +7,7 @@ edition.workspace = true asap-planner.workspace = true asap-executor.workspace = true sha2 = "0.10" +tikv-jemallocator = "0.6" serde.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 638d517c..d035da37 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -1,6 +1,19 @@ -//! The engine thread: owns the installed plan and everything that executes -//! it. Planner IR holds `Rc`s, so the plan cannot be shared across the HTTP -//! server's threads; handlers send the engine commands instead. +//! The data plane's threads: ingestion workers that build stored summaries, +//! and a pool of query threads that answer planned queries. +//! +//! The plan is compiled once, when it is installed: its ingestion side +//! ([`Ingestion`]) and its query roots ([`QueryEngine`]). Every thread +//! shares them, with the stored summaries, the kept query-run summaries and +//! the watermark ([`Shared`]). +//! +//! - **Ingestion.** Series are sharded by a hash of their label set; each +//! worker owns the open panes of its shard. The global watermark is the +//! minimum over the workers; the worker or barrier that moves it closes +//! the panes it passed, from every worker's samples. +//! - **Queries.** One query per thread, from a pool: a query runs on the +//! thread that takes it, while other queries run on the others and +//! ingestion continues. Before reading, a query waits until the workers +//! have run every write acknowledged before it, so it sees those writes. use std::collections::BTreeMap; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -11,36 +24,52 @@ use asap_planner::InstallablePlan; use tokio::sync::oneshot; -use crate::precompute::{Precompute, PrecomputeOptions, PrecomputeStats, RawSample}; +use crate::precompute::{ + shard_of, wall_clock_ms, Ingestion, PrecomputeOptions, PrecomputeStats, RawSample, Shared, +}; use crate::query::{ IngestionProgress, KeptResultStats, Prepared, QueryEngine, ResultSample, Unanswered, }; use crate::sds::Sds; -/// How the engine thread ingests samples. +/// How the engine runs: its threads and how it ingests samples. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct EngineOptions { /// Lateness, future bound and pane checks of the precompute engine. pub precompute: PrecomputeOptions, - /// Samples queued for the engine, not yet ingested, beyond which Remote - /// Write is refused with 503 so senders retry. + /// Samples queued for the workers, not yet ingested, beyond which + /// Remote Write is refused with 503 so senders retry. pub max_queued_samples: u64, + /// Ingestion workers; series are sharded across them. + pub ingestion_workers: usize, + /// Query threads; each runs one query at a time. + pub query_threads: usize, } impl Default for EngineOptions { + /// One thread of each kind per available core. fn default() -> Self { + let cores = std::thread::available_parallelism().map_or(1, |n| n.get()); Self { precompute: PrecomputeOptions::default(), // Two Remote Write requests of the default sample limit. max_queued_samples: 2_000_000, + ingestion_workers: cores, + query_threads: cores, } } } type Reply = oneshot::Sender>; -enum Command { +enum WorkerCommand { Ingest(Vec), + /// Reply once every command queued before has run. + Flush(oneshot::Sender<()>), + Stop, +} + +enum QueryJob { RawRanges { query: String, time_ms: i64, @@ -53,11 +82,6 @@ enum Command { raw: BTreeMap>, reply: Reply>, }, - AdvanceWatermark { - barrier_ms: i64, - reply: tokio::sync::oneshot::Sender, - }, - Stop, } /// Why the engine did not take samples. @@ -66,27 +90,26 @@ pub enum IngestError { /// More samples are queued than `max_queued_samples`; retry later. #[error("the engine is busy: too many samples queued")] Busy, - /// The engine thread has stopped; no sample will be ingested again. + /// An engine thread has stopped; no sample will be ingested again. #[error("the engine has stopped")] Stopped, } -/// The HTTP side of the engine thread. +/// The HTTP side of the engine's threads. #[derive(Clone)] pub struct EngineHandle { - commands: Sender, + workers: Vec>, + queries: Sender, + precompute: Arc, + query: Arc, /// Samples sent and not yet ingested. queued_samples: Arc, max_queued_samples: u64, - /// False once the engine thread has exited, for any reason. + /// False once any engine thread has exited, for any reason. running: Arc, - /// The counters as of the last command the engine ran. - stats: Arc>, - /// The query-run summary counters as of the last evaluation. - kept_stats: Arc>, } -/// Marks the engine stopped when its thread exits, panics included. +/// Marks the engine stopped when one of its threads exits, panics included. struct RunningGuard(Arc); impl Drop for RunningGuard { @@ -95,51 +118,82 @@ impl Drop for RunningGuard { } } +/// Reads and validates the plan serialized in `plan_json`, and its stored +/// summaries. +fn install(plan_json: &str) -> anyhow::Result<(InstallablePlan, Sds)> { + let plan: InstallablePlan = serde_json::from_str(plan_json)?; + plan.validate()?; + let sds = Sds::from_plan(&plan)?; + tracing::info!( + definitions = sds.definitions.len(), + outputs = sds.outputs.len(), + "stored summaries" + ); + Ok((plan, sds)) +} + impl EngineHandle { - /// Starts the engine for the plan serialized in `plan_json`, parsed once - /// on the engine thread, and returns the plan's queries. Fails if the - /// plan is invalid or the data plane cannot execute it. + /// Starts the engine's threads for the plan serialized in `plan_json`, + /// compiled once, and returns the plan's queries. Fails if the plan is + /// invalid or the data plane cannot execute it. pub fn spawn(plan_json: String, options: EngineOptions) -> anyhow::Result<(Self, Vec)> { - let (commands, receiver) = channel(); - let (ready, started) = std::sync::mpsc::channel(); + let (plan, sds) = install(&plan_json)?; + let queries = plan.queries.iter().map(|q| q.query.clone()).collect(); + let ingestion = Ingestion::new(&plan.dag.dag, &sds, options.precompute)?; + let precompute = Arc::new(Shared::new(ingestion, options.ingestion_workers)); + let engine = Arc::new(QueryEngine::new(&plan, &sds)); + tracing::info!( + ingestion_workers = precompute.workers(), + query_threads = options.query_threads.max(1), + "engine threads" + ); let running = Arc::new(AtomicBool::new(true)); let queued_samples = Arc::new(AtomicU64::new(0)); - let stats = Arc::new(Mutex::new(PrecomputeStats::default())); - let kept_stats = Arc::new(Mutex::new(KeptResultStats::default())); - let shared = Shared { - queued_samples: queued_samples.clone(), - stats: stats.clone(), - kept_stats: kept_stats.clone(), - }; - let guard = RunningGuard(running.clone()); - std::thread::Builder::new() - .name("engine".into()) - .spawn(move || { - let _guard = guard; - match Engine::new(&plan_json, options.precompute) { - Ok((engine, queries)) => { - let _ = ready.send(Ok(queries)); - engine.serve(receiver, shared); - } - Err(error) => { - let _ = ready.send(Err(error)); - } - } - })?; - let queries = started.recv()??; + let mut workers = Vec::new(); + for index in 0..precompute.workers() { + let (sender, commands) = channel(); + let (precompute, queued_samples) = (precompute.clone(), queued_samples.clone()); + let guard = RunningGuard(running.clone()); + std::thread::Builder::new() + .name(format!("ingest-{index}")) + .spawn(move || { + let _guard = guard; + serve_worker(index, &precompute, &queued_samples, commands); + })?; + workers.push(sender); + } + let (queries_sender, jobs) = channel(); + let jobs = Arc::new(Mutex::new(jobs)); + for index in 0..options.query_threads.max(1) { + let (engine, precompute, jobs) = (engine.clone(), precompute.clone(), jobs.clone()); + let guard = RunningGuard(running.clone()); + std::thread::Builder::new() + .name(format!("query-{index}")) + .spawn(move || { + let _guard = guard; + serve_queries(&engine, &precompute, &jobs); + })?; + } let handle = Self { - commands, + workers, + queries: queries_sender, + precompute, + query: engine, queued_samples, max_queued_samples: options.max_queued_samples, running, - stats, - kept_stats, }; Ok((handle, queries)) } - /// Queues `samples` for precompute. `Busy` when the queue would hold - /// more than `max_queued_samples`; `Stopped` when the engine has exited. + /// How many ingestion workers series are sharded across. + pub fn ingestion_workers(&self) -> usize { + self.workers.len() + } + + /// Queues each worker's shard of `samples`, all or none. `Busy` when + /// the queue would hold more than `max_queued_samples`; `Stopped` when + /// an engine thread has exited. pub fn try_ingest(&self, samples: Vec) -> Result<(), IngestError> { if !self.is_running() { return Err(IngestError::Stopped); @@ -158,16 +212,31 @@ impl EngineHandle { Err(current) => queued = current, } } - self.commands.send(Command::Ingest(samples)).map_err(|_| { - self.queued_samples.fetch_sub(count, Ordering::SeqCst); - IngestError::Stopped - }) + let mut shards = vec![Vec::new(); self.workers.len()]; + for sample in samples { + shards[shard_of(&sample.labels, self.workers.len())].push(sample); + } + // The channels are unbounded: once admitted, a request reaches + // every worker, so a retry never doubles samples. + for (worker, shard) in self.workers.iter().zip(shards) { + if shard.is_empty() { + continue; + } + let n = shard.len() as u64; + if worker.send(WorkerCommand::Ingest(shard)).is_err() { + self.queued_samples.fetch_sub(n, Ordering::SeqCst); + return Err(IngestError::Stopped); + } + } + Ok(()) } /// The plan root answering `query` at `time_ms` and the raw samples it - /// needs; see [`QueryEngine::raw_ranges`]. + /// needs, once the writes acknowledged before are ingested; see + /// [`QueryEngine::raw_ranges`]. pub async fn raw_ranges(&self, query: String, time_ms: i64) -> Result { - self.ask(|reply| Command::RawRanges { + self.sync().await; + self.ask(|reply| QueryJob::RawRanges { query, time_ms, reply, @@ -183,7 +252,7 @@ impl EngineHandle { time_ms: i64, raw: BTreeMap>, ) -> Result, Unanswered> { - self.ask(|reply| Command::Evaluate { + self.ask(|reply| QueryJob::Evaluate { root: prepared.root, time_ms, progress: prepared.progress, @@ -193,19 +262,36 @@ impl EngineHandle { .await } - /// The query-run summary counters as of the last evaluation. + /// The query-run summary counters as of now. pub fn kept_stats(&self) -> KeptResultStats { - *self.kept_stats.lock().unwrap_or_else(|e| e.into_inner()) + self.query.kept_stats() } - async fn ask(&self, command: impl FnOnce(Reply) -> Command) -> Result { + async fn ask(&self, job: impl FnOnce(Reply) -> QueryJob) -> Result { let stopped = || Unanswered::Execution("the engine has stopped".into()); let (reply, answer) = oneshot::channel(); - self.commands.send(command(reply)).map_err(|_| stopped())?; + self.queries.send(job(reply)).map_err(|_| stopped())?; answer.await.map_err(|_| stopped())? } - /// Whether the engine thread is still running. + /// Returns once every worker has run every command queued before, so + /// a query sees the writes acknowledged before it. A worker that takes + /// samples closes the panes they let the watermark pass before it runs + /// its next command, so one round suffices. + pub async fn sync(&self) { + let mut flushed = Vec::new(); + for worker in &self.workers { + let (reply, done) = oneshot::channel(); + if worker.send(WorkerCommand::Flush(reply)).is_ok() { + flushed.push(done); + } + } + for done in flushed { + let _ = done.await; + } + } + + /// Whether every engine thread is still running. pub fn is_running(&self) -> bool { self.running.load(Ordering::SeqCst) } @@ -215,33 +301,37 @@ impl EngineHandle { self.queued_samples.load(Ordering::SeqCst) } - /// The precompute counters as of the last batch the engine ingested. + /// The precompute counters as of now. pub fn stats(&self) -> PrecomputeStats { - *self.stats.lock().unwrap_or_else(|e| e.into_inner()) + self.precompute.stats() } - /// Queues an event-time barrier behind every write already acknowledged - /// and waits until it has run: every pane ending at or before + /// Runs an event-time barrier behind every write already acknowledged, + /// on every worker's shard: every pane ending at or before /// `barrier_ms` (capped at the wall clock) is then closed. Returns the - /// barrier in force; see [`Precompute::advance_watermark_at`]. + /// barrier in force; see [`Shared::advance_watermark_at`]. pub async fn advance_watermark(&self, barrier_ms: i64) -> Result { if !self.is_running() { return Err(IngestError::Stopped); } - let (reply, done) = tokio::sync::oneshot::channel(); - self.commands - .send(Command::AdvanceWatermark { barrier_ms, reply }) - .map_err(|_| IngestError::Stopped)?; - done.await.map_err(|_| IngestError::Stopped) + self.sync().await; + let precompute = self.precompute.clone(); + tokio::task::spawn_blocking(move || { + precompute.advance_watermark_at(barrier_ms, wall_clock_ms()) + }) + .await + .map_err(|_| IngestError::Stopped) } - /// For tests: stops the engine thread after the commands queued before. - /// Samples queued between this call and the stop are acknowledged and - /// then lost; ingestion after the stop fails with + /// For tests: stops the ingestion workers after the commands queued + /// before. Samples queued between this call and the stop are + /// acknowledged and then lost; ingestion after the stop fails with /// [`IngestError::Stopped`]. #[doc(hidden)] pub fn stop(&self) { - let _ = self.commands.send(Command::Stop); + for worker in &self.workers { + let _ = worker.send(WorkerCommand::Stop); + } } } @@ -258,88 +348,56 @@ fn admit(queued: u64, count: u64, max_queued: u64) -> Result { } } -/// What the engine thread publishes to the HTTP side. -struct Shared { - queued_samples: Arc, - stats: Arc>, - kept_stats: Arc>, -} - -struct Engine { - precompute: Precompute, - query: QueryEngine, +/// Runs worker `index`'s commands until it is stopped or the handle is +/// dropped. +fn serve_worker( + index: usize, + precompute: &Shared, + queued_samples: &AtomicU64, + commands: Receiver, +) { + for command in commands { + match command { + WorkerCommand::Ingest(samples) => { + precompute.ingest_at(index, &samples, wall_clock_ms()); + queued_samples.fetch_sub(samples.len() as u64, Ordering::SeqCst); + } + WorkerCommand::Flush(reply) => { + let _ = reply.send(()); + } + WorkerCommand::Stop => break, + } + } + tracing::warn!(worker = index, "ingestion worker stopped"); } -impl Engine { - fn new(plan_json: &str, options: PrecomputeOptions) -> anyhow::Result<(Self, Vec)> { - let plan: InstallablePlan = serde_json::from_str(plan_json)?; - plan.validate()?; - let sds = Sds::from_plan(&plan)?; - tracing::info!( - definitions = sds.definitions.len(), - outputs = sds.outputs.len(), - "stored summaries" - ); - 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, +/// Answers jobs from the pool's queue, one at a time, until the queue closes. +fn serve_queries(engine: &QueryEngine, precompute: &Shared, jobs: &Mutex>) { + loop { + let job = match jobs.lock().unwrap_or_else(|e| e.into_inner()).recv() { + Ok(job) => job, + Err(_) => return, }; - Ok((engine, queries)) - } - - fn serve(mut self, commands: Receiver, shared: Shared) { - for command in commands { - match command { - Command::Ingest(samples) => { - self.precompute.ingest(&samples); - *shared.stats.lock().unwrap_or_else(|e| e.into_inner()) = self.precompute.stats; - shared - .queued_samples - .fetch_sub(samples.len() as u64, Ordering::SeqCst); - } - Command::RawRanges { - query, - time_ms, - reply, - } => { - let _ = reply.send(self.query.raw_ranges( - &query, - time_ms, - &self.precompute.store, - self.precompute.progress(), - )); - } - Command::Evaluate { - root, - time_ms, - progress, - raw, - reply, - } => { - let answer = self.query.evaluate( - root, - time_ms, - &mut self.precompute.store, - &progress, - &raw, - ); - // Published before the reply, so a caller that awaited - // the evaluation reads its counters. - *shared.kept_stats.lock().unwrap_or_else(|e| e.into_inner()) = - self.query.kept_stats; - let _ = reply.send(answer); - } - Command::AdvanceWatermark { barrier_ms, reply } => { - let barrier_ms = self.precompute.advance_watermark(barrier_ms); - *shared.stats.lock().unwrap_or_else(|e| e.into_inner()) = self.precompute.stats; - let _ = reply.send(barrier_ms); - } - Command::Stop => break, + let store = &precompute.store; + match job { + QueryJob::RawRanges { + query, + time_ms, + reply, + } => { + let progress = precompute.progress(); + let _ = reply.send(engine.raw_ranges(&query, time_ms, store, progress)); + } + QueryJob::Evaluate { + root, + time_ms, + progress, + raw, + reply, + } => { + let _ = reply.send(engine.evaluate(root, time_ms, store, &progress, &raw)); } } - tracing::warn!("engine stopped"); } } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 5b7b0a43..379fc313 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -1,10 +1,11 @@ //! Backend data plane: receives Prometheus Remote Write samples and serves //! 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 and answers -//! the plan's instant queries from them; every other query, and any planned -//! query it cannot answer, is forwarded to Prometheus. +//! The installed plan is ASAPPlanner's `InstallablePlan`. Ingestion +//! workers build its ingestion-time outputs pane by pane, and a pool of +//! query threads answers the plan's instant queries from them ([`engine`]); +//! every other query, and any planned query it cannot answer, is forwarded +//! to Prometheus. pub mod engine; pub mod precompute; diff --git a/data_plane/src/main.rs b/data_plane/src/main.rs index 88186e41..96297832 100644 --- a/data_plane/src/main.rs +++ b/data_plane/src/main.rs @@ -7,6 +7,9 @@ use data_plane::engine::EngineOptions; use data_plane::precompute::PrecomputeOptions; use data_plane::{router, DataPlane, DEFAULT_FORWARD_TIMEOUT}; +#[global_allocator] +static ALLOCATOR: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + /// Serve Remote Write ingestion and the Prometheus query API for a plan. #[derive(Parser)] struct Args { @@ -30,6 +33,12 @@ struct Args { /// Samples queued for the engine beyond which Remote Write gets 503. #[arg(long, default_value_t = EngineOptions::default().max_queued_samples)] max_queued_samples: u64, + /// Ingestion workers series are sharded across; default one per core. + #[arg(long, default_value_t = EngineOptions::default().ingestion_workers)] + ingestion_workers: usize, + /// Query threads, each running one query at a time; default one per core. + #[arg(long, default_value_t = EngineOptions::default().query_threads)] + query_threads: usize, } #[tokio::main] @@ -49,6 +58,8 @@ async fn main() -> anyhow::Result<()> { max_future_ms: args.max_future_ms, }, max_queued_samples: args.max_queued_samples, + ingestion_workers: args.ingestion_workers, + query_threads: args.query_threads, }, )? .with_forward_timeout(Duration::from_secs(args.forward_timeout_seconds)), diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index 1eba1f2d..d2112ad4 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -1,24 +1,28 @@ //! Precompute engine: builds each stored summary pane once, when the //! watermark passes its end. //! -//! For every [`SummaryDefinition`] the engine buffers the samples of its -//! `Scan`'s metric by tumbling pane. A pane `(end - pane, end]` closes when -//! the watermark (the newest sample time minus the allowed lateness) reaches -//! `end`. Closing it selects the pane's rows from the buffered samples with -//! the executor's `RawSampleSelector` (the `Scan`'s label matchers), then runs -//! the definition's build from its `TimeRange` input up to its output, under -//! ingestion scope. The output is stored and never changes; samples older -//! than a closed pane are dropped. The first pane a definition sees may have -//! started before the data plane did, so it is never built: queries that -//! need it are not answered from stored panes. +//! The plan's ingestion side is compiled once ([`Ingestion`]) and shared by +//! the ingestion workers ([`Shared`]), each of which owns a shard of the +//! series. For every [`SummaryDefinition`] a worker buffers its shard's +//! samples of the definition's `Scan` metric by tumbling pane. A pane +//! `(end - pane, end]` closes when the global watermark (the minimum over +//! the workers of their newest sample, minus the allowed lateness) reaches +//! `end`. Closing it takes every worker's samples of the pane, selects the +//! pane's rows with the executor's `RawSampleSelector` (the `Scan`'s label +//! matchers), then runs the definition's build from its `TimeRange` input +//! up to its output, under ingestion scope. The output is stored and never +//! changes; samples older than a closed pane are dropped. The first pane a +//! definition sees may have started before the data plane did, so it is +//! never built: queries that need it are not answered from stored panes. //! //! Only samples of a metric some definition reads move the watermark, and a //! sample further ahead of the wall clock than `max_future_ms` is dropped, so //! one bad timestamp cannot close every pane up to it. When the watermark //! jumps, panes that would already be evicted are skipped, not built. -use std::collections::BTreeMap; -use std::sync::Arc; +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use asap_executor::operators::Operator; use asap_executor::physical_planner::{compile, CompiledPhysicalDAG, InputContract, Source}; @@ -177,15 +181,19 @@ impl SummaryBuild { } } -struct Builder { - summary: SummaryBuild, - /// Samples of open panes, by pane end. - open: BTreeMap>, - /// End of the first pane seen, which is skipped. `None` until then. - first_pane_end: Option, - /// End of the newest pane built or skipped; samples in it or before it - /// are dropped. `None` until the first sample. - built_through: Option, +/// No value yet in an `AtomicI64` slot. +const UNSET: i64 = i64::MIN; + +fn get(slot: &AtomicI64) -> Option { + Some(slot.load(Ordering::Acquire)).filter(|&v| v != UNSET) +} + +/// The slot's value, set to `value` first if it had none. +fn get_or_init(slot: &AtomicI64, value: i64) -> i64 { + match slot.compare_exchange(UNSET, value, Ordering::AcqRel, Ordering::Acquire) { + Ok(_) => value, + Err(current) => current, + } } /// What one definition did with one sample. @@ -196,57 +204,44 @@ enum Outcome { Buffered, } -pub struct Precompute { - builders: Vec, - /// Metrics read by query-run summaries, with each one's newest routed - /// sample time. - query_run_metrics: BTreeMap>, +/// The plan's ingestion side, compiled once and shared by every ingestion +/// worker: one build per definition built in an ingestion run, in +/// dependency order. +pub struct Ingestion { + builds: Vec, + /// Metrics read by query-run summaries. + query_run_metrics: BTreeSet, allowed_lateness_ms: i64, max_future_ms: i64, - newest_sample_ms: Option, - /// The newest wall-clock time an ingest ran at: the watermark never - /// passes it, so a series stamped ahead of the others cannot close - /// panes early. - newest_clock_ms: Option, - /// The newest event-time barrier a producer declared: every pane ending - /// at or before it is closed. - barrier_ms: Option, - pub store: SummaryStore, - pub stats: PrecomputeStats, } -impl Precompute { - /// Compiles one build per definition of `sds` built in an ingestion run, - /// in dependency order; `sds` describes `dag`. +impl Ingestion { + /// Compiles one build per definition of `sds` built in an ingestion run; + /// `sds` describes `dag`. pub fn new( - dag: PhysicalASAPDAG, + dag: &PhysicalASAPDAG, sds: &Sds, options: PrecomputeOptions, ) -> Result { - let mut builders = Vec::new(); + let mut builds = Vec::new(); // Ingestion-run definitions in dependency order; query-run ones are // the query engine's. let order = sds .ingestion_build_order() .map_err(|e| PrecomputeError::Plan(e.to_string()))?; for id in &order { - builders.push(Builder { - summary: SummaryBuild::new(&dag, &sds.definitions[id])?, - open: BTreeMap::new(), - first_pane_end: None, - built_through: None, - }); + builds.push(SummaryBuild::new(dag, &sds.definitions[id])?); } // Samples of a metric only query-run summaries read move the // watermark, which says when such a summary's coverage is closed. // A query-run definition whose selector does not build is not // tracked here; the query engine forwards the queries that read it. - let mut query_run_metrics = BTreeMap::new(); + let mut query_run_metrics = BTreeSet::new(); for definition in sds.definitions.values() { if definition.built_at == StoredOutputBuild::QueryRun { - match RawSampleSelector::new(&dag, &definition.window) { + match RawSampleSelector::new(dag, &definition.window) { Ok(selector) => { - query_run_metrics.insert(selector.metric().to_string(), None); + query_run_metrics.insert(selector.metric().to_string()); } Err(error) => tracing::warn!( definition = ?definition.id, %error, @@ -257,17 +252,137 @@ impl Precompute { } let millis = |value: u64| i64::try_from(value).unwrap_or(i64::MAX); Ok(Self { - builders, + builds, query_run_metrics, allowed_lateness_ms: millis(options.allowed_lateness_ms), max_future_ms: millis(options.max_future_ms), - newest_sample_ms: None, - newest_clock_ms: None, - barrier_ms: None, - store: SummaryStore::new(), - stats: PrecomputeStats::default(), }) } +} + +/// One ingestion worker's open panes: per definition, its shard's samples +/// of each open pane, by pane end. Only the worker writes them; the thread +/// that closes a pane takes every worker's samples of it. +struct OpenPanes { + samples: Vec>>, +} + +/// The precompute counters, shared by every worker. +#[derive(Default)] +struct Counters { + samples_accepted: AtomicU64, + samples_dropped_late: AtomicU64, + samples_skipped_first_pane: AtomicU64, + samples_dropped_future: AtomicU64, + samples_unrouted: AtomicU64, + samples_query_run_only: AtomicU64, + panes_built: AtomicU64, + panes_failed: AtomicU64, + panes_skipped_expired: AtomicU64, +} + +fn add(counter: &AtomicU64, n: u64) { + if n > 0 { + counter.fetch_add(n, Ordering::Relaxed); + } +} + +/// Precompute state every ingestion worker and query thread shares: the +/// compiled builds, each worker's open panes and newest sample, the +/// watermark's other inputs, the stored summaries and the counters. +/// +/// Series are sharded across workers ([`shard_of`]), so each worker's open +/// panes hold only its shard. The global watermark is the minimum over the +/// workers that have seen a sample of their newest sample, capped at the +/// wall clock, minus the allowed lateness, or the newest barrier if that is +/// later. The worker (or barrier) that moves it past a pane end closes the +/// pane: under every worker's lock, it takes each worker's samples of the +/// pane and builds it once, so a pane never waits for an idle worker and a +/// group spanning workers is built whole. +pub struct Shared { + ingestion: Ingestion, + workers: Vec>, + /// Each worker's newest routed sample time. + newest_sample_ms: Vec, + /// Each metric query-run summaries read, with its newest sample time. + query_run_newest_ms: BTreeMap, + /// The newest wall-clock time an ingest or barrier ran at: the + /// watermark never passes it, so a series stamped ahead of the others + /// cannot close panes early. + newest_clock_ms: AtomicI64, + /// The newest event-time barrier a producer declared. + barrier_ms: AtomicI64, + /// Per definition, the end of the first pane any worker saw, which is + /// never built. + first_pane_end: Vec, + /// Per definition, the end of the newest pane built or skipped; samples + /// in it or before it are late. Written only under every worker's lock. + built_through: Vec, + /// The watermark panes are closed through, and whether a definition saw + /// its first sample since, which may leave panes to close at the same + /// watermark. + closed_through_ms: AtomicI64, + new_definition: AtomicBool, + /// Serializes closing. + closing: Mutex<()>, + /// Stored summaries, shared with the query threads. + pub store: Arc>, + counters: Counters, +} + +impl Shared { + /// `workers` ingestion workers (at least one) running `ingestion`. + pub fn new(ingestion: Ingestion, workers: usize) -> Self { + let definitions = ingestion.builds.len(); + let slots = |n: usize| (0..n).map(|_| AtomicI64::new(UNSET)).collect::>(); + Self { + workers: (0..workers.max(1)) + .map(|_| { + Mutex::new(OpenPanes { + samples: vec![BTreeMap::new(); definitions], + }) + }) + .collect(), + newest_sample_ms: slots(workers.max(1)), + query_run_newest_ms: ingestion + .query_run_metrics + .iter() + .map(|metric| (metric.clone(), AtomicI64::new(UNSET))) + .collect(), + newest_clock_ms: AtomicI64::new(UNSET), + barrier_ms: AtomicI64::new(UNSET), + first_pane_end: slots(definitions), + built_through: slots(definitions), + closed_through_ms: AtomicI64::new(UNSET), + new_definition: AtomicBool::new(false), + closing: Mutex::new(()), + store: Arc::new(SummaryStore::new()), + counters: Counters::default(), + ingestion, + } + } + + /// How many ingestion workers series are sharded across. + pub fn workers(&self) -> usize { + self.workers.len() + } + + /// The global watermark: no sample older than this is expected any + /// more. `None` before any routed sample or barrier. + pub fn watermark_ms(&self) -> Option { + let newest = self + .newest_sample_ms + .iter() + .filter_map(get) + .min() + .map(|newest| match get(&self.newest_clock_ms) { + Some(clock) => newest.min(clock), + None => newest, + }); + let from_samples = + newest.map(|newest| newest.saturating_sub(self.ingestion.allowed_lateness_ms)); + from_samples.max(get(&self.barrier_ms)) + } /// What ingestion has seen, for the query engine's keep decisions: the /// watermark and each query-run metric's newest sample. @@ -275,129 +390,172 @@ impl Precompute { crate::query::IngestionProgress { watermark_ms: self.watermark_ms(), newest_sample_ms: self - .query_run_metrics + .query_run_newest_ms .iter() - .filter_map(|(metric, newest)| Some((metric.clone(), (*newest)?))) + .filter_map(|(metric, newest)| Some((metric.clone(), get(newest)?))) .collect(), } } - /// The watermark: no sample older than this is expected any more. It is - /// the newest routed sample, capped at the wall clock, minus the allowed - /// lateness, or the newest barrier if that is later. - pub fn watermark_ms(&self) -> Option { - let newest = match self.newest_clock_ms { - Some(clock) => self.newest_sample_ms.map(|newest| newest.min(clock)), - None => self.newest_sample_ms, - }; - let from_samples = newest.map(|newest| newest.saturating_sub(self.allowed_lateness_ms)); - from_samples.max(self.barrier_ms) - } - - /// [`Self::advance_watermark_at`] the current wall-clock time. - pub fn advance_watermark(&mut self, barrier_ms: i64) -> i64 { - self.advance_watermark_at(barrier_ms, wall_clock_ms()) + pub fn stats(&self) -> PrecomputeStats { + let c = &self.counters; + let load = |counter: &AtomicU64| counter.load(Ordering::Relaxed); + PrecomputeStats { + samples_accepted: load(&c.samples_accepted), + samples_dropped_late: load(&c.samples_dropped_late), + samples_skipped_first_pane: load(&c.samples_skipped_first_pane), + samples_dropped_future: load(&c.samples_dropped_future), + samples_unrouted: load(&c.samples_unrouted), + samples_query_run_only: load(&c.samples_query_run_only), + panes_built: load(&c.panes_built), + panes_failed: load(&c.panes_failed), + panes_skipped_expired: load(&c.panes_skipped_expired), + } } - /// An event-time barrier: the producer states that every sample at or - /// before `barrier_ms` has been written. Closes every pane with - /// `end <= barrier_ms` (panes are `(end - width, end]`), so a pane whose - /// series stopped sending closes without a later sample. The barrier is - /// capped at `now_ms`; a repeated or lower barrier does nothing. A sample - /// that arrives later for a pane the barrier closed counts as late. - /// Returns the barrier in force. - pub fn advance_watermark_at(&mut self, barrier_ms: i64, now_ms: i64) -> i64 { - self.newest_clock_ms = self.newest_clock_ms.max(Some(now_ms)); + /// An event-time barrier at `now_ms`: the producers state that every + /// sample at or before `barrier_ms` has been written. Closes every pane + /// with `end <= barrier_ms` (panes are `(end - width, end]`), so a pane + /// whose series stopped sending closes without a later sample, for + /// every worker's shard. The barrier is capped at `now_ms`; a repeated + /// or lower barrier does nothing. Samples at or before it that arrive + /// later count as late. Returns the barrier in force. The caller first + /// lets every worker run the writes acknowledged before. + pub fn advance_watermark_at(&self, barrier_ms: i64, now_ms: i64) -> i64 { + self.newest_clock_ms.fetch_max(now_ms, Ordering::AcqRel); let barrier_ms = barrier_ms.min(now_ms); - if self.barrier_ms.is_some_and(|current| current >= barrier_ms) { - return self.barrier_ms.unwrap_or(barrier_ms); + let before = self.barrier_ms.fetch_max(barrier_ms, Ordering::AcqRel); + if before != UNSET && before >= barrier_ms { + return before; } - self.barrier_ms = Some(barrier_ms); self.close_panes(); barrier_ms } - /// [`Self::ingest_at`] the current wall-clock time. - pub fn ingest(&mut self, samples: &[RawSample]) { - self.ingest_at(samples, wall_clock_ms()); + /// Worker `worker` takes `samples`, all of its shard, into their open + /// panes at wall-clock time `now_ms`, which bounds sample times from + /// above; then every pane the global watermark has passed is closed. + /// + /// Until a worker has seen a sample it does not hold the watermark + /// back, so a worker with no series never blocks it. Right after + /// startup, the first worker to take a request may therefore close a + /// pane before another worker takes its part of the same request; the + /// allowed lateness covers that. + pub fn ingest_at(&self, worker: usize, samples: &[RawSample], now_ms: i64) { + self.take_at(worker, samples, now_ms); + self.close_panes(); } - /// Buffers `samples` into their panes, then builds every pane the - /// watermark has passed. `now_ms` is the wall-clock time that bounds - /// sample times from above. - pub fn ingest_at(&mut self, samples: &[RawSample], now_ms: i64) { - self.newest_clock_ms = self.newest_clock_ms.max(Some(now_ms)); - let latest_allowed_ms = now_ms.saturating_add(self.max_future_ms); - for sample in samples { - let metric = sample.labels.get("__name__").map(String::as_str); - let query_run = metric.is_some_and(|m| self.query_run_metrics.contains_key(m)); - if !query_run - && !self - .builders - .iter() - .any(|b| metric == Some(b.summary.metric())) - { - self.stats.samples_unrouted += 1; - continue; - } - if sample.timestamp_ms > latest_allowed_ms { - self.stats.samples_dropped_future += 1; - continue; - } - let mut outcome: Option = None; - for builder in &mut self.builders { - if metric != Some(builder.summary.metric()) { + fn take_at(&self, worker: usize, samples: &[RawSample], now_ms: i64) { + self.newest_clock_ms.fetch_max(now_ms, Ordering::AcqRel); + let ingestion = &self.ingestion; + let latest_allowed_ms = now_ms.saturating_add(ingestion.max_future_ms); + let mut outcomes = [0u64; 6]; + let (late, first_pane, buffered, query_run_only, unrouted, future) = (0, 1, 2, 3, 4, 5); + let mut newest: Option = None; + { + let mut open = self.workers[worker] + .lock() + .unwrap_or_else(|e| e.into_inner()); + for sample in samples { + let metric = sample.labels.get("__name__").map(String::as_str); + let query_run = metric.is_some_and(|m| ingestion.query_run_metrics.contains(m)); + if !query_run && !ingestion.builds.iter().any(|b| metric == Some(b.metric())) { + outcomes[unrouted] += 1; continue; } - let this = match pane_end(sample.timestamp_ms, builder.summary.definition.pane_ms) { - None => Outcome::Late, - Some(end) => { - let first = *builder.first_pane_end.get_or_insert(end); - let built = *builder.built_through.get_or_insert(end); - if end == first { - Outcome::FirstPane - } else if end <= built { - Outcome::Late - } else { - builder.open.entry(end).or_default().push(sample.clone()); - Outcome::Buffered - } + if sample.timestamp_ms > latest_allowed_ms { + outcomes[future] += 1; + continue; + } + let mut outcome: Option = None; + for (index, build) in ingestion.builds.iter().enumerate() { + if metric != Some(build.metric()) { + continue; } - }; - outcome = outcome.max(Some(this)); - } - match outcome { - None => self.stats.samples_query_run_only += 1, - Some(Outcome::Buffered) => self.stats.samples_accepted += 1, - Some(Outcome::FirstPane) => self.stats.samples_skipped_first_pane += 1, - Some(Outcome::Late) => self.stats.samples_dropped_late += 1, + let this = match pane_end(sample.timestamp_ms, build.definition.pane_ms) { + None => Outcome::Late, + Some(end) => { + if get(&self.first_pane_end[index]).is_none() { + self.new_definition.store(true, Ordering::Release); + } + let first = get_or_init(&self.first_pane_end[index], end); + let built = get_or_init(&self.built_through[index], first); + if end == first { + Outcome::FirstPane + } else if end <= built { + Outcome::Late + } else { + open.samples[index] + .entry(end) + .or_default() + .push(sample.clone()); + Outcome::Buffered + } + } + }; + outcome = outcome.max(Some(this)); + } + outcomes[match outcome { + None => query_run_only, + Some(Outcome::Buffered) => buffered, + Some(Outcome::FirstPane) => first_pane, + Some(Outcome::Late) => late, + }] += 1; + newest = newest.max(Some(sample.timestamp_ms)); + if let Some(newest) = metric.and_then(|m| self.query_run_newest_ms.get(m)) { + newest.fetch_max(sample.timestamp_ms, Ordering::AcqRel); + } } - self.newest_sample_ms = self.newest_sample_ms.max(Some(sample.timestamp_ms)); - if let Some(newest) = metric.and_then(|m| self.query_run_metrics.get_mut(m)) { - *newest = (*newest).max(Some(sample.timestamp_ms)); + if let Some(newest) = newest { + self.newest_sample_ms[worker].fetch_max(newest, Ordering::AcqRel); } } - self.close_panes(); + let c = &self.counters; + add(&c.samples_dropped_late, outcomes[late]); + add(&c.samples_skipped_first_pane, outcomes[first_pane]); + add(&c.samples_accepted, outcomes[buffered]); + add(&c.samples_query_run_only, outcomes[query_run_only]); + add(&c.samples_unrouted, outcomes[unrouted]); + add(&c.samples_dropped_future, outcomes[future]); } - fn close_panes(&mut self) { + /// Builds every pane the global watermark has passed, from every + /// worker's samples of it. + fn close_panes(&self) { + let closed = |watermark: i64| { + get(&self.closed_through_ms).is_some_and(|closed| watermark <= closed) + && !self.new_definition.load(Ordering::Acquire) + }; + match self.watermark_ms() { + Some(watermark) if !closed(watermark) => {} + _ => return, + } + let _closing = self.closing.lock().unwrap_or_else(|e| e.into_inner()); + let mut open: Vec<_> = self + .workers + .iter() + .map(|worker| worker.lock().unwrap_or_else(|e| e.into_inner())) + .collect(); + // Under every worker's lock: no sample is being taken, so the + // watermark read now covers every sample taken. let Some(watermark) = self.watermark_ms() else { return; }; - for builder in &mut self.builders { - let Some(built) = builder.built_through else { + self.new_definition.store(false, Ordering::Release); + for (index, build) in self.ingestion.builds.iter().enumerate() { + let Some(built) = get(&self.built_through[index]) else { continue; }; - let pane_ms = builder.summary.definition.pane_ms; - let id = &builder.summary.definition.id; + let pane_ms = build.definition.pane_ms; + let id = &build.definition.id; // The newest evaluation stored panes can serve is at the end of // the newest closed pane. Panes are kept for one extra pane // width beyond the plan's retention, so an evaluation one pane // behind it is still served; older panes are no longer read: // skip them. let newest_closed_end = watermark.div_euclid(pane_ms).saturating_mul(pane_ms); - let oldest_kept_end = builder - .summary + let oldest_kept_end = build .definition .oldest_kept_end_ms(newest_closed_end.saturating_sub(pane_ms)); let Some(mut end) = built.checked_add(pane_ms) else { @@ -407,34 +565,42 @@ impl Precompute { let Some(first_kept) = pane_end(oldest_kept_end, pane_ms) else { continue; }; - self.stats.panes_skipped_expired += - (first_kept.saturating_sub(end) / pane_ms) as u64; - builder.open = builder.open.split_off(&first_kept); + add( + &self.counters.panes_skipped_expired, + (first_kept.saturating_sub(end) / pane_ms) as u64, + ); + for worker in open.iter_mut() { + let samples = &mut worker.samples[index]; + *samples = samples.split_off(&first_kept); + } end = first_kept; } // Every pane after the first through the watermark is built, so // a pane without samples is stored as empty, not missing. while end <= watermark { - let samples = builder.open.remove(&end).unwrap_or_default(); + let mut samples = Vec::new(); + for worker in open.iter_mut() { + samples.extend(worker.samples[index].remove(&end).unwrap_or_default()); + } let scope = Scope::Ingestion { window_start_ms: end - pane_ms, window_end_ms: end, revision: 1, }; - match builder.summary.run(end, &samples, scope) { + match build.run(end, &samples, scope) { Ok(state) => { self.store.insert( SummaryInstanceKey::new(id.clone(), Selection::ending_at(end)), state, ); - self.stats.panes_built += 1; + add(&self.counters.panes_built, 1); } Err(error) => { tracing::error!(definition = ?id, end, %error, "pane build failed"); - self.stats.panes_failed += 1; + add(&self.counters.panes_failed, 1); } } - builder.built_through = Some(end); + self.built_through[index].store(end, Ordering::Release); match end.checked_add(pane_ms) { Some(next) => end = next, None => break, @@ -443,11 +609,103 @@ impl Precompute { self.store .evict_before(id, Selection::ending_at(oldest_kept_end)); } + self.closed_through_ms + .fetch_max(watermark, Ordering::AcqRel); + } +} + +/// The ingestion worker of `workers` that owns the series `labels`, by a +/// hash of its label set. +pub fn shard_of(labels: &BTreeMap, workers: usize) -> usize { + use std::hash::{Hash, Hasher}; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + labels.hash(&mut hasher); + (hasher.finish() % workers.max(1) as u64) as usize +} + +/// Every ingestion worker of a plan on the caller's thread, each fed its +/// shard of the series: the threaded engine's precompute without threads. +pub struct Precompute { + shared: Arc, + /// Stored summaries. + pub store: Arc>, + /// The counters as of the last call. + pub stats: PrecomputeStats, +} + +impl Precompute { + /// One worker for every series; `sds` describes `dag`. + pub fn new( + dag: PhysicalASAPDAG, + sds: &Sds, + options: PrecomputeOptions, + ) -> Result { + Self::with_shards(&dag, sds, options, 1) + } + + /// `shards` workers, each owning the series [`shard_of`] assigns it. + pub fn with_shards( + dag: &PhysicalASAPDAG, + sds: &Sds, + options: PrecomputeOptions, + shards: usize, + ) -> Result { + let shared = Arc::new(Shared::new(Ingestion::new(dag, sds, options)?, shards)); + Ok(Self { + store: shared.store.clone(), + stats: PrecomputeStats::default(), + shared, + }) + } + + /// See [`Shared::progress`]. + pub fn progress(&self) -> crate::query::IngestionProgress { + self.shared.progress() + } + + /// The global watermark; see [`Shared::watermark_ms`]. + pub fn watermark_ms(&self) -> Option { + self.shared.watermark_ms() + } + + /// [`Self::advance_watermark_at`] the current wall-clock time. + pub fn advance_watermark(&mut self, barrier_ms: i64) -> i64 { + self.advance_watermark_at(barrier_ms, wall_clock_ms()) + } + + /// See [`Shared::advance_watermark_at`]. + pub fn advance_watermark_at(&mut self, barrier_ms: i64, now_ms: i64) -> i64 { + let barrier_ms = self.shared.advance_watermark_at(barrier_ms, now_ms); + self.stats = self.shared.stats(); + barrier_ms + } + + /// [`Self::ingest_at`] the current wall-clock time. + pub fn ingest(&mut self, samples: &[RawSample]) { + self.ingest_at(samples, wall_clock_ms()); + } + + /// Routes `samples` to their workers, which take them into their panes + /// and close every pane the global watermark has passed. `now_ms` is + /// the wall-clock time that bounds sample times from above. + pub fn ingest_at(&mut self, samples: &[RawSample], now_ms: i64) { + let workers = self.shared.workers(); + let mut shards = vec![Vec::new(); workers]; + for sample in samples { + shards[shard_of(&sample.labels, workers)].push(sample.clone()); + } + // Every worker takes its shard before any closes a pane, as if the + // request reached them all at once. + for (worker, samples) in shards.iter().enumerate() { + self.shared.take_at(worker, samples, now_ms); + } + self.shared.close_panes(); + self.stats = self.shared.stats(); } } /// Milliseconds since the Unix epoch. -fn wall_clock_ms() -> i64 { +pub fn wall_clock_ms() -> i64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX)) diff --git a/data_plane/src/query.rs b/data_plane/src/query.rs index f0999097..43ce6cbc 100644 --- a/data_plane/src/query.rs +++ b/data_plane/src/query.rs @@ -18,6 +18,7 @@ //! forwards it to Prometheus. use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use asap_executor::operators::Operator; @@ -153,6 +154,19 @@ pub struct KeptResultStats { pub reused: u64, } +/// [`KeptResultStats`] as counters every query thread adds to. +#[derive(Default)] +struct KeptResultCounters { + built: AtomicU64, + kept: AtomicU64, + kept_while_metric_behind: AtomicU64, + reused: AtomicU64, +} + +fn add(counter: &AtomicU64) { + counter.fetch_add(1, Ordering::Relaxed); +} + /// What ingestion had seen when an evaluation decided what to fetch. The /// keep decision uses this snapshot, not the state when the evaluation /// runs: samples ingested in between were not in the fetched data. @@ -175,9 +189,11 @@ pub struct Prepared { pub progress: IngestionProgress, } +/// The plan's compiled roots, compiled once and shared by every query +/// thread. pub struct QueryEngine { /// How many query-run summaries evaluations built, kept and reused. - pub kept_stats: KeptResultStats, + kept_stats: KeptResultCounters, /// How each query-run definition's summaries are built, or why they /// cannot be. kept_builds: BTreeMap>, @@ -221,7 +237,7 @@ impl QueryEngine { }) .collect(); Self { - kept_stats: KeptResultStats::default(), + kept_stats: KeptResultCounters::default(), kept_builds, kept_matchers, roots, @@ -301,10 +317,10 @@ impl QueryEngine { /// measured from the newest closed pane at or before the watermark, are /// evicted. pub fn evaluate( - &mut self, + &self, root: usize, time_ms: i64, - store: &mut SummaryStore, + store: &SummaryStore, progress: &IngestionProgress, raw: &BTreeMap>, ) -> Result, Unanswered> { @@ -320,8 +336,9 @@ impl QueryEngine { let mut sources: BTreeMap> = BTreeMap::new(); for input in &program.stored { let pane = pane(input, time_ms, store)?; + // Cloning a batch shares its buffers: the pane is not copied. let source = - Operator::source(pane.schema().clone(), vec![pane.clone()]).map_err(failed)?; + Operator::source(pane.schema().clone(), vec![(*pane).clone()]).map_err(failed)?; sources.insert(input.node, Box::new(source)); } for input in &program.kept { @@ -332,8 +349,8 @@ impl QueryEngine { let selection = Selection::ending_at(end_ms); let summary = match store.get(&input.definition, selection) { Some(summary) => { - self.kept_stats.reused += 1; - summary.clone() + add(&self.kept_stats.reused); + (*summary).clone() } None => { // `raw_ranges` asked for these samples unless the @@ -355,21 +372,22 @@ impl QueryEngine { let summary = build .run(end_ms, samples, build_scope) .map_err(|e| Unanswered::Execution(e.to_string()))?; - self.kept_stats.built += 1; + add(&self.kept_stats.built); // Kept only once closed, as the precompute engine closes // panes: a sample may still arrive for an open one. if progress.watermark_ms.is_some_and(|w| end_ms <= w) { - self.kept_stats.kept += 1; + add(&self.kept_stats.kept); let metric_newest = progress.newest_sample_ms.get(build.metric()); if metric_newest.is_none_or(|newest| *newest < end_ms) { - self.kept_stats.kept_while_metric_behind += 1; + add(&self.kept_stats.kept_while_metric_behind); } - store.insert( - SummaryInstanceKey::new(input.definition.clone(), selection), - summary.clone(), - ); + // Another query thread may have stored it meanwhile: + // the first stored is the one every evaluation reads. + let key = SummaryInstanceKey::new(input.definition.clone(), selection); + (*store.insert(key, summary)).clone() + } else { + summary } - summary } }; let source = @@ -405,6 +423,18 @@ impl QueryEngine { result_samples(&program.output, rows) } + /// How many query-run summaries evaluations built, kept and reused. + pub fn kept_stats(&self) -> KeptResultStats { + let load = |counter: &AtomicU64| counter.load(Ordering::Relaxed); + let c = &self.kept_stats; + KeptResultStats { + built: load(&c.built), + kept: load(&c.kept), + kept_while_metric_behind: load(&c.kept_while_metric_behind), + reused: load(&c.reused), + } + } + /// The build of a query-run definition a compiled root reads. fn kept_build(&self, definition: &DefinitionId) -> &SummaryBuild { self.kept_builds[definition] @@ -437,11 +467,11 @@ fn window(time_ms: i64, offset_ms: i64, width_ms: i64) -> Result<(i64, i64), Una } /// The stored pane `input` reads at `time_ms`. -fn pane<'a>( +fn pane( input: &StoredInput, time_ms: i64, - store: &'a SummaryStore, -) -> Result<&'a Batch, Unanswered> { + store: &SummaryStore, +) -> Result, 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!( diff --git a/data_plane/src/sds.rs b/data_plane/src/sds.rs index f45b37f2..04364964 100644 --- a/data_plane/src/sds.rs +++ b/data_plane/src/sds.rs @@ -22,6 +22,7 @@ //! all its stored inputs are sealed. use std::collections::{BTreeMap, HashMap}; +use std::sync::{Arc, RwLock}; use asap_executor::planner::ir::physical_export::{ PhysicalASAPDAG, PhysicalASAPDAGNode, RawSampleWindow, @@ -320,11 +321,15 @@ impl SummaryInstanceKey { } } -/// Stored summary instances, by definition and selection. An instance is -/// stored once, when its coverage is complete, and never changes. +/// Stored summary instances, by definition and selection, shared by every +/// thread that builds or reads them. An instance is stored once, when its +/// coverage is complete, and never changes: a read takes the read lock only +/// to clone the instance's `Arc`, so readers share it without copying and +/// hold no lock while they compute. Inserts and evictions take the write +/// lock briefly. #[derive(Debug)] pub struct SummaryStore { - instances: BTreeMap>, + instances: RwLock>>>, } impl Default for SummaryStore { @@ -336,36 +341,59 @@ impl Default for SummaryStore { impl SummaryStore { pub fn new() -> Self { Self { - instances: BTreeMap::new(), + instances: RwLock::new(BTreeMap::new()), } } - pub fn insert(&mut self, key: SummaryInstanceKey, state: S) { - self.instances + fn read( + &self, + ) -> std::sync::RwLockReadGuard<'_, BTreeMap>>> { + self.instances.read().unwrap_or_else(|e| e.into_inner()) + } + + fn write( + &self, + ) -> std::sync::RwLockWriteGuard<'_, BTreeMap>>> { + self.instances.write().unwrap_or_else(|e| e.into_inner()) + } + + /// Stores `state` unless an instance of `key` is already stored, and + /// returns the stored instance: an instance is never rewritten, so two + /// threads that build the same one read the same. + pub fn insert(&self, key: SummaryInstanceKey, state: S) -> Arc { + self.write() .entry(key.definition) .or_default() - .insert(key.selection, state); + .entry(key.selection) + .or_insert_with(|| Arc::new(state)) + .clone() } /// The instance of `definition` covering `selection`, if stored. Borrows /// the definition id, so a lookup allocates nothing. - pub fn get(&self, definition: &DefinitionId, selection: Selection) -> Option<&S> { - self.instances.get(definition)?.get(&selection) + pub fn get(&self, definition: &DefinitionId, selection: Selection) -> Option> { + self.read().get(definition)?.get(&selection).cloned() } /// Drops the instances of `definition` whose coverage ends before /// `oldest`'s. - pub fn evict_before(&mut self, definition: &DefinitionId, oldest: Selection) { - if let Some(instances) = self.instances.get_mut(definition) { - *instances = instances.split_off(&oldest); + pub fn evict_before(&self, definition: &DefinitionId, oldest: Selection) { + let mut instances = self.write(); + if let Some(instances) = instances.get_mut(definition) { + if instances + .first_key_value() + .is_some_and(|(first, _)| *first < oldest) + { + *instances = instances.split_off(&oldest); + } } } pub fn len(&self) -> usize { - self.instances.values().map(BTreeMap::len).sum() + self.read().values().map(BTreeMap::len).sum() } pub fn is_empty(&self) -> bool { - self.instances.values().all(BTreeMap::is_empty) + self.read().values().all(BTreeMap::is_empty) } } diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs index c6ac3a27..e66da5ca 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -493,3 +493,106 @@ async fn barrier_serves_the_window_of_a_stopped_series() { ); } } + +/// A data plane for Example 4b on `workers` ingestion workers and `threads` +/// query threads, with no lateness. +async fn sharded_example4b(workers: usize, threads: usize) -> (Router, Arc) { + let mut options = data_plane::engine::EngineOptions { + ingestion_workers: workers, + query_threads: threads, + ..Default::default() + }; + options.precompute.allowed_lateness_ms = 0; + let state = Arc::new( + DataPlane::with_options( + common::plan_json("example4b", false), + fake_prometheus(vec![]).await, + options, + ) + .unwrap(), + ); + (router(state.clone()), state) +} + +/// Writes one sample every 15 s in `[from, to)` of each job's `latency_ms`. +async fn write_jobs(app: &Router, jobs: &[String], from: i64, to: i64) { + let timeseries = jobs + .iter() + .enumerate() + .map(|(i, job)| TimeSeries { + labels: vec![ + Label { + name: "__name__".into(), + value: "latency_ms".into(), + }, + Label { + name: "job".into(), + value: job.clone(), + }, + ], + samples: (from..to) + .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); +} + +/// 32 queries at once, on 4 query threads over 4 ingestion workers, get +/// the answers they get one at a time. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn answers_concurrent_queries_like_one_at_a_time() { + let (app, _) = sharded_example4b(4, 4).await; + let jobs: Vec = (0..16).map(|i| format!("job{i}")).collect(); + write_jobs(&app, &jobs, MINUTE, 95 * MINUTE).await; + let query = "quantile_over_time(0.99, latency_ms[60m])"; + let times: Vec = (0..32).map(|i| (80 + 10 * (i % 2)) * MINUTE).collect(); + let mut sequential = Vec::new(); + for &time in × { + let response = instant(&app, query, time).await; + assert_eq!(response.execution.as_deref(), Some("plan"), "{time}"); + assert_eq!(by_job(&response.body).len(), 16); + sequential.push(response.body); + } + let concurrent = futures::future::join_all(times.iter().map(|&time| { + let app = app.clone(); + async move { instant(&app, query, time).await.body } + })) + .await; + assert_eq!(concurrent, sequential); +} + +/// A query sees the panes of every write acknowledged before it, with no +/// wait in between, though the write's series go to different workers. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn a_query_reads_the_panes_of_the_writes_before_it() { + let jobs: Vec = (0..8).map(|i| format!("job{i}")).collect(); + let query = "quantile_over_time(0.99, latency_ms[60m])"; + for _ in 0..8 { + let (app, _) = sharded_example4b(4, 2).await; + for start in (0..8).map(|i| MINUTE + i * 10 * MINUTE) { + write_jobs(&app, &jobs, start, start + 10 * MINUTE).await; + } + let response = instant(&app, query, 80 * MINUTE).await; + assert_eq!( + response.execution.as_deref(), + Some("plan"), + "{:?}", + response.reason + ); + assert_eq!(by_job(&response.body).len(), 8); + } +} diff --git a/data_plane/tests/kept.rs b/data_plane/tests/kept.rs index 47496865..e9031afe 100644 --- a/data_plane/tests/kept.rs +++ b/data_plane/tests/kept.rs @@ -177,7 +177,7 @@ fn kept_summaries_answer_like_the_reference_and_are_reused() { assert_eq!(first, answer(&mut reference, 10 * MINUTE, &all).0); assert_eq!(first.len(), 2); assert_eq!( - k.engine.kept_stats, + k.engine.kept_stats(), KeptResultStats { built: 5, kept: 5, @@ -191,7 +191,7 @@ fn kept_summaries_answer_like_the_reference_and_are_reused() { assert_eq!(second, answer(&mut reference, 11 * MINUTE, &all).0); assert_ne!(second, first, "the answer moves with the window"); assert_eq!( - (k.engine.kept_stats.built, k.engine.kept_stats.reused), + (k.engine.kept_stats().built, k.engine.kept_stats().reused), (6, 4) ); @@ -200,11 +200,11 @@ fn kept_summaries_answer_like_the_reference_and_are_reused() { assert!(ranges.is_empty()); assert_eq!(again, second); assert_eq!( - (k.engine.kept_stats.built, k.engine.kept_stats.reused), + (k.engine.kept_stats().built, k.engine.kept_stats().reused), (6, 9) ); assert_eq!( - reference.engine.kept_stats, + reference.engine.kept_stats(), KeptResultStats::default(), "the reference keeps nothing" ); @@ -351,6 +351,10 @@ fn counts_summaries_kept_while_their_metric_is_behind() { &raw, ) .unwrap(); - assert_eq!(k.engine.kept_stats.kept, 5); - assert_eq!(k.engine.kept_stats.kept_while_metric_behind, 2, "9 and 10"); + assert_eq!(k.engine.kept_stats().kept, 5); + assert_eq!( + k.engine.kept_stats().kept_while_metric_behind, + 2, + "9 and 10" + ); } diff --git a/data_plane/tests/precompute.rs b/data_plane/tests/precompute.rs index ce848e79..43065b1c 100644 --- a/data_plane/tests/precompute.rs +++ b/data_plane/tests/precompute.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use asap_executor::planner::workload::{PlanningWorkload, Query}; use asap_planner::InstallablePlan; use control_plane::plan_workload; -use data_plane::precompute::{pane_end, Precompute, PrecomputeOptions, RawSample}; +use data_plane::precompute::{pane_end, shard_of, Precompute, PrecomputeOptions, RawSample}; use data_plane::sds::{Sds, Selection}; const MINUTE: i64 = 60_000; @@ -600,3 +600,111 @@ fn input_stays_open_after_a_barrier() { engine.ingest_at(&[sample("latency_ms", "a", 29 * MINUTE, 1.0)], 36 * MINUTE); assert_eq!(engine.stats.samples_dropped_late, 1); } + +fn sharded(plan: &InstallablePlan, shards: usize) -> (Precompute, Sds) { + let sds = Sds::from_plan(plan).unwrap(); + let options = PrecomputeOptions { + allowed_lateness_ms: 0, + ..PrecomputeOptions::default() + }; + let engine = Precompute::with_shards(&plan.dag.dag, &sds, options, shards).unwrap(); + (engine, sds) +} + +/// Two jobs whose series two shards own one each. +fn jobs_on_two_shards() -> (String, String) { + let shard = |job: &str| shard_of(&sample("latency_ms", job, 0, 0.0).labels, 2); + let a = (0..) + .map(|i| format!("j{i}")) + .find(|j| shard(j) == 0) + .unwrap(); + let b = (0..) + .map(|i| format!("j{i}")) + .find(|j| shard(j) == 1) + .unwrap(); + (a, b) +} + +/// With two shards, a pane closes only once both have passed it: the +/// global watermark is the minimum over the shards' newest samples. +#[test] +fn a_pane_closes_once_every_shard_has_passed_it() { + let (mut engine, sds) = sharded(&common::plan("example4b", false), 2); + let id = sds.definitions.keys().next().unwrap().clone(); + let (a, b) = jobs_on_two_shards(); + engine.ingest(&samples("latency_ms", &[&a, &b], MINUTE, 15 * MINUTE)); + engine.ingest(&samples("latency_ms", &[&a], 15 * MINUTE, 35 * MINUTE)); + assert_eq!(engine.watermark_ms(), Some(15 * MINUTE - 15_000)); + assert!( + engine + .store + .get(&id, Selection::ending_at(2 * PANE)) + .is_none(), + "the shard of {b} is still in (10, 20 min]" + ); + engine.ingest(&[sample("latency_ms", &b, 21 * MINUTE, 1.0)]); + let pane = engine + .store + .get(&id, Selection::ending_at(2 * PANE)) + .unwrap(); + assert_eq!(pane.num_rows(), 2, "both shards' series"); + assert!(engine + .store + .get(&id, Selection::ending_at(3 * PANE)) + .is_none()); +} + +/// A barrier closes every shard's panes, the idle shard's included. +#[test] +fn a_barrier_closes_every_shards_panes() { + let (mut engine, sds) = sharded(&common::plan("example4b", false), 2); + let id = sds.definitions.keys().next().unwrap().clone(); + let (a, b) = jobs_on_two_shards(); + engine.ingest_at(&samples("latency_ms", &[&b], MINUTE, 15 * MINUTE), DAY); + engine.ingest_at(&samples("latency_ms", &[&a], MINUTE, 35 * MINUTE), DAY); + assert!(engine + .store + .get(&id, Selection::ending_at(2 * PANE)) + .is_none()); + engine.advance_watermark_at(3 * PANE, DAY); + for (end, rows) in [(2 * PANE, 2), (3 * PANE, 1)] { + let pane = engine.store.get(&id, Selection::ending_at(end)).unwrap(); + assert_eq!(pane.num_rows(), rows, "pane ending at {end}"); + } +} + +/// Panes built on four shards equal those built on one, row for row: +/// Example 4b's per-series KLL and Example 1's rates and sums, with series +/// spread over the shards. +#[test] +fn sharded_panes_answer_like_one_shard() { + for (example, metric) in [ + ("example4b", "latency_ms"), + ("example1", "http_requests_total"), + ] { + let plan = common::plan(example, false); + let jobs: Vec = (0..12).map(|i| format!("job{i}")).collect(); + let jobs: Vec<&str> = jobs.iter().map(String::as_str).collect(); + let all = samples(metric, &jobs, MINUTE, 75 * MINUTE); + let (mut one, sds) = sharded(&plan, 1); + let (mut four, _) = sharded(&plan, 4); + // Whole timestamps per request, as a scraper writes them. + for chunk in all.chunks(4 * jobs.len()) { + one.ingest_at(chunk, DAY); + four.ingest_at(chunk, DAY); + } + assert_eq!(one.stats, four.stats, "{example}"); + assert!(one.stats.panes_built > 0, "{example}"); + for id in sds.definitions.keys() { + for end in (0..=75 * MINUTE).step_by(10_000) { + let rows = |engine: &Precompute| { + engine + .store + .get(id, Selection::ending_at(end)) + .map(|pane| format!("{:?}", pane.to_rows())) + }; + assert_eq!(rows(&one), rows(&four), "{example} {id:?} {end}"); + } + } + } +} diff --git a/data_plane/tests/sds.rs b/data_plane/tests/sds.rs index cdfecf37..819451bc 100644 --- a/data_plane/tests/sds.rs +++ b/data_plane/tests/sds.rs @@ -164,7 +164,7 @@ fn build_order_follows_stored_inputs() { #[test] fn store_finds_and_evicts_panes() { let (a, b) = (DefinitionId("a".into()), DefinitionId("b".into())); - let mut store = SummaryStore::new(); + let store = SummaryStore::new(); for end in [60, 120, 180] { store.insert( SummaryInstanceKey::new(a.clone(), Selection::ending_at(end)), @@ -176,18 +176,27 @@ fn store_finds_and_evicts_panes() { "b60".to_string(), ); assert_eq!( - store.get(&a, Selection::ending_at(120)).map(String::as_str), + store + .get(&a, Selection::ending_at(120)) + .as_deref() + .map(String::as_str), Some("a120") ); assert_eq!(store.get(&a, Selection::ending_at(90)), None); store.evict_before(&a, Selection::ending_at(180)); assert_eq!(store.get(&a, Selection::ending_at(120)), None); assert_eq!( - store.get(&a, Selection::ending_at(180)).map(String::as_str), + store + .get(&a, Selection::ending_at(180)) + .as_deref() + .map(String::as_str), Some("a180") ); assert_eq!( - store.get(&b, Selection::ending_at(60)).map(String::as_str), + store + .get(&b, Selection::ending_at(60)) + .as_deref() + .map(String::as_str), Some("b60") ); assert_eq!(store.len(), 2); @@ -275,3 +284,54 @@ fn kept_summaries_share_one_query_run_definition() { offsets.sort(); assert_eq!(offsets, (0..5).map(|i| i * 60_000).collect::>()); } + +/// An instance is stored once and never rewritten: a second insert of the +/// same key returns the first instance, and readers share it. +#[test] +fn store_shares_instances_and_never_rewrites_them() { + let a = DefinitionId("a".into()); + let store = SummaryStore::new(); + let key = || SummaryInstanceKey::new(a.clone(), Selection::ending_at(60)); + let first = store.insert(key(), "first".to_string()); + let second = store.insert(key(), "second".to_string()); + assert!(std::sync::Arc::ptr_eq(&first, &second)); + let read = store.get(&a, Selection::ending_at(60)).unwrap(); + assert!(std::sync::Arc::ptr_eq(&first, &read)); + assert_eq!(read.as_str(), "first"); +} + +/// Readers on several threads find every instance while another thread +/// inserts and evicts. +#[test] +fn store_is_read_concurrently() { + let a = DefinitionId("a".into()); + let store = SummaryStore::new(); + for end in 0..100 { + store.insert( + SummaryInstanceKey::new(a.clone(), Selection::ending_at(end)), + end, + ); + } + std::thread::scope(|scope| { + for _ in 0..8 { + scope.spawn(|| { + for _ in 0..100 { + for end in 50..100 { + let found = store.get(&a, Selection::ending_at(end)); + assert_eq!(found.as_deref(), Some(&end)); + } + } + }); + } + scope.spawn(|| { + for end in 100..200 { + store.insert( + SummaryInstanceKey::new(a.clone(), Selection::ending_at(end)), + end, + ); + store.evict_before(&a, Selection::ending_at(end.min(150) - 100)); + } + }); + }); + assert_eq!(store.len(), 150); +}