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

README.md

@@ -33,14 +33,18 @@ The backend is being rebuilt on the Planner's physical DAG:
 | Data plane: load and validate `plan.json`, accept Remote Write | done |
 | Stored summaries (SDS): one definition per computation, panes by end time | done |
 | Precompute engine: build ingestion-time outputs once per tumbling pane | done |
-| Query engine: run query-time nodes over stored panes and raw data | planned |
+| Query engine: answer planned instant queries from stored panes and raw samples | done |
 
 Panes close `--allowed-lateness-ms` (default 15 s) after the newest sample
 passes their end; later samples for them are dropped. The first pane after
 startup is skipped because it may be partial.
 
-Until the query engine lands, the data plane forwards every query to
-Prometheus and marks the response `x-asap-execution: forwarded`. SQL and
+An instant query whose text is one of the plan's queries is answered by the
+plan (`x-asap-execution: plan`). Its stored inputs read the panes ending at
+`T - offset`, so `T` must be on the pane grid and those panes stored; its
+query-time raw ranges are read from Prometheus. Anything else is forwarded
+to Prometheus (`x-asap-execution: forwarded`); a planned query that could not
+be answered also says why in `x-asap-forward-reason`. Range queries, SQL and
 ClickHouse are not supported yet.
 
 ## Running

data_plane/Cargo.toml

@@ -21,7 +21,9 @@ reqwest.workspace = true
 prost = "0.13"
 snap = "1.1"
 futures = "0.3"
+form_urlencoded = "1"
 
 [dev-dependencies]
 tower = { version = "0.5", features = ["util"] }
 http-body-util = "0.1"
+serde_json.workspace = true

data_plane/src/engine.rs

@@ -2,21 +2,36 @@
 //! it. Planner IR holds `Rc`s, so the plan cannot be shared across the HTTP
 //! server's threads; handlers send the engine commands instead.
 
+use std::collections::BTreeMap;
 use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TrySendError};
 
 use asap_executor::physical_planner::frontier_from_timing;
 use control_plane::InstalledPlan;
 use tokio::sync::oneshot;
 
 use crate::precompute::{Precompute, PrecomputeStats, RawSample};
+use crate::query::{QueryEngine, RawRange, ResultSample, Unanswered};
 use crate::sds::Sds;
 
 /// Commands queued for the engine before ingestion pushes back.
 const QUEUE: usize = 1024;
 
+type Reply<T> = oneshot::Sender<Result<T, Unanswered>>;
+
 pub enum Command {
     Ingest(Vec<RawSample>),
     Stats(oneshot::Sender<PrecomputeStats>),
+    RawRanges {
+        query: String,
+        time_ms: i64,
+        reply: Reply<(usize, Vec<RawRange>)>,
+    },
+    Evaluate {
+        root: usize,
+        time_ms: i64,
+        raw: BTreeMap<u64, Vec<RawSample>>,
+        reply: Reply<Vec<ResultSample>>,
+    },
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -58,32 +73,74 @@ impl EngineHandle {
         }
     }
 
+    /// The plan root answering `query` at `time_ms` and the raw samples it
+    /// needs; see [`QueryEngine::raw_ranges`].
+    pub async fn raw_ranges(
+        &self,
+        query: String,
+        time_ms: i64,
+    ) -> Result<(usize, Vec<RawRange>), 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<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 busy = || Unanswered::Execution("the engine is busy".into());
+        let (reply, answer) = oneshot::channel();
+        self.commands.try_send(command(reply)).map_err(|_| busy())?;
+        answer.await.map_err(|_| busy())?
+    }
+
     /// Precompute counters, once every command queued before has run.
     pub async fn stats(&self) -> PrecomputeStats {
         let (reply, stats) = oneshot::channel();
-        let _ = self.commands.send(Command::Stats(reply));
+        if self.commands.try_send(Command::Stats(reply)).is_err() {
+            return PrecomputeStats::default();
+        }
         stats.await.unwrap_or_default()
     }
 }
 
 struct Engine {
     precompute: Precompute,
+    query: QueryEngine,
 }
 
 impl Engine {
     fn new(plan_json: &str, allowed_lateness_ms: i64) -> anyhow::Result<Self> {
         let plan: InstalledPlan = serde_json::from_str(plan_json)?;
         plan.validate()?;
-        let dag = plan.dag.dag;
-        let frontier = frontier_from_timing(&dag)?;
-        let sds = Sds::from_dag(&dag, &frontier)?;
+        let frontier = frontier_from_timing(&plan.dag.dag)?;
+        let sds = Sds::from_dag(&plan.dag.dag, &frontier)?;
         tracing::info!(
             definitions = sds.definitions.len(),
             outputs = sds.outputs.len(),
             "stored summaries"
         );
+        let query = QueryEngine::new(&plan, &sds);
         Ok(Self {
-            precompute: Precompute::new(dag, &sds, allowed_lateness_ms)?,
+            precompute: Precompute::new(plan.dag.dag, &sds, allowed_lateness_ms)?,
+            query,
         })
     }
 
@@ -94,6 +151,30 @@ impl Engine {
                 Command::Stats(reply) => {
                     let _ = reply.send(self.precompute.stats);
                 }
+                Command::RawRanges {
+                    query,
+                    time_ms,
+                    reply,
+                } => {
+                    let _ = reply.send(self.query.raw_ranges(
+                        &query,
+                        time_ms,
+                        &self.precompute.store,
+                    ));
+                }
+                Command::Evaluate {
+                    root,
+                    time_ms,
+                    raw,
+                    reply,
+                } => {
+                    let _ = reply.send(self.query.evaluate(
+                        root,
+                        time_ms,
+                        &self.precompute.store,
+                        &raw,
+                    ));
+                }
             }
         }
     }

data_plane/src/lib.rs

@@ -2,11 +2,14 @@
 //! the Prometheus query API for an installed plan.
 //!
 //! The installed plan is ASAPPlanner's selected `PhysicalASAPDAG`. The
-//! engine thread builds its ingestion-time outputs pane by pane; queries are
-//! still forwarded to Prometheus.
+//! engine thread builds its ingestion-time outputs pane by pane and answers
+//! the plan's instant queries from them; every other query, and any planned
+//! query it cannot answer, is forwarded to Prometheus.
 
 pub mod engine;
 pub mod precompute;
+pub mod prometheus;
+pub mod query;
 pub mod remote_write;
 pub mod sds;
 
@@ -23,12 +26,16 @@ use axum::Router;
 use control_plane::InstalledPlan;
 use engine::EngineHandle;
 use precompute::RawSample;
+use prometheus::InstantQuery;
+use query::Unanswered;
 
 /// Prometheus' staleness marker: a NaN that ends a series, not a sample.
 const STALE_NAN_BITS: u64 = 0x7ff0_0000_0000_0002;
 
-/// Response header naming what answered a query.
+/// Response header naming what answered a query: `plan` or `forwarded`.
 pub const EXECUTION_HEADER: &str = "x-asap-execution";
+/// Response header of a forwarded planned query: why the plan did not answer.
+pub const FORWARD_REASON_HEADER: &str = "x-asap-forward-reason";
 
 pub struct DataPlane {
     /// The installed plan's queries, in root order.
@@ -110,14 +117,54 @@ async fn write(State(state): State<Arc<DataPlane>>, body: Bytes) -> Response {
     }
 }
 
+/// 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 Some(instant) = InstantQuery::parse(query.0.as_deref(), &body) else {
+        return forward(state, "query", method, query, headers, body).await;
+    };
+    let reason = match answer(&state, instant).await {
+        Ok(response) => {
+            let mut response = axum::Json(response).into_response();
+            response
+                .headers_mut()
+                .insert(EXECUTION_HEADER, HeaderValue::from_static("plan"));
+            return response;
+        }
+        Err(Unanswered::NotPlanned(_)) => None,
+        Err(reason) => {
+            tracing::debug!(%reason, "forwarding a planned query");
+            HeaderValue::from_str(&reason.to_string()).ok()
+        }
+    };
+    let mut response = forward(state, "query", method, query, headers, body).await;
+    if let Some(reason) = reason {
+        response.headers_mut().insert(FORWARD_REASON_HEADER, reason);
+    }
+    response
+}
+
+async fn answer(state: &DataPlane, query: InstantQuery) -> Result<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.raw_ranges(query.query, time_ms).await?;
+    let mut raw = BTreeMap::new();
+    for range in &ranges {
+        let samples = prometheus::fetch(&state.client, &state.prometheus_url, range)
+            .await
+            .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))?;
+        raw.insert(range.node, samples);
+    }
+    let samples = state.engine.evaluate(root, time_ms, raw).await?;
+    Ok(prometheus::vector_response(time_ms, &samples))
 }
 
 async fn query_range(