diff --git a/Cargo.lock b/Cargo.lock index 303579a0..a218ee69 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1043,6 +1043,7 @@ dependencies = [ "axum", "clap", "control_plane", + "form_urlencoded", "futures", "http-body-util", "prost", @@ -1052,6 +1053,7 @@ dependencies = [ "sha2 0.10.9", "snap", "thiserror 1.0.69", + "tikv-jemallocator", "tokio", "tower", "tracing", @@ -1097,7 +1099,7 @@ dependencies = [ "flate2", "futures", "indexmap", - "itertools 0.14.0", + "itertools", "liblzma", "log", "object_store", @@ -1129,7 +1131,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", "parking_lot", @@ -1154,7 +1156,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", ] @@ -1173,7 +1175,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "libc", "log", "object_store", @@ -1220,7 +1222,7 @@ dependencies = [ "flate2", "futures", "glob", - "itertools 0.14.0", + "itertools", "liblzma", "log", "object_store", @@ -1251,7 +1253,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "object_store", "tokio", ] @@ -1325,7 +1327,7 @@ dependencies = [ "datafusion-pruning", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", "parking_lot", @@ -1378,7 +1380,7 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr-common", "indexmap", - "itertools 0.14.0", + "itertools", "recursive", "serde_json", "sqlparser", @@ -1393,7 +1395,7 @@ dependencies = [ "arrow", "datafusion-common", "indexmap", - "itertools 0.14.0", + "itertools", ] [[package]] @@ -1417,7 +1419,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hex", - "itertools 0.14.0", + "itertools", "log", "md-5", "memchr", @@ -1480,7 +1482,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hashbrown 0.17.1", - "itertools 0.14.0", + "itertools", "itoa", "log", "memchr", @@ -1553,7 +1555,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools 0.14.0", + "itertools", "log", "recursive", "regex", @@ -1575,7 +1577,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "parking_lot", "petgraph", "recursive", @@ -1594,7 +1596,7 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools 0.14.0", + "itertools", ] [[package]] @@ -1609,7 +1611,7 @@ dependencies = [ "datafusion-expr-common", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "parking_lot", "pin-project", ] @@ -1629,7 +1631,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", - "itertools 0.14.0", + "itertools", "recursive", ] @@ -1658,7 +1660,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "log", "num-traits", "parking_lot", @@ -2339,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" @@ -2788,7 +2781,7 @@ dependencies = [ "futures-util", "http", "humantime", - "itertools 0.14.0", + "itertools", "parking_lot", "percent-encoding", "thiserror 2.0.20", @@ -3069,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", @@ -3386,7 +3379,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3786,7 +3779,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3849,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" @@ -4333,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 4794bf52..a7d4233c 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,41 @@ 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 +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 +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 +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 ClickHouse are not supported yet. ## Running @@ -71,8 +103,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/Cargo.toml b/data_plane/Cargo.toml index 5914ecd0..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 @@ -20,8 +21,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..d035da37 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -1,66 +1,115 @@ -//! 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}; use std::sync::mpsc::{channel, Receiver, Sender}; use std::sync::{Arc, Mutex}; use asap_planner::InstallablePlan; -use crate::precompute::{Precompute, PrecomputeOptions, PrecomputeStats, RawSample}; +use tokio::sync::oneshot; + +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, } } } -enum Command { +type Reply = oneshot::Sender>; + +enum WorkerCommand { Ingest(Vec), + /// Reply once every command queued before has run. + Flush(oneshot::Sender<()>), Stop, } +enum QueryJob { + RawRanges { + query: String, + time_ms: i64, + reply: Reply, + }, + Evaluate { + root: usize, + time_ms: i64, + progress: IngestionProgress, + raw: BTreeMap>, + reply: Reply>, + }, +} + /// 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. + /// 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>, } -/// 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 { @@ -69,48 +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 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 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, }; 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); @@ -129,13 +212,86 @@ 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, once the writes acknowledged before are ingested; see + /// [`QueryEngine::raw_ranges`]. + pub async fn raw_ranges(&self, query: String, time_ms: i64) -> Result { + self.sync().await; + self.ask(|reply| QueryJob::RawRanges { + query, + time_ms, + reply, + }) + .await + } + + /// Runs a prepared root over stored summaries and `raw`; see + /// [`QueryEngine::evaluate`]. + pub async fn evaluate( + &self, + prepared: Prepared, + time_ms: i64, + raw: BTreeMap>, + ) -> Result, Unanswered> { + self.ask(|reply| QueryJob::Evaluate { + root: prepared.root, + time_ms, + progress: prepared.progress, + raw, + reply, }) + .await + } + + /// The query-run summary counters as of now. + pub fn kept_stats(&self) -> KeptResultStats { + self.query.kept_stats() } - /// Whether the engine thread is still running. + 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.queries.send(job(reply)).map_err(|_| stopped())?; + answer.await.map_err(|_| stopped())? + } + + /// 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) } @@ -145,18 +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() + } + + /// 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 [`Shared::advance_watermark_at`]. + pub async fn advance_watermark(&self, barrier_ms: i64) -> Result { + if !self.is_running() { + return 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); + } } } @@ -173,47 +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>, -} - -struct Engine { - precompute: Precompute, +/// 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 engine = Self { - precompute: Precompute::new(plan.dag.dag, &sds, options)?, +/// 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::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 fe037eb7..379fc313 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -1,12 +1,16 @@ //! 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; queries are -//! still 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; +pub mod prometheus; +pub mod query; pub mod remote_write; pub mod sds; @@ -23,9 +27,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. @@ -99,6 +107,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) } @@ -152,6 +161,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 { @@ -169,6 +226,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")], @@ -177,6 +235,23 @@ 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>, method: Method, @@ -184,7 +259,65 @@ async fn query( 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. 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 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); + // 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(prepared, time_ms, raw).await?; + Ok(prometheus::vector_response(time_ms, &samples)) } async fn query_range( @@ -288,3 +421,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/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 9658a4fe..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}; @@ -29,6 +33,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 +74,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,20 +105,95 @@ 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, - /// 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, +} + +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)?) + } +} + +/// 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. @@ -119,173 +204,358 @@ enum Outcome { Buffered, } -pub struct Precompute { - builders: Vec, +/// 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, - 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 { - 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))); + 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 = BTreeSet::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()); + } + Err(error) => tracing::warn!( + definition = ?definition.id, %error, + "query-run summary cannot be built; its queries are forwarded" + ), + } } - 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, - open: BTreeMap::new(), - first_pane_end: None, - built_through: None, - }); } 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, - store: SummaryStore::new(), - stats: PrecomputeStats::default(), }) } +} - /// 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. - 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(newest.saturating_sub(self.allowed_lateness_ms)) +/// 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); } +} - /// [`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); +/// 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, + } } - /// 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); - if !self - .builders + /// 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. + pub fn progress(&self) -> crate::query::IngestionProgress { + crate::query::IngestionProgress { + watermark_ms: self.watermark_ms(), + newest_sample_ms: self + .query_run_newest_ms .iter() - .any(|b| metric == Some(b.selector.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.selector.metric()) { + .filter_map(|(metric, newest)| Some((metric.clone(), get(newest)?))) + .collect(), + } + } + + 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 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); + let before = self.barrier_ms.fetch_max(barrier_ms, Ordering::AcqRel); + if before != UNSET && before >= barrier_ms { + return before; + } + self.close_panes(); + barrier_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(); + } + + 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.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 - } - } - }; - 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, + 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; } - self.newest_sample_ms = self.newest_sample_ms.max(Some(sample.timestamp_ms)); + 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); + } + } + 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.definition.pane_ms; - let id = &builder.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 + 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 { @@ -295,29 +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(); - match build_pane(builder, end, &samples) { + 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 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, @@ -326,9 +609,108 @@ 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. +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)) +} + /// 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 @@ -341,39 +723,17 @@ 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, - }; - // The selector takes owned label maps, so each sample's labels are - // copied once per pane build. - let samples: Vec = samples +/// `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(); - let selected = builder.selector.window_batch(end, &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)?) + .collect() } /// Executes `root` of `dag` to completion and returns its rows. diff --git a/data_plane/src/prometheus.rs b/data_plane/src/prometheus.rs new file mode 100644 index 00000000..2e18a871 --- /dev/null +++ b/data_plane/src/prometheus.rs @@ -0,0 +1,267 @@ +//! 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, +} + +/// 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). `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)) + .map(|(k, v)| (k.into_owned(), v.into_owned())) + .collect(); + let time_ms = match params.get("time") { + None => None, + 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" + )) + })?), + }, + }; + 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 +} + +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 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 +/// `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 = range_selector(range); + 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() + .unwrap(); + assert_eq!((q.query.as_str(), q.time_ms), ("up", Some(1500))); + 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"").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. + #[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..43ce6cbc --- /dev/null +++ b/data_plane/src/query.rs @@ -0,0 +1,662 @@ +//! 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 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. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::atomic::{AtomicU64, Ordering}; +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, 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, StoredOutputBuild}; + +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. +#[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 { + 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, + matchers: Vec, + /// The plan's `Scan` label matchers over the range's samples. + selector: RawSampleSelector, + schema: SchemaRef, + range_ms: i64, + 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, +} + +/// [`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. +#[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, +} + +/// 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. + kept_stats: KeptResultCounters, + /// 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)>, +} + +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 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, &kept_builds, &windows, q.root_node_id as u64); + (q.query.trim().to_string(), program) + }) + .collect(); + 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: KeptResultCounters::default(), + kept_builds, + kept_matchers, + roots, + } + } + + /// 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, + progress: IngestionProgress, + ) -> Result { + let (root, program) = self.program(query)?; + for input in &program.stored { + pane(input, time_ms, store)?; + } + let mut ranges: Vec = program + .raw + .iter() + .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::>()?; + // 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 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, + root: usize, + time_ms: i64, + store: &SummaryStore, + progress: &IngestionProgress, + 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)?; + // Cloning a batch shares its buffers: the pane is not copied. + let source = + 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) => { + add(&self.kept_stats.reused); + (*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()))?; + 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) { + 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) { + add(&self.kept_stats.kept_while_metric_behind); + } + // 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 + } + } + }; + 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 + .selector + .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)); + } + 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) + } + + /// 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] + .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 + .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())), + } + } +} + +/// `(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( + input: &StoredInput, + time_ms: i64, + 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!( + "{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, + kept_builds: &BTreeMap>, + 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 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() { + if !seen.insert(id) { + continue; + } + if let Some((definition, binding)) = sds.outputs.get(&id) { + 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 { + 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, + matchers: label_matchers(dag, window.scan_node_id as u64), + 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 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(), + }) +} + +/// 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> { + 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/src/sds.rs b/data_plane/src/sds.rs index b5354310..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, @@ -85,6 +86,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 +157,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)); } @@ -316,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 { @@ -332,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 new file mode 100644 index 00000000..e66da5ca --- /dev/null +++ b/data_plane/tests/answers.rs @@ -0,0 +1,598 @@ +//! 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")); +} + +/// 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()); +} + +/// 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}" + ); + } +} + +/// 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 new file mode 100644 index 00000000..e9031afe --- /dev/null +++ b/data_plane/tests/kept.rs @@ -0,0 +1,360 @@ +//! 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..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; @@ -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,238 @@ 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)); +} + +/// 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); +} + +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/query.rs b/data_plane/tests/query.rs new file mode 100644 index 00000000..ee1a6d1c --- /dev/null +++ b/data_plane/tests/query.rs @@ -0,0 +1,347 @@ +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::{LabelMatcher, MatchOp, 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: &mut Fixture, query: &str, time_ms: i64, all: &[RawSample]) -> Vec { + let prepared = f + .query + .raw_ranges(query, time_ms, &f.precompute.store, f.precompute.progress()) + .unwrap(); + let raw = fetch(&prepared.ranges, all); + let mut result = f + .query + .evaluate( + prepared.root, + time_ms, + &mut f.precompute.store, + &prepared.progress, + &raw, + ) + .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, + Default::default() + ) + .is_err(), + "no panes before ingestion" + ); + panes.precompute.ingest(&all); + let ranges = panes + .query + .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(&mut panes, P99, 70 * MINUTE, &all); + + let mut single = fixture(common::plan("example4b", true)); + let ranges = single + .query + .raw_ranges( + P99, + 70 * MINUTE, + &single.precompute.store, + Default::default(), + ) + .unwrap() + .ranges; + 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(&mut 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 mut f = fixture(common::plan("example1", true)); + let result = answer( + &mut 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( + &mut 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, Default::default()), + Err(Unanswered::NotPlanned(_)) + )); + assert!(matches!( + 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, Default::default()) + .is_ok()); + assert!(matches!( + 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, Default::default()), + 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, 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, Default::default()) + .is_ok()); + let panes = fixture(common::plan("example4b", false)); + assert!(matches!( + panes + .query + .raw_ranges(P99, i64::MIN, &panes.precompute.store, Default::default()), + 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 mut by_reference = fixture(reference); + for query in plan.queries.iter().map(|q| q.query.as_str()) { + let expected = answer(&mut by_reference, query, 70 * MINUTE, &all); + assert!(!expected.is_empty(), "{example} raw={raw} {query}"); + assert_eq!( + answer(&mut 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, + Default::default(), + ) + .unwrap() + .ranges; + assert_eq!(ranges.len(), 1); + assert_eq!( + ranges[0].matchers, + vec![LabelMatcher { + name: "job".into(), + op: MatchOp::Equal, + value: "a".into(), + }] + ); +} diff --git a/data_plane/tests/sds.rs b/data_plane/tests/sds.rs index 1e27b416..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); @@ -251,3 +260,78 @@ 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::>()); +} + +/// 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); +} 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 + } + } + } +}