feat(precompute): build stored summary panes from Remote Write samples

Cargo.toml

@@ -9,11 +9,11 @@ version = "0.1.0"
 [workspace.dependencies]
 # Keep every Planner crate on the same immutable revision. Alias upstream
 # asap-types as planner-types to name its origin at use sites.
-planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "f9dd41e095997fe2c34c1a1278f9f20c50a814ff" }
-asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "f9dd41e095997fe2c34c1a1278f9f20c50a814ff" }
-asap-logical-asap-optimization = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "f9dd41e095997fe2c34c1a1278f9f20c50a814ff" }
-asap-plan-selection = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "f9dd41e095997fe2c34c1a1278f9f20c50a814ff" }
-asap-executor = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "f9dd41e095997fe2c34c1a1278f9f20c50a814ff" }
+planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "ecea897f78640c377ce1d80846156f330f62dee2" }
+asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "ecea897f78640c377ce1d80846156f330f62dee2" }
+asap-logical-asap-optimization = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "ecea897f78640c377ce1d80846156f330f62dee2" }
+asap-plan-selection = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "ecea897f78640c377ce1d80846156f330f62dee2" }
+asap-executor = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "ecea897f78640c377ce1d80846156f330f62dee2" }
 
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"

README.md

@@ -31,18 +31,25 @@ The backend is being rebuilt on the Planner's physical DAG:
 | Part | State |
 |---|---|
 | Control plane: plan a PromQL workload, write `plan.json` | done |
-| Data plane: load and validate `plan.json`; accept, count and drop Remote Write samples | done |
-| Stored summaries (SDS) | next |
-| Precompute engine: run ingestion-time nodes per pane | planned |
+| 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 |
 
-Until the precompute engine lands, the data plane validates and counts Remote
-Write samples, then drops them: Prometheus must receive Remote Write itself
-(for example, configure the sender with both endpoints). 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 ClickHouse are not supported
-yet.
+Panes close `--allowed-lateness-ms` (default 15 s) after the newest sample
+of a stored summary's metric, capped at the wall clock, passes their end; later samples for them are
+dropped. The first pane after startup is skipped because it may be partial.
+Samples more than `--max-future-ms` (default 5 min) ahead of the wall clock
+are dropped. Remote Write gets 503 when admitting it would leave more than
+`--max-queued-samples` (default 2,000,000) samples waiting for the engine
+(a request into an empty queue is always admitted), and 500 if the engine has
+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
+ClickHouse are not supported yet.
 
 ## Running
 
@@ -51,8 +58,7 @@ yet.
 cargo run -p control_plane -- \
   --workload examples/planner-layering-example1.json --out plan.json
 
-# Serve Remote Write on /api/v1/write (samples are counted and dropped for now)
-# and the query API on /api/v1/query[_range], forwarded to Prometheus.
+# Serve Remote Write on /api/v1/write and the query API on /api/v1/query[_range].
 cargo run -p data_plane -- --plan plan.json --prometheus-url http://localhost:9090
 ```
 

data_plane/Cargo.toml

@@ -20,9 +20,10 @@ axum.workspace = true
 reqwest.workspace = true
 prost = "0.13"
 snap = "1.1"
+futures = "0.3"
 
 [dev-dependencies]
 tower = { version = "0.5", features = ["util"] }
 http-body-util = "0.1"
-# The tests find the frontier the way the Planner does; the runtime does not.
+# The tests check the SDS's frontier against the executor's own.
 asap-executor = { workspace = true, features = ["test-support"] }

data_plane/src/engine.rs

@@ -0,0 +1,247 @@
+//! 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.
+
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+use std::sync::mpsc::{channel, Receiver, Sender};
+use std::sync::{Arc, Mutex};
+
+use control_plane::InstalledPlan;
+
+use crate::precompute::{Precompute, PrecomputeOptions, PrecomputeStats, RawSample};
+use crate::sds::Sds;
+
+/// How the engine thread 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.
+    pub max_queued_samples: u64,
+}
+
+impl Default for EngineOptions {
+    fn default() -> Self {
+        Self {
+            precompute: PrecomputeOptions::default(),
+            // Two Remote Write requests of the default sample limit.
+            max_queued_samples: 2_000_000,
+        }
+    }
+}
+
+enum Command {
+    Ingest(Vec<RawSample>),
+    Stop,
+}
+
+/// Why the engine did not take samples.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
+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.
+    #[error("the engine has stopped")]
+    Stopped,
+}
+
+/// The HTTP side of the engine thread.
+#[derive(Clone)]
+pub struct EngineHandle {
+    commands: Sender<Command>,
+    /// Samples sent and not yet ingested.
+    queued_samples: Arc<AtomicU64>,
+    max_queued_samples: u64,
+    /// False once the engine thread has exited, for any reason.
+    running: Arc<AtomicBool>,
+    /// The counters as of the last command the engine ran.
+    stats: Arc<Mutex<PrecomputeStats>>,
+}
+
+/// Marks the engine stopped when its thread exits, panics included.
+struct RunningGuard(Arc<AtomicBool>);
+
+impl Drop for RunningGuard {
+    fn drop(&mut self) {
+        self.0.store(false, Ordering::SeqCst);
+    }
+}
+
+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.
+    pub fn spawn(plan_json: String, options: EngineOptions) -> anyhow::Result<(Self, Vec<String>)> {
+        let (commands, receiver) = channel();
+        let (ready, started) = std::sync::mpsc::channel();
+        let running = Arc::new(AtomicBool::new(true));
+        let queued_samples = Arc::new(AtomicU64::new(0));
+        let stats = Arc::new(Mutex::new(PrecomputeStats::default()));
+        let shared = Shared {
+            queued_samples: queued_samples.clone(),
+            stats: 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 handle = Self {
+            commands,
+            queued_samples,
+            max_queued_samples: options.max_queued_samples,
+            running,
+            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.
+    pub fn try_ingest(&self, samples: Vec<RawSample>) -> Result<(), IngestError> {
+        if !self.is_running() {
+            return Err(IngestError::Stopped);
+        }
+        let count = samples.len() as u64;
+        let mut queued = self.queued_samples.load(Ordering::SeqCst);
+        loop {
+            let total = admit(queued, count, self.max_queued_samples)?;
+            match self.queued_samples.compare_exchange(
+                queued,
+                total,
+                Ordering::SeqCst,
+                Ordering::SeqCst,
+            ) {
+                Ok(_) => break,
+                Err(current) => queued = current,
+            }
+        }
+        self.commands.send(Command::Ingest(samples)).map_err(|_| {
+            self.queued_samples.fetch_sub(count, Ordering::SeqCst);
+            IngestError::Stopped
+        })
+    }
+
+    /// Whether the engine thread is still running.
+    pub fn is_running(&self) -> bool {
+        self.running.load(Ordering::SeqCst)
+    }
+
+    /// Samples queued and not yet ingested.
+    pub fn queued_samples(&self) -> u64 {
+        self.queued_samples.load(Ordering::SeqCst)
+    }
+
+    /// The precompute counters as of the last batch the engine ingested.
+    pub fn stats(&self) -> PrecomputeStats {
+        *self.stats.lock().unwrap_or_else(|e| e.into_inner())
+    }
+
+    /// 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
+    /// [`IngestError::Stopped`].
+    #[doc(hidden)]
+    pub fn stop(&self) {
+        let _ = self.commands.send(Command::Stop);
+    }
+}
+
+/// The queue length after admitting `count` samples to a queue of
+/// `queued`: admitted while the total stays within `max_queued`, and always
+/// into an empty queue, so a bound smaller than one request cannot refuse
+/// that request forever.
+fn admit(queued: u64, count: u64, max_queued: u64) -> Result<u64, IngestError> {
+    let total = queued.checked_add(count).ok_or(IngestError::Busy)?;
+    if queued == 0 || total <= max_queued {
+        Ok(total)
+    } else {
+        Err(IngestError::Busy)
+    }
+}
+
+/// What the engine thread publishes to the HTTP side.
+struct Shared {
+    queued_samples: Arc<AtomicU64>,
+    stats: Arc<Mutex<PrecomputeStats>>,
+}
+
+struct Engine {
+    precompute: Precompute,
+}
+
+impl Engine {
+    fn new(plan_json: &str, options: PrecomputeOptions) -> anyhow::Result<(Self, Vec<String>)> {
+        let plan: InstalledPlan = serde_json::from_str(plan_json)?;
+        plan.validate()?;
+        let dag = plan.dag.dag;
+        let frontier = Sds::ingestion_frontier(&dag)?;
+        let sds = Sds::from_dag(&dag, &frontier)?;
+        tracing::info!(
+            definitions = sds.definitions.len(),
+            outputs = sds.outputs.len(),
+            "stored summaries"
+        );
+        let engine = Self {
+            precompute: Precompute::new(dag, &sds, options)?,
+        };
+        Ok((engine, plan.queries))
+    }
+
+    fn serve(mut self, commands: Receiver<Command>, 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::Stop => break,
+            }
+        }
+        tracing::warn!("engine stopped");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A request is admitted while the queue stays within its bound, and
+    /// always into an empty queue; otherwise it is `Busy`.
+    #[test]
+    fn admission() {
+        assert_eq!(admit(0, 3, 2), Ok(3));
+        assert_eq!(admit(1, 1, 2), Ok(2));
+        assert_eq!(admit(1, 2, 2), Err(IngestError::Busy));
+        assert_eq!(admit(u64::MAX, 1, u64::MAX), Err(IngestError::Busy));
+    }
+
+    /// The running flag clears when the engine thread panics.
+    #[test]
+    fn running_guard_clears_on_panic() {
+        let running = Arc::new(AtomicBool::new(true));
+        let guard = RunningGuard(running.clone());
+        let thread = std::thread::spawn(move || {
+            let _guard = guard;
+            panic!("engine bug");
+        });
+        assert!(thread.join().is_err());
+        assert!(!running.load(Ordering::SeqCst));
+    }
+}