feat(data-plane): --no-forward: every answer computed by ASAP

README.md

@@ -106,6 +106,16 @@ with cost model 1 and illustrative unit costs.
 `--forward-timeout-seconds` bounds each forwarded query (default 120, matching
 Prometheus's default query timeout); a query that runs longer gets 502.
 
+`--no-forward` makes every answer ASAP's: the data plane never sends a query
+to Prometheus for computation. It refuses to start with a plan that has a
+query it cannot run, and answers any query the plan does not answer (not
+planned, a pane not stored, an evaluation time off the pane grid, a range
+query) with a Prometheus `execution` error (422, `x-asap-execution: refused`)
+whose `error` is the reason a forwarded query's `x-asap-forward-reason` would
+give. Raw samples for query-time windows are still read from Prometheus. The
+control plane writes a plan only when every workload query has one the data
+plane can execute; otherwise it fails.
+
 ## Development
 
 ```sh

control_plane/tests/plan_workload.rs

@@ -106,3 +106,20 @@ fn selects_stored_outputs_built_in_a_query_run() {
     assert_eq!(query_run.len(), 5, "five one-minute summaries");
     assert!(query_run.iter().all(|o| o.pane.width_ms == 60_000));
 }
+
+/// A workload is planned whole or refused: when one query has no plan the
+/// data plane can execute (here a raw instant selector, without raw data at
+/// query time), no plan is written, so no query is left for Prometheus to
+/// compute.
+#[test]
+fn refuses_a_workload_with_a_query_it_cannot_plan() {
+    let mut workload = example("example4b");
+    let queries = workload.query_workload.repeating_queries.as_mut().unwrap();
+    let mut unplannable = queries[0].clone();
+    unplannable.query = serde_json::from_value(serde_json::json!("latency_ms")).unwrap();
+    queries.push(unplannable);
+    let mut inputs = DeploymentInputs::default();
+    inputs.deployment_capabilities.raw_data_retained = false;
+    let error = plan_workload(&workload, &inputs).unwrap_err();
+    assert!(matches!(error, PlanError::Planning(_)), "{error}");
+}

data_plane/src/engine.rs

@@ -26,6 +26,10 @@ pub struct EngineOptions {
     /// Samples queued for the engine, not yet ingested, beyond which Remote
     /// Write is refused with 503 so senders retry.
     pub max_queued_samples: u64,
+    /// Every answer must be computed from the plan: a plan with a query the
+    /// query engine cannot run is refused at install, and the data plane
+    /// never forwards a query to Prometheus (`--no-forward`).
+    pub no_forward: bool,
 }
 
 impl Default for EngineOptions {
@@ -34,6 +38,7 @@ impl Default for EngineOptions {
             precompute: PrecomputeOptions::default(),
             // Two Remote Write requests of the default sample limit.
             max_queued_samples: 2_000_000,
+            no_forward: false,
         }
     }
 }
@@ -102,7 +107,7 @@ impl EngineHandle {
     /// plan is invalid or the data plane cannot execute it.
     pub fn spawn(plan_json: String, options: EngineOptions) -> anyhow::Result<(Self, Vec<String>)> {
         let (commands, receiver) = channel();
-        let (engine, queries) = Engine::new(&plan_json, options.precompute)?;
+        let (engine, queries) = Engine::new(&plan_json, options)?;
         let running = Arc::new(AtomicBool::new(true));
         let queued_samples = Arc::new(AtomicU64::new(0));
         let stats = Arc::new(Mutex::new(PrecomputeStats::default()));
@@ -270,7 +275,7 @@ struct Engine {
 }
 
 impl Engine {
-    fn new(plan_json: &str, options: PrecomputeOptions) -> anyhow::Result<(Self, Vec<String>)> {
+    fn new(plan_json: &str, options: EngineOptions) -> anyhow::Result<(Self, Vec<String>)> {
         let plan: InstallablePlan = serde_json::from_str(plan_json)?;
         plan.validate()?;
         let sds = Sds::from_plan(&plan)?;
@@ -281,8 +286,19 @@ impl Engine {
         );
         let queries = plan.queries.iter().map(|q| q.query.clone()).collect();
         let query = QueryEngine::new(&plan, &sds);
+        let unexecutable = query.unexecutable();
+        if options.no_forward && !unexecutable.is_empty() {
+            let reasons: Vec<String> = unexecutable
+                .iter()
+                .map(|(query, reason)| format!("{query:?}: {reason}"))
+                .collect();
+            anyhow::bail!(
+                "--no-forward refuses a plan with queries the data plane cannot run: {}",
+                reasons.join("; ")
+            );
+        }
         let engine = Self {
-            precompute: Precompute::new(plan.dag.dag, &sds, options)?,
+            precompute: Precompute::new(plan.dag.dag, &sds, options.precompute)?,
             query,
         };
         Ok((engine, queries))

data_plane/src/lib.rs

@@ -32,7 +32,8 @@ use precompute::RawSample;
 use prometheus::InstantQuery;
 use query::Unanswered;
 
-/// Response header naming what answered a query: `plan` or `forwarded`.
+/// Response header naming what answered a query: `plan` or `forwarded`, or
+/// `refused` when `--no-forward` refused to forward it.
 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";
@@ -50,6 +51,9 @@ pub struct DataPlane {
     pub limits: remote_write::Limits,
     /// Samples accepted since startup.
     pub samples: AtomicU64,
+    /// Never forward a query to Prometheus: a query the plan cannot answer
+    /// gets an error with the reason ([`EngineOptions::no_forward`]).
+    pub no_forward: bool,
     client: reqwest::Client,
 }
 
@@ -76,6 +80,7 @@ impl DataPlane {
     ) -> anyhow::Result<Self> {
         let (engine, queries) = EngineHandle::spawn(plan_json, options)?;
         Ok(Self {
+            no_forward: options.no_forward,
             queries,
             engine,
             prometheus_url: prometheus_url.trim_end_matches('/').to_string(),
@@ -262,6 +267,9 @@ async fn query(
 ) -> Response {
     let instant = match InstantQuery::parse(query.0.as_deref(), &body) {
         Ok(Some(instant)) => instant,
+        Ok(None) if state.no_forward => {
+            return not_answered("the request has no query parameter");
+        }
         Ok(None) => return forward(state, "query", method, query, headers, body).await,
         Err(bad) => {
             let body = axum::Json(prometheus::bad_data(&bad.to_string()));
@@ -276,6 +284,7 @@ async fn query(
                 .insert(EXECUTION_HEADER, HeaderValue::from_static("plan"));
             return response;
         }
+        Err(reason) if state.no_forward => return not_answered(&reason.to_string()),
         Err(Unanswered::NotPlanned(_)) => None,
         Err(reason) => {
             tracing::debug!(%reason, "forwarding a planned query");
@@ -331,9 +340,25 @@ async fn query_range(
     headers: HeaderMap,
     body: Bytes,
 ) -> Response {
+    if state.no_forward {
+        return not_answered("the data plane does not answer range queries yet");
+    }
     forward(state, "query_range", method, query, headers, body).await
 }
 
+/// With `--no-forward`, the answer to a query the plan does not answer: a
+/// Prometheus `execution` error (422) with `reason`, the reason a
+/// forwarded query's `x-asap-forward-reason` would carry, and
+/// `x-asap-execution: refused`. Prometheus receives no query.
+fn not_answered(reason: &str) -> Response {
+    let body = serde_json::json!({"status": "error", "errorType": "execution", "error": reason});
+    let mut response = (StatusCode::UNPROCESSABLE_ENTITY, axum::Json(body)).into_response();
+    response
+        .headers_mut()
+        .insert(EXECUTION_HEADER, HeaderValue::from_static("refused"));
+    response
+}
+
 /// Request headers not forwarded: hop-by-hop headers (RFC 9110 section 7.6.1),
 /// `Host` and `Content-Length`, which the client sets for the new request,
 /// and `Accept-Encoding`, because only the response's `Content-Type` is

data_plane/src/main.rs

@@ -30,6 +30,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,
+    /// Never forward a query to Prometheus: refuse a plan with a query the
+    /// data plane cannot run, and answer any query the plan does not answer
+    /// with an error giving the reason. Raw samples are still read from
+    /// Prometheus.
+    #[arg(long)]
+    no_forward: bool,
 }
 
 #[tokio::main]
@@ -49,6 +55,7 @@ async fn main() -> anyhow::Result<()> {
                     max_future_ms: args.max_future_ms,
                 },
                 max_queued_samples: args.max_queued_samples,
+                no_forward: args.no_forward,
             },
         )?
         .with_forward_timeout(Duration::from_secs(args.forward_timeout_seconds)),

data_plane/src/precompute.rs

@@ -634,7 +634,14 @@ mod tests {
             ..PrecomputeOptions::default()
         };
         let mut engine = Precompute::new(plan.dag.dag, &sds, options).unwrap();
-        assert_eq!((sds.definitions.len(), engine.builders.len(), engine.per_sample_definitions()), (1, 1, 1));
+        assert_eq!(
+            (
+                sds.definitions.len(),
+                engine.builders.len(),
+                engine.per_sample_definitions()
+            ),
+            (1, 1, 1)
+        );
         let jobs: Vec<String> = (0..10).map(|i| format!("job{i}")).collect();
         let mut samples = Vec::new();
         for t in (MINUTE..25 * MINUTE).step_by(15_000) {

data_plane/src/query.rs

@@ -407,6 +407,14 @@ impl QueryEngine {
         result_samples(&program.output, &batches)
     }
 
+    /// The planned queries the engine cannot run, with why.
+    pub fn unexecutable(&self) -> Vec<(String, String)> {
+        self.roots
+            .iter()
+            .filter_map(|(query, program)| Some((query.clone(), program.as_ref().err()?.clone())))
+            .collect()
+    }
+
     /// The build of a `QueryRun` definition a compiled root reads.
     fn query_run_build(&self, definition: &DefinitionId) -> &SummaryBuild {
         self.query_run_builds[definition]

data_plane/tests/no_forward.rs

@@ -0,0 +1,156 @@
+//! `--no-forward`: every answer is computed from the plan. A query the plan
+//! cannot answer gets an error with the reason, and Prometheus receives no
+//! query for it.
+
+mod common;
+
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+use asap_executor::planner::ir::{NonASAPOp, Operator, Predicate, ScalarExpr};
+use axum::body::Body;
+use axum::http::{Request, StatusCode};
+use axum::Router;
+use data_plane::engine::EngineOptions;
+use data_plane::{router, DataPlane, EXECUTION_HEADER};
+use http_body_util::BodyExt;
+use serde_json::Value;
+use tower::ServiceExt;
+
+const MINUTE: i64 = 60_000;
+
+/// A stand-in Prometheus that counts every request it receives.
+async fn counting_prometheus() -> (String, Arc<AtomicUsize>) {
+    let requests = Arc::new(AtomicUsize::new(0));
+    let counter = requests.clone();
+    let app = Router::new().fallback(move || {
+        let counter = counter.clone();
+        async move {
+            counter.fetch_add(1, Ordering::SeqCst);
+            "{}"
+        }
+    });
+    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, requests)
+}
+
+fn options(no_forward: bool) -> EngineOptions {
+    EngineOptions {
+        no_forward,
+        ..EngineOptions::default()
+    }
+}
+
+async fn get(app: &Router, uri: &str) -> (StatusCode, Option<String>, Value) {
+    let response = app
+        .clone()
+        .oneshot(Request::get(uri).body(Body::empty()).unwrap())
+        .await
+        .unwrap();
+    let status = response.status();
+    let execution = response
+        .headers()
+        .get(EXECUTION_HEADER)
+        .map(|v| v.to_str().unwrap().to_string());
+    let body = response.into_body().collect().await.unwrap().to_bytes();
+    (
+        status,
+        execution,
+        serde_json::from_slice(&body).unwrap_or(Value::Null),
+    )
+}
+
+fn instant(query: &str, time_ms: i64) -> String {
+    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()
+    )
+}
+
+/// With the flag, an unplanned query, an evaluation whose pane is missing
+/// and a range query each get a Prometheus `execution` error with the
+/// reason, and the stand-in Prometheus receives no request.
+#[tokio::test]
+async fn refuses_what_the_plan_cannot_answer_without_asking_prometheus() {
+    let (url, requests) = counting_prometheus().await;
+    let state =
+        DataPlane::with_options(common::plan_json("example4b", false), url, options(true)).unwrap();
+    let app = router(Arc::new(state));
+
+    let (status, execution, body) = get(&app, &instant("up", 70 * MINUTE)).await;
+    assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
+    assert_eq!(execution.as_deref(), Some("refused"));
+    assert_eq!(body["status"], "error");
+    assert_eq!(body["errorType"], "execution");
+    assert!(
+        body["error"].as_str().unwrap().contains("has no query"),
+        "{body}"
+    );
+
+    // Nothing was written, so the panes the query reads are not stored.
+    let query = "quantile_over_time(0.99, latency_ms[60m])";
+    let (status, execution, body) = get(&app, &instant(query, 70 * MINUTE)).await;
+    assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
+    assert_eq!(execution.as_deref(), Some("refused"));
+    assert!(
+        body["error"].as_str().unwrap().contains("not stored"),
+        "{body}"
+    );
+
+    let range = format!(
+        "/api/v1/query_range?query={}&start=0&end=60&step=15",
+        form_urlencoded::byte_serialize(query.as_bytes()).collect::<String>()
+    );
+    let (status, execution, _) = get(&app, &range).await;
+    assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY);
+    assert_eq!(execution.as_deref(), Some("refused"));
+
+    assert_eq!(requests.load(Ordering::SeqCst), 0, "Prometheus got a query");
+}
+
+/// Without the flag the same unplanned query is forwarded, as before.
+#[tokio::test]
+async fn forwards_without_the_flag() {
+    let (url, requests) = counting_prometheus().await;
+    let state = DataPlane::with_options(common::plan_json("example4b", false), url, options(false))
+        .unwrap();
+    let app = router(Arc::new(state));
+    let (status, execution, _) = get(&app, &instant("up", 70 * MINUTE)).await;
+    assert_eq!(status, StatusCode::OK);
+    assert_eq!(execution.as_deref(), Some("forwarded"));
+    assert_eq!(requests.load(Ordering::SeqCst), 1);
+}
+
+/// With the flag, a plan with a query the data plane cannot run is refused
+/// at install; without it, the plan installs and that query is forwarded.
+#[test]
+fn refuses_to_install_a_plan_with_a_query_it_cannot_run() {
+    let mut plan = common::plan("example1", true);
+    // A label matcher on a column the `Scan` does not have: the selector of
+    // its raw-sample window does not compile, so the queries that read it
+    // cannot run.
+    let scan = plan.query_time_raw_windows[0].scan_node_id;
+    let node = plan
+        .dag
+        .dag
+        .nodes
+        .iter_mut()
+        .find(|n| n.id == scan)
+        .unwrap();
+    let Operator::NonASAP(NonASAPOp::Scan { predicates, .. }) = &mut node.payload else {
+        panic!("not a Scan");
+    };
+    predicates.push(Predicate(ScalarExpr::Column(99)));
+    plan.validate().unwrap();
+    let plan_json = serde_json::to_string(&plan).unwrap();
+    let error = DataPlane::with_options(plan_json.clone(), "http://unused".into(), options(true))
+        .err()
+        .expect("refused");
+    assert!(error.to_string().contains("--no-forward"), "{error}");
+    assert!(DataPlane::with_options(plan_json, "http://unused".into(), options(false)).is_ok());
+}