feat(query): answer planned instant queries by running the plan's roots

README.md

@@ -35,7 +35,7 @@ The backend is being rebuilt on the Planner's physical DAG:
 | Data plane: load and validate `plan.json`, accept Remote Write | done |
 | Stored summaries (SDS): one definition per computation, panes by end time | done |
 | Precompute engine: build ingestion-time outputs once per tumbling pane | done |
-| Query engine: run query-time nodes over stored panes and raw data | planned |
+| Query engine: answer planned instant queries from stored panes and raw samples | done |
 
 The data plane stores the plan's `stored_outputs`. Outputs that differ only
 in which pane they read share one stored summary, kept for the largest
@@ -53,9 +53,13 @@ are dropped. Remote Write gets 503 when admitting it would leave more than
 stopped, which also makes `/-/ready` fail. `/api/v1/status/asap` reports the
 engine's sample and pane counters as JSON.
 
-Until the query engine lands, the data plane forwards every query to
-Prometheus, with its end-to-end headers such as `Authorization` and
-`X-Scope-OrgID`, and marks the response `x-asap-execution: forwarded`. SQL and
+An instant query whose text is one of the plan's queries is answered by the
+plan (`x-asap-execution: plan`). Its stored inputs read the panes ending at
+`T - offset`, so `T` must be on the pane grid and those panes stored; its
+query-time bound raw-sample windows are read from Prometheus. Anything else is forwarded
+to Prometheus with its end-to-end headers, such as `Authorization` and
+`X-Scope-OrgID` (`x-asap-execution: forwarded`); a planned query that could not
+be answered also says why in `x-asap-forward-reason`. Range queries, SQL and
 ClickHouse are not supported yet.
 
 ## Running
@@ -71,8 +75,8 @@ cargo run -p data_plane -- --plan plan.json --prometheus-url http://localhost:90
 
 `--deployment-capabilities caps.json` replaces the default
 `DeploymentCapabilities`: every summary the Planner can build, raw data kept
-in Prometheus, and no query-time result kept across evaluations
-(`MatAtQueryKeptAcrossEvaluations` is not implemented). For example
+in Prometheus, and no stored outputs built in a `QueryRun`
+(`MatAtQueryKeptAcrossEvaluations`). For example
 `{"raw_data_retained": false}` makes the Planner build every input at
 ingestion time. `--measured-operation-costs costs.json` gives Stage 3 the
 deployment's measured sketch operation costs; otherwise it prices every node

control_plane/src/lib.rs

@@ -83,7 +83,7 @@ pub fn load_plan(path: &Path) -> Result<InstallablePlan, PlanError> {
 
 /// What this backend's data plane can execute and keep: every summary the
 /// Planner executor builds, raw samples kept in Prometheus, and no
-/// query-time result kept across evaluations: the SDS describes such
+/// stored outputs built in a `QueryRun` (`MatAtQueryKeptAcrossEvaluations`): the SDS describes such
 /// outputs, but nothing builds them until the query engine does.
 pub fn backend_capabilities() -> DeploymentCapabilities {
     DeploymentCapabilities {

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

data_plane/src/engine.rs

@@ -2,13 +2,17 @@
 //! it. Planner IR holds `Rc`s, so the plan cannot be shared across the HTTP
 //! server's threads; handlers send the engine commands instead.
 
+use std::collections::BTreeMap;
 use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::sync::mpsc::{channel, Receiver, Sender};
 use std::sync::{Arc, Mutex};
 
 use asap_planner::InstallablePlan;
 
+use tokio::sync::oneshot;
+
 use crate::precompute::{Precompute, PrecomputeOptions, PrecomputeStats, RawSample};
+use crate::query::{BoundRawSampleWindow, QueryEngine, ResultSample, Unanswered};
 use crate::sds::Sds;
 
 /// How the engine thread ingests samples.
@@ -31,8 +35,21 @@ impl Default for EngineOptions {
     }
 }
 
+type Reply<T> = oneshot::Sender<Result<T, Unanswered>>;
+
 enum Command {
     Ingest(Vec<RawSample>),
+    BoundRawSampleWindows {
+        query: String,
+        time_ms: i64,
+        reply: Reply<(usize, Vec<BoundRawSampleWindow>)>,
+    },
+    Evaluate {
+        root: usize,
+        time_ms: i64,
+        raw: BTreeMap<u64, Vec<RawSample>>,
+        reply: Reply<Vec<ResultSample>>,
+    },
     Stop,
 }
 
@@ -135,6 +152,44 @@ impl EngineHandle {
         })
     }
 
+    /// The plan root answering `query` at `time_ms` and the raw samples it
+    /// needs; see [`QueryEngine::bound_raw_sample_windows`].
+    pub async fn bound_raw_sample_windows(
+        &self,
+        query: String,
+        time_ms: i64,
+    ) -> Result<(usize, Vec<BoundRawSampleWindow>), Unanswered> {
+        self.ask(|reply| Command::BoundRawSampleWindows {
+            query,
+            time_ms,
+            reply,
+        })
+        .await
+    }
+
+    /// Runs a root over stored panes and `raw`; see [`QueryEngine::evaluate`].
+    pub async fn evaluate(
+        &self,
+        root: usize,
+        time_ms: i64,
+        raw: BTreeMap<u64, Vec<RawSample>>,
+    ) -> Result<Vec<ResultSample>, Unanswered> {
+        self.ask(|reply| Command::Evaluate {
+            root,
+            time_ms,
+            raw,
+            reply,
+        })
+        .await
+    }
+
+    async fn ask<T>(&self, command: impl FnOnce(Reply<T>) -> Command) -> Result<T, Unanswered> {
+        let stopped = || Unanswered::Execution("the engine has stopped".into());
+        let (reply, answer) = oneshot::channel();
+        self.commands.send(command(reply)).map_err(|_| stopped())?;
+        answer.await.map_err(|_| stopped())?
+    }
+
     /// Whether the engine thread is still running.
     pub fn is_running(&self) -> bool {
         self.running.load(Ordering::SeqCst)
@@ -181,6 +236,7 @@ struct Shared {
 
 struct Engine {
     precompute: Precompute,
+    query: QueryEngine,
 }
 
 impl Engine {
@@ -193,9 +249,23 @@ impl Engine {
             outputs = sds.outputs.len(),
             "stored summaries"
         );
+        let query_run = sds
+            .definitions
+            .values()
+            .filter(|d| d.built_at == asap_planner::StoredOutputBuild::QueryRun)
+            .count();
+        if query_run > 0 {
+            tracing::warn!(
+                definitions = query_run,
+                "the plan has stored outputs built in a QueryRun, which this data plane does not build yet; \
+                 queries that read them are forwarded"
+            );
+        }
         let queries = plan.queries.iter().map(|q| q.query.clone()).collect();
+        let query = QueryEngine::new(&plan, &sds);
         let engine = Self {
             precompute: Precompute::new(plan.dag.dag, &sds, options)?,
+            query,
         };
         Ok((engine, queries))
     }
@@ -210,6 +280,30 @@ impl Engine {
                         .queued_samples
                         .fetch_sub(samples.len() as u64, Ordering::SeqCst);
                 }
+                Command::BoundRawSampleWindows {
+                    query,
+                    time_ms,
+                    reply,
+                } => {
+                    let _ = reply.send(self.query.bound_raw_sample_windows(
+                        &query,
+                        time_ms,
+                        &self.precompute.store,
+                    ));
+                }
+                Command::Evaluate {
+                    root,
+                    time_ms,
+                    raw,
+                    reply,
+                } => {
+                    let _ = reply.send(self.query.evaluate(
+                        root,
+                        time_ms,
+                        &self.precompute.store,
+                        &raw,
+                    ));
+                }
                 Command::Stop => break,
             }
         }

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,14 +184,91 @@ async fn status(State(state): State<Arc<DataPlane>>) -> 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<Arc<DataPlane>>,
     method: Method,
     query: RawQuery,
     headers: HeaderMap,
     body: Bytes,
 ) -> Response {
-    forward(state, "query", method, query, headers, body).await
+    let instant = match InstantQuery::parse(query.0.as_deref(), &body) {
+        Ok(Some(instant)) => instant,
+        Ok(None) => return forward(state, "query", method, query, headers, body).await,
+        Err(bad) => {
+            let body = axum::Json(prometheus::bad_data(&bad.to_string()));
+            return (StatusCode::BAD_REQUEST, body).into_response();
+        }
+    };
+    let reason = match answer(&state, instant, &headers).await {
+        Ok(response) => {
+            let mut response = axum::Json(response).into_response();
+            response
+                .headers_mut()
+                .insert(EXECUTION_HEADER, HeaderValue::from_static("plan"));
+            return response;
+        }
+        Err(Unanswered::NotPlanned(_)) => None,
+        Err(reason) => {
+            tracing::debug!(%reason, "forwarding a planned query");
+            Some(header_value(&reason.to_string()))
+        }
+    };
+    let mut response = forward(state, "query", method, query, headers, body).await;
+    if let Some(reason) = reason {
+        response.headers_mut().insert(FORWARD_REASON_HEADER, reason);
+    }
+    response
+}
+
+/// Runs a planned instant query. Bound raw-sample windows are read from Prometheus with
+/// the caller's end-to-end `headers`, as a forwarded query would be.
+async fn answer(
+    state: &DataPlane,
+    query: InstantQuery,
+    headers: &HeaderMap,
+) -> Result<serde_json::Value, Unanswered> {
+    let time_ms = query.time_ms.unwrap_or_else(|| {
+        std::time::SystemTime::now()
+            .duration_since(std::time::UNIX_EPOCH)
+            .map_or(0, |d| d.as_millis() as i64)
+    });
+    let (root, ranges) = state
+        .engine
+        .bound_raw_sample_windows(query.query, time_ms)
+        .await?;
+    let mut fetch_headers = forwarded_headers(headers);
+    // The raw read is a GET without the caller's body.
+    fetch_headers.remove(header::CONTENT_TYPE);
+    // The bound raw-sample windows are fetched concurrently.
+    let fetches = ranges.iter().map(|range| async {
+        prometheus::fetch(&state.client, &state.prometheus_url, &fetch_headers, range)
+            .await
+            .map(|samples| (range.node, samples))
+    });
+    let raw: BTreeMap<u64, Vec<RawSample>> = futures::future::try_join_all(fetches)
+        .await
+        .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))?
+        .into_iter()
+        .collect();
+    let samples = state.engine.evaluate(root, time_ms, raw).await?;
+    Ok(prometheus::vector_response(time_ms, &samples))
 }
 
 async fn query_range(
@@ -288,3 +372,19 @@ async fn forward(
         .insert(EXECUTION_HEADER, HeaderValue::from_static("forwarded"));
     response
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A forward reason with characters a header cannot carry is still sent,
+    /// with those characters replaced.
+    #[test]
+    fn sanitizes_header_values() {
+        assert_eq!(
+            header_value("pane ending at 10 is missing"),
+            "pane ending at 10 is missing"
+        );
+        assert_eq!(header_value("bad\nlabel é"), "bad?label ?");
+    }
+}