From d32a8bb0b4a29f8bf553b74682f2a46464d1acaa Mon Sep 17 00:00:00 2001 From: zzylol Date: Mon, 19 Oct 2026 06:24:57 +0000 Subject: [PATCH 1/2] feat(data-plane): sharded ingestion workers, a query thread pool and a shared pane store - Ingestion: N workers (`--ingestion-workers`, default one per core). Remote Write series are sharded by a hash of their label set; each worker owns its shard's open panes. The global watermark is the minimum over the workers that have seen a sample; when it passes a pane, every worker builds its part, and the last part completes the pane: parts are joined and the states of a group several workers saw are merged with the family's merge. A request is queued to every worker or refused whole. - Queries: a pool of `--query-threads` threads, one query (and every step of a range query) per thread. The executor's compiled plans hold `Rc`s, so each thread reads plan.json and compiles its own. - SDS: `RwLock>>`. Panes are inserted once and never rewritten; readers clone the `Arc`. Kept query-time results use the same store, shared by the query threads. A stored pane binds to a root as a shared source, copied once per run (the executor's sources emit owned batches) instead of twice. - jemalloc is the data plane's global allocator. --- Cargo.lock | 76 ++++--- README.md | 12 +- data_plane/Cargo.toml | 1 + data_plane/src/engine.rs | 380 ++++++++++++++++++++----------- data_plane/src/lib.rs | 47 ++-- data_plane/src/main.rs | 20 +- data_plane/src/precompute.rs | 370 +++++++++++++++++++++++++----- data_plane/src/query.rs | 144 +++++++++--- data_plane/src/sds.rs | 46 ++-- data_plane/tests/answers.rs | 25 ++ data_plane/tests/e2e_examples.rs | 2 +- data_plane/tests/kept.rs | 16 +- data_plane/tests/precompute.rs | 125 ++++++++-- data_plane/tests/query.rs | 14 +- data_plane/tests/range.rs | 10 +- data_plane/tests/sds.rs | 68 +++++- 16 files changed, 1031 insertions(+), 325 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f794f62e..c55c0c8b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1052,6 +1052,7 @@ dependencies = [ "sha2 0.10.9", "snap", "thiserror 1.0.69", + "tikv-jemallocator", "tokio", "tower", "tracing", @@ -1097,7 +1098,7 @@ dependencies = [ "flate2", "futures", "indexmap", - "itertools 0.14.0", + "itertools", "liblzma", "log", "object_store", @@ -1129,7 +1130,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", "parking_lot", @@ -1154,7 +1155,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", ] @@ -1173,7 +1174,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "libc", "log", "object_store", @@ -1220,7 +1221,7 @@ dependencies = [ "flate2", "futures", "glob", - "itertools 0.14.0", + "itertools", "liblzma", "log", "object_store", @@ -1251,7 +1252,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "object_store", "tokio", ] @@ -1325,7 +1326,7 @@ dependencies = [ "datafusion-pruning", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools", "log", "object_store", "parking_lot", @@ -1378,7 +1379,7 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr-common", "indexmap", - "itertools 0.14.0", + "itertools", "recursive", "serde_json", "sqlparser", @@ -1393,7 +1394,7 @@ dependencies = [ "arrow", "datafusion-common", "indexmap", - "itertools 0.14.0", + "itertools", ] [[package]] @@ -1417,7 +1418,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hex", - "itertools 0.14.0", + "itertools", "log", "md-5", "memchr", @@ -1480,7 +1481,7 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr-common", "hashbrown 0.17.1", - "itertools 0.14.0", + "itertools", "itoa", "log", "memchr", @@ -1553,7 +1554,7 @@ dependencies = [ "datafusion-expr-common", "datafusion-physical-expr", "indexmap", - "itertools 0.14.0", + "itertools", "log", "recursive", "regex", @@ -1575,7 +1576,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "parking_lot", "petgraph", "recursive", @@ -1594,7 +1595,7 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools 0.14.0", + "itertools", ] [[package]] @@ -1609,7 +1610,7 @@ dependencies = [ "datafusion-expr-common", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "parking_lot", "pin-project", ] @@ -1629,7 +1630,7 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", - "itertools 0.14.0", + "itertools", "recursive", ] @@ -1658,7 +1659,7 @@ dependencies = [ "half", "hashbrown 0.17.1", "indexmap", - "itertools 0.14.0", + "itertools", "log", "num-traits", "parking_lot", @@ -2339,15 +2340,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 +2780,7 @@ dependencies = [ "futures-util", "http", "humantime", - "itertools 0.14.0", + "itertools", "parking_lot", "percent-encoding", "thiserror 2.0.20", @@ -3069,7 +3061,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 +3378,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3783,10 +3775,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3849,6 +3841,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 +4345,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 f31039f8..afbd541f 100644 --- a/README.md +++ b/README.md @@ -39,9 +39,15 @@ The backend is being rebuilt on the Planner's physical DAG: | Query engine: answer planned instant queries from stored panes and raw samples | done | | Query engine: keep query-time panes across evaluations (`MatAtQueryKeptAcrossEvaluations`) | done | -Panes close `--allowed-lateness-ms` (default 15 s) after the newest sample -passes their end; later samples for them are dropped. The first pane after -startup is skipped because it may be partial. +Ingestion runs on `--ingestion-workers` threads and queries on +`--query-threads` threads (each default one per core). Series are sharded +across the ingestion workers by a hash of their label set; a pane closes +`--allowed-lateness-ms` (default 15 s) after every worker's newest sample +passes its end, and later samples for it are dropped. The first pane after +startup is skipped because it may be partial. Each query, including every +step of a range query, runs on one query thread; queries read the stored +panes concurrently, by shared reference. The data plane uses jemalloc as its +allocator. 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 diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index b82d13b7..0632b8ad 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 diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index ba44b400..8150a63e 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -1,25 +1,46 @@ -//! 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 panes, and +//! a pool of query threads that answer planned queries. +//! +//! Planner IR and the executor's compiled plans hold `Rc`s, so they cannot +//! be shared across threads: every worker and query thread reads the plan +//! and compiles its own. What threads share is plain data: the stored panes, +//! the kept query-time results, and the watermarks. +//! +//! - **Ingestion.** Series are sharded by a hash of their label set; each +//! worker owns the open panes of its shard. A worker's watermark is its +//! newest sample; a pane closes when the minimum over the workers passes +//! it. Each worker then builds its part of the pane, and the last part +//! completes it ([`crate::precompute::Shared`]). +//! - **Queries.** One query per thread, from a pool: a query, and every step +//! of a range query, runs on the thread that takes it, while other queries +//! run on the others. They read stored panes under the store's read lock. use std::collections::BTreeMap; -use std::sync::mpsc::{sync_channel, Receiver, SyncSender, TrySendError}; +use std::sync::mpsc::{sync_channel, Receiver, SyncSender}; +use std::sync::{Arc, Mutex}; use asap_planner::InstallablePlan; -use tokio::sync::oneshot; +use tokio::sync::{mpsc, oneshot}; -use crate::precompute::{Precompute, PrecomputeStats, RawSample}; -use crate::query::{QueryEngine, RawRange, ResultSample, Unanswered}; +use crate::precompute::{PrecomputeStats, RawSample, ShardWorker, Shared}; +use crate::query::{KeptResults, QueryEngine, RawRange, ResultSample, Unanswered}; use crate::sds::Sds; -/// Commands queued for the engine before ingestion pushes back. +/// Messages queued for each ingestion worker, and queries for the pool, +/// before the data plane pushes back. const QUEUE: usize = 1024; type Reply = oneshot::Sender>; -pub enum Command { +enum WorkerCommand { Ingest(Vec), - Stats(oneshot::Sender), + /// Another worker moved the global watermark: build the panes it passed. + Advance, + /// Reply once every command queued before has run. + Flush(oneshot::Sender<()>), +} + +enum QueryJob { RawRanges { query: String, time_ms: i64, @@ -48,39 +69,173 @@ pub enum Command { #[error("the engine is busy")] pub struct Busy; -/// The HTTP side of the engine thread. +/// How many threads the engine runs. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Threads { + /// 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 Threads { + /// One of each per available core. + fn default() -> Self { + let cores = std::thread::available_parallelism().map_or(1, |n| n.get()); + Self { + ingestion_workers: cores, + query_threads: cores, + } + } +} + +/// The HTTP side of the engine's threads. #[derive(Clone)] pub struct EngineHandle { - commands: SyncSender, + workers: Vec>, + queries: SyncSender, + precompute: Arc, +} + +/// 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)?; + Ok((plan, sds)) +} + +/// Runs `thread` on a new thread named `name`; returns once it has started, +/// or with the error it failed to start with. +fn spawn( + name: String, + thread: impl FnOnce(std::sync::mpsc::Sender>) + Send + 'static, +) -> anyhow::Result<()> { + let (ready, started) = std::sync::mpsc::channel(); + std::thread::Builder::new() + .name(name) + .spawn(move || thread(ready))?; + started.recv()? } impl EngineHandle { - /// Starts the engine for the plan serialized in `plan_json`. Fails if the - /// plan is invalid or the data plane cannot execute it. - pub fn spawn(plan_json: String, allowed_lateness_ms: i64) -> anyhow::Result { - let (commands, receiver) = sync_channel(QUEUE); - let (ready, started) = std::sync::mpsc::channel(); - std::thread::Builder::new() - .name("engine".into()) - .spawn(move || match Engine::new(&plan_json, allowed_lateness_ms) { - Ok(engine) => { - let _ = ready.send(Ok(())); - engine.serve(receiver); - } - Err(error) => { - let _ = ready.send(Err(error)); + /// Starts the engine's threads for the plan serialized in `plan_json`. + /// Fails if the plan is invalid or the data plane cannot execute it. + pub fn spawn( + plan_json: String, + allowed_lateness_ms: i64, + threads: Threads, + ) -> anyhow::Result { + let (plan, sds) = install(&plan_json)?; + tracing::info!( + definitions = sds.definitions.len(), + outputs = sds.outputs.len(), + kept_definitions = sds.kept_definitions.len(), + kept_outputs = sds.kept_outputs.len(), + ?threads, + "stored summaries" + ); + drop(plan); + let plan_json = Arc::new(plan_json); + let precompute = Arc::new(Shared::new(threads.ingestion_workers, allowed_lateness_ms)); + let (senders, receivers): (Vec<_>, Vec<_>) = (0..precompute.workers()) + .map(|_| mpsc::channel(QUEUE)) + .unzip(); + // Workers tell each other when the global watermark moves; weak + // senders let the channels close when the handle is dropped. + let peers: Vec<_> = senders.iter().map(mpsc::Sender::downgrade).collect(); + for (index, mut commands) in receivers.into_iter().enumerate() { + let (plan_json, shared, peers) = (plan_json.clone(), precompute.clone(), peers.clone()); + spawn(format!("ingest-{index}"), move |ready| { + let started = install(&plan_json) + .and_then(|(plan, sds)| Ok(ShardWorker::new(&plan, &sds, index, shared)?)); + let mut worker = match started { + Ok(worker) => { + let _ = ready.send(Ok(())); + worker + } + Err(error) => { + let _ = ready.send(Err(error)); + return; + } + }; + while let Some(command) = commands.blocking_recv() { + match command { + WorkerCommand::Ingest(samples) => { + if worker.ingest(&samples) { + for (peer, sender) in peers.iter().enumerate() { + if peer != index { + // A full queue runs `advance` with + // its next command anyway. + if let Some(sender) = sender.upgrade() { + let _ = sender.try_send(WorkerCommand::Advance); + } + } + } + } + } + WorkerCommand::Advance => worker.advance(), + WorkerCommand::Flush(reply) => { + let _ = reply.send(()); + } + } } })?; - started.recv()??; - Ok(Self { commands }) + } + let (queries, jobs) = sync_channel(QUEUE); + let jobs = Arc::new(Mutex::new(jobs)); + let kept = Arc::new(KeptResults::default()); + for index in 0..threads.query_threads.max(1) { + let (plan_json, shared, kept, jobs) = ( + plan_json.clone(), + precompute.clone(), + kept.clone(), + jobs.clone(), + ); + spawn(format!("query-{index}"), move |ready| { + let engine = match install(&plan_json) { + Ok((plan, sds)) => QueryEngine::with_kept(&plan, &sds, kept), + Err(error) => { + let _ = ready.send(Err(error)); + return; + } + }; + let _ = ready.send(Ok(())); + serve_queries(&engine, &shared, &jobs); + })?; + } + Ok(Self { + workers: senders, + queries, + precompute, + }) + } + + /// How many ingestion workers series are sharded across. + pub fn ingestion_workers(&self) -> usize { + self.workers.len() } - /// Queues `samples` for precompute, or fails if the engine is behind. - pub fn try_ingest(&self, samples: Vec) -> Result<(), Busy> { - match self.commands.try_send(Command::Ingest(samples)) { - Ok(()) => Ok(()), - Err(TrySendError::Full(_) | TrySendError::Disconnected(_)) => Err(Busy), + /// Queues each worker's shard of the samples (`shards[i]` for worker + /// `i`), or none if any worker is behind. + pub fn try_ingest(&self, shards: Vec>) -> Result<(), Busy> { + // Reserve every worker's slot first, so a request is taken whole or + // not at all: Remote Write retries a refused request. + let permits = self + .workers + .iter() + .zip(&shards) + .filter(|(_, samples)| !samples.is_empty()) + .map(|(sender, _)| sender.try_reserve().map_err(|_| Busy)) + .collect::, _>>()?; + for (permit, samples) in permits + .into_iter() + .zip(shards.into_iter().filter(|samples| !samples.is_empty())) + { + permit.send(WorkerCommand::Ingest(samples)); } + Ok(()) } /// The plan root answering `query` at `time_ms` and the raw samples it @@ -90,7 +245,7 @@ impl EngineHandle { query: String, time_ms: i64, ) -> Result<(usize, Vec), Unanswered> { - self.ask(|reply| Command::RawRanges { + self.ask(|reply| QueryJob::RawRanges { query, time_ms, reply, @@ -105,7 +260,7 @@ impl EngineHandle { time_ms: i64, raw: BTreeMap>, ) -> Result, Unanswered> { - self.ask(|reply| Command::Evaluate { + self.ask(|reply| QueryJob::Evaluate { root, time_ms, raw, @@ -121,7 +276,7 @@ impl EngineHandle { query: String, steps: Vec, ) -> Result<(usize, Vec), Unanswered> { - self.ask(|reply| Command::RawRangesForSteps { + self.ask(|reply| QueryJob::RawRangesForSteps { query, steps, reply, @@ -129,14 +284,15 @@ impl EngineHandle { .await } - /// Runs a root at every one of `steps`; see [`QueryEngine::evaluate_steps`]. + /// Runs a root at every one of `steps`, all on one query thread; see + /// [`QueryEngine::evaluate_steps`]. pub async fn evaluate_steps( &self, root: usize, steps: Vec, raw: BTreeMap>, ) -> Result>, Unanswered> { - self.ask(|reply| Command::EvaluateSteps { + self.ask(|reply| QueryJob::EvaluateSteps { root, steps, raw, @@ -145,105 +301,75 @@ impl EngineHandle { .await } - async fn ask(&self, command: impl FnOnce(Reply) -> Command) -> Result { + async fn ask(&self, job: impl FnOnce(Reply) -> QueryJob) -> Result { let busy = || Unanswered::Execution("the engine is busy".into()); let (reply, answer) = oneshot::channel(); - self.commands.try_send(command(reply)).map_err(|_| busy())?; + self.queries.try_send(job(reply)).map_err(|_| busy())?; answer.await.map_err(|_| busy())? } - /// Precompute counters, once every command queued before has run. + /// Precompute counters, once every sample queued before has been + /// ingested and every pane it closed has been built. pub async fn stats(&self) -> PrecomputeStats { - let (reply, stats) = oneshot::channel(); - if self.commands.try_send(Command::Stats(reply)).is_err() { - return PrecomputeStats::default(); + // The first round runs every queued ingest, which may tell other + // workers to advance; the second runs those advances. + for _ in 0..2 { + let mut flushed = Vec::new(); + for worker in &self.workers { + let (reply, done) = oneshot::channel(); + if worker.send(WorkerCommand::Flush(reply)).await.is_ok() { + flushed.push(done); + } + } + for done in flushed { + let _ = done.await; + } } - stats.await.unwrap_or_default() + self.precompute.stats() } } -struct Engine { - precompute: Precompute, - query: QueryEngine, -} - -impl Engine { - fn new(plan_json: &str, allowed_lateness_ms: i64) -> anyhow::Result { - 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(), - kept_definitions = sds.kept_definitions.len(), - kept_outputs = sds.kept_outputs.len(), - "stored summaries" - ); - let query = QueryEngine::new(&plan, &sds); - Ok(Self { - precompute: Precompute::new(&plan, &sds, allowed_lateness_ms)?, - query, - }) - } - - fn serve(mut self, commands: Receiver) { - for command in commands { - match command { - Command::Ingest(samples) => self.precompute.ingest(&samples), - Command::Stats(reply) => { - let _ = reply.send(self.precompute.stats); - } - Command::RawRanges { - query, - time_ms, - reply, - } => { - let _ = reply.send(self.query.raw_ranges( - &query, - time_ms, - &self.precompute.store, - )); - } - Command::Evaluate { - root, - time_ms, - raw, - reply, - } => { - let _ = reply.send(self.query.evaluate( - root, - time_ms, - &self.precompute.store, - self.precompute.watermark_ms(), - &raw, - )); - } - Command::RawRangesForSteps { - query, - steps, - reply, - } => { - let _ = reply.send(self.query.raw_ranges_for_steps( - &query, - &steps, - &self.precompute.store, - )); - } - Command::EvaluateSteps { - root, - steps, - raw, - reply, - } => { - // One watermark for every step of the request. - let _ = reply.send(self.query.evaluate_steps( - root, - &steps, - &self.precompute.store, - self.precompute.watermark_ms(), - &raw, - )); - } +/// 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().expect("query queue lock").recv() { + Ok(job) => job, + Err(_) => return, + }; + let store = &precompute.store; + match job { + QueryJob::RawRanges { + query, + time_ms, + reply, + } => { + let _ = reply.send(engine.raw_ranges(&query, time_ms, store)); + } + QueryJob::Evaluate { + root, + time_ms, + raw, + reply, + } => { + let watermark = precompute.watermark_ms(); + let _ = reply.send(engine.evaluate(root, time_ms, store, watermark, &raw)); + } + QueryJob::RawRangesForSteps { + query, + steps, + reply, + } => { + let _ = reply.send(engine.raw_ranges_for_steps(&query, &steps, store)); + } + QueryJob::EvaluateSteps { + root, + steps, + raw, + reply, + } => { + // One watermark for every step of the request. + let watermark = precompute.watermark_ms(); + let _ = reply.send(engine.evaluate_steps(root, &steps, store, watermark, &raw)); } } } diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index c813c4cb..7e6954f2 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -2,10 +2,11 @@ //! the Prometheus query API for an installed plan. //! //! The installed plan is ASAPPlanner's `InstallablePlan`: its selected -//! physical DAG, the outputs it stores and the raw windows it reads. The -//! engine thread builds its ingestion-time outputs pane by pane and answers -//! the plan's instant and range queries from them; every other query, and -//! any planned query it cannot answer, is forwarded to Prometheus. +//! physical DAG, the outputs it stores and the raw windows it reads. +//! Ingestion workers build its ingestion-time outputs pane by pane, and a +//! pool of query threads answers the plan's instant and range queries from +//! them ([`engine`]); every other query, and any planned query it cannot +//! answer, is forwarded to Prometheus. pub mod engine; pub mod precompute; @@ -25,8 +26,8 @@ use axum::http::{header, HeaderMap, HeaderValue, Method, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::Router; -use engine::EngineHandle; -use precompute::RawSample; +use engine::{EngineHandle, Threads}; +use precompute::{shard_of, RawSample}; use prometheus::{InstantQuery, RangeQuery}; use query::Unanswered; @@ -51,18 +52,34 @@ pub struct DataPlane { } impl DataPlane { - /// Installs the plan serialized in `plan_json` (an [`InstallablePlan`]). - /// Panes close `allowed_lateness_ms` after the newest sample passes them. + /// Installs the plan serialized in `plan_json` (an [`InstallablePlan`]) + /// with [`Threads::default`]. Panes close `allowed_lateness_ms` after + /// every ingestion worker's newest sample passes them. pub fn new( plan_json: String, prometheus_url: String, allowed_lateness_ms: i64, + ) -> anyhow::Result { + Self::with_threads( + plan_json, + prometheus_url, + allowed_lateness_ms, + Threads::default(), + ) + } + + /// [`Self::new`] with `threads` ingestion workers and query threads. + pub fn with_threads( + plan_json: String, + prometheus_url: String, + allowed_lateness_ms: i64, + threads: Threads, ) -> anyhow::Result { let plan: InstallablePlan = serde_json::from_str(&plan_json)?; plan.validate()?; Ok(Self { queries: plan.queries.into_iter().map(|q| q.query).collect(), - engine: EngineHandle::spawn(plan_json, allowed_lateness_ms)?, + engine: EngineHandle::spawn(plan_json, allowed_lateness_ms, threads)?, prometheus_url: prometheus_url.trim_end_matches('/').to_string(), limits: remote_write::Limits::default(), samples: AtomicU64::new(0), @@ -83,7 +100,9 @@ pub fn router(state: Arc) -> Router { async fn write(State(state): State>, body: Bytes) -> Response { match remote_write::decode(&body, &state.limits) { Ok(request) => { - let mut samples = Vec::new(); + let workers = state.engine.ingestion_workers(); + let mut shards = vec![Vec::new(); workers]; + let mut count = 0; for series in request.timeseries { let labels: Arc> = Arc::new( series @@ -92,7 +111,9 @@ async fn write(State(state): State>, body: Bytes) -> Response { .map(|label| (label.name, label.value)) .collect(), ); - samples.extend( + let shard = &mut shards[shard_of(&labels, workers)]; + let before = shard.len(); + shard.extend( series .samples .into_iter() @@ -103,9 +124,9 @@ async fn write(State(state): State>, body: Bytes) -> Response { value: s.value, }), ); + count += (shard.len() - before) as u64; } - let count = samples.len() as u64; - match state.engine.try_ingest(samples) { + match state.engine.try_ingest(shards) { Ok(()) => { state.samples.fetch_add(count, Ordering::Relaxed); StatusCode::NO_CONTENT.into_response() diff --git a/data_plane/src/main.rs b/data_plane/src/main.rs index d8a24abf..990f4283 100644 --- a/data_plane/src/main.rs +++ b/data_plane/src/main.rs @@ -2,8 +2,14 @@ use std::path::PathBuf; use std::sync::Arc; use clap::Parser; +use data_plane::engine::Threads; use data_plane::{router, DataPlane}; +// jemalloc, as ASAPQuery's engine uses: the executor allocates per row, and +// glibc's malloc and free took a third of the data plane's CPU. +#[global_allocator] +static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + /// Serve Remote Write ingestion and the Prometheus query API for a plan. #[derive(Parser)] struct Args { @@ -22,6 +28,12 @@ struct Args { /// executor's default is 64); an evaluation over it is forwarded. #[arg(long, default_value_t = 64)] execution_memory_limit_mib: usize, + /// Ingestion workers series are sharded across; default one per core. + #[arg(long)] + ingestion_workers: Option, + /// Query threads, each running one query at a time; default one per core. + #[arg(long)] + query_threads: Option, } #[tokio::main] @@ -32,10 +44,16 @@ async fn main() -> anyhow::Result<()> { let args = Args::parse(); data_plane::precompute::set_execution_memory_limit(args.execution_memory_limit_mib << 20); let plan = std::fs::read_to_string(&args.plan)?; - let state = Arc::new(DataPlane::new( + let defaults = Threads::default(); + let threads = Threads { + ingestion_workers: args.ingestion_workers.unwrap_or(defaults.ingestion_workers), + query_threads: args.query_threads.unwrap_or(defaults.query_threads), + }; + let state = Arc::new(DataPlane::with_threads( plan, args.prometheus_url, args.allowed_lateness_ms, + threads, )?); tracing::info!(queries = state.queries.len(), "installed plan"); let listener = tokio::net::TcpListener::bind(&args.listen).await?; diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index 7f9a10af..4ac4fd56 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -13,12 +13,15 @@ //! did, so it is never built: queries that need it are not answered from //! stored panes. -use std::collections::BTreeMap; -use std::sync::{Arc, OnceLock}; +use std::collections::{BTreeMap, HashMap}; +use std::hash::{Hash, Hasher}; +use std::sync::atomic::{AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; use asap_executor::operators::Operator; use asap_executor::physical_planner::{compile, CompiledPhysicalDAG, InputContract, Source}; use asap_executor::plan::PhysicalDAG; +use asap_executor::planner::ir::schema::FieldDataType; use asap_executor::runtime::{Limits, RunContext, Scope}; use asap_executor::values::{Batch, SchemaRef, Value}; use asap_executor::{RawSampleSelector, Sample}; @@ -120,30 +123,198 @@ impl PaneBuild { } } +/// One worker's part of a pane, or why its build failed. +type PanePart = Result; + +/// No value yet in an `AtomicI64` slot. +const UNSET: i64 = i64::MIN; + +/// Precompute state every ingestion worker and query thread shares: the +/// stored panes, each worker's watermark, and the per-worker partial panes +/// not yet complete. +pub struct Shared { + /// Stored panes, by definition and pane end. + pub store: SummaryStore, + allowed_lateness_ms: i64, + /// Each worker's newest sample time, `UNSET` before its first. + newest_sample_ms: Vec, + /// The first sample's time: every worker skips the pane it falls in, so + /// all workers build the same panes. + start_ms: AtomicI64, + /// The newest global watermark any worker has announced to the others. + announced_ms: AtomicI64, + /// Partial panes by definition and end, until every worker has built + /// its part. + pending: Mutex>>, + accepted: AtomicU64, + dropped: AtomicU64, + panes_built: AtomicU64, + panes_failed: AtomicU64, +} + +impl Shared { + pub fn new(workers: usize, allowed_lateness_ms: i64) -> Self { + Self { + store: SummaryStore::new(), + allowed_lateness_ms, + newest_sample_ms: (0..workers.max(1)).map(|_| AtomicI64::new(UNSET)).collect(), + start_ms: AtomicI64::new(UNSET), + announced_ms: AtomicI64::new(UNSET), + pending: Mutex::new(HashMap::new()), + accepted: AtomicU64::new(0), + dropped: AtomicU64::new(0), + panes_built: AtomicU64::new(0), + panes_failed: AtomicU64::new(0), + } + } + + pub fn workers(&self) -> usize { + self.newest_sample_ms.len() + } + + /// The global watermark: no sample older than this is expected any + /// more. It is the minimum over the workers that have seen a sample of + /// their newest sample time, minus the allowed lateness; `None` before + /// any sample. + pub fn watermark_ms(&self) -> Option { + self.newest_sample_ms + .iter() + .map(|newest| newest.load(Ordering::Acquire)) + .filter(|&newest| newest != UNSET) + .min() + .map(|newest| newest - self.allowed_lateness_ms) + } + + pub fn stats(&self) -> PrecomputeStats { + PrecomputeStats { + accepted: self.accepted.load(Ordering::Relaxed), + dropped: self.dropped.load(Ordering::Relaxed), + panes_built: self.panes_built.load(Ordering::Relaxed), + panes_failed: self.panes_failed.load(Ordering::Relaxed), + } + } + + /// Adds one worker's part of the pane of `definition` ending at `end_ms`. + /// The worker that adds the last part merges them and stores the pane. + fn add_part(&self, definition: &SummaryDefinition, end_ms: i64, part: Result) { + let parts = { + let mut pending = self.pending.lock().expect("pending panes lock"); + let key = (definition.id.clone(), end_ms); + let parts = pending.entry(key.clone()).or_default(); + parts.push(part); + if parts.len() < self.workers() { + return; + } + pending.remove(&key).expect("present") + }; + match parts + .into_iter() + .collect::, _>>() + .and_then(|parts| merge_parts(definition, parts)) + { + Ok(pane) => { + self.store.insert(&definition.id, end_ms, pane); + self.panes_built.fetch_add(1, Ordering::Relaxed); + // The newest evaluation that reads stored panes is at most + // the newest stored pane's end, so retention is measured + // from it, not from the watermark, which may be past it. + self.store + .evict_before(&definition.id, definition.oldest_kept_end_ms(end_ms)); + } + Err(error) => { + tracing::error!(definition = ?definition.id, end_ms, %error, "pane build failed"); + self.panes_failed.fetch_add(1, Ordering::Relaxed); + } + } + } +} + +/// One pane from the parts the workers built over their own series: their +/// rows, with the states of a group several workers saw merged by the +/// family's merge. Each worker's part has one row per group, so a group +/// present in one part only is kept as built. +fn merge_parts(definition: &SummaryDefinition, parts: Vec) -> Result { + let schema = Arc::new(definition.output_schema.clone()); + let state = definition + .output_schema + .fields + .iter() + .position(|f| !matches!(f.dtype, FieldDataType::Plain(_))); + let mut rows: BTreeMap>, Vec> = BTreeMap::new(); + for part in &parts { + for row in part.rows() { + let key = row + .iter() + .enumerate() + .filter(|(i, _)| Some(*i) != state) + .map(|(_, value)| value.key()) + .collect::, _>>() + .map_err(|e| e.to_string())?; + match (rows.get_mut(&key), state) { + (Some(existing), Some(state)) => { + let ( + Value::Summary { state: left, .. }, + Value::Summary { + family, + state: right, + }, + ) = (&existing[state], &row[state]) + else { + return Err("summary state column without a state".into()); + }; + let merged = left.merge_with(right.as_ref()).map_err(|e| e.to_string())?; + existing[state] = Value::Summary { + family: family.clone(), + state: Arc::from(merged), + }; + } + (Some(_), None) => return Err("duplicate group without a state".into()), + (None, _) => { + rows.insert(key, row.clone()); + } + } + } + } + Batch::try_new(schema, rows.into_values().collect()).map_err(|e| e.to_string()) +} + +/// The ingestion worker of `workers` that owns the series with `labels`. +pub fn shard_of(labels: &BTreeMap, workers: usize) -> usize { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + labels.hash(&mut hasher); + (hasher.finish() % workers.max(1) as u64) as usize +} + struct Builder { pane: PaneBuild, - /// Samples of open panes, by pane end. + /// This worker's samples of open panes, by pane end. open: BTreeMap>, - /// End of the newest pane built or skipped; samples in it or before it - /// are dropped. `None` until the first sample. + /// End of the newest pane this worker built its part of, or skipped; + /// samples in it or before it are dropped. `None` before the first + /// sample anywhere. built_through: Option, } -pub struct Precompute { +/// One ingestion worker: buffers the samples of its shard of the series and +/// builds its part of each pane once the global watermark passes it. It +/// holds the plan's compiled builds, which are not `Send`, so each worker +/// thread compiles its own. +pub struct ShardWorker { + index: usize, builders: Vec, - allowed_lateness_ms: i64, - newest_sample_ms: Option, - pub store: SummaryStore, - pub stats: PrecomputeStats, + shared: Arc, + /// The global watermark this worker has built its parts through. + closed_through_ms: Option, } -impl Precompute { - /// Compiles one build per ingestion-time definition of `sds`, the - /// stored summaries of `plan`. +impl ShardWorker { + /// Worker `index` of `shared`'s workers, with one build per + /// ingestion-time definition of `sds`, the stored summaries of `plan`. pub fn new( plan: &InstallablePlan, sds: &Sds, - allowed_lateness_ms: i64, + index: usize, + shared: Arc, ) -> Result { let builders = sds .definitions @@ -157,35 +328,43 @@ impl Precompute { }) .collect::>()?; Ok(Self { + index, builders, - allowed_lateness_ms, - newest_sample_ms: None, - store: SummaryStore::new(), - stats: PrecomputeStats::default(), + shared, + closed_through_ms: None, }) } - /// The watermark: no sample older than this is expected any more. - pub fn watermark_ms(&self) -> Option { - self.newest_sample_ms - .map(|newest| newest - self.allowed_lateness_ms) - } - - /// Buffers `samples` into their panes, then builds every pane the - /// watermark has passed. - pub fn ingest(&mut self, samples: &[RawSample]) { + /// Buffers `samples`, all of this worker's shard, into their panes, then + /// builds this worker's part of every pane the global watermark has + /// passed. Returns whether this moved the global watermark past what + /// was announced, so the caller tells the other workers to + /// [`Self::advance`]. + pub fn ingest(&mut self, samples: &[RawSample]) -> bool { + let Some(first) = samples.iter().map(|s| s.timestamp_ms).min() else { + return false; + }; + let shared = self.shared.clone(); + let _ = shared + .start_ms + .compare_exchange(UNSET, first, Ordering::AcqRel, Ordering::Acquire); + let start_ms = shared.start_ms.load(Ordering::Acquire); + let (mut accepted, mut dropped) = (0, 0); for sample in samples { let metric = sample.labels.get("__name__").map(String::as_str); - let mut accepted = false; + let mut buffered = false; for builder in &mut self.builders { if metric != Some(builder.pane.selector.metric()) { continue; } - let end = pane_end(sample.timestamp_ms, builder.pane.definition.pane_ms); - // The first pane seen may be partial: skip it. - let built = *builder.built_through.get_or_insert(end); + let pane_ms = builder.pane.definition.pane_ms; + let end = pane_end(sample.timestamp_ms, pane_ms); + // The pane of the first sample may be partial: skip it. + let built = *builder + .built_through + .get_or_insert(pane_end(start_ms, pane_ms)); if end <= built { - self.stats.dropped += 1; + dropped += 1; continue; } builder @@ -193,27 +372,48 @@ impl Precompute { .entry(end) .or_default() .push(sample.to_sample()); - accepted = true; + buffered = true; } - if accepted { - self.stats.accepted += 1; + accepted += u64::from(buffered); + } + shared.accepted.fetch_add(accepted, Ordering::Relaxed); + shared.dropped.fetch_add(dropped, Ordering::Relaxed); + let newest = samples + .iter() + .map(|s| s.timestamp_ms) + .max() + .expect("nonempty"); + shared.newest_sample_ms[self.index].fetch_max(newest, Ordering::AcqRel); + self.advance(); + match shared.watermark_ms() { + Some(watermark) => { + shared.announced_ms.fetch_max(watermark, Ordering::AcqRel) < watermark } - self.newest_sample_ms = self.newest_sample_ms.max(Some(sample.timestamp_ms)); + None => false, } - self.close_panes(); } - fn close_panes(&mut self) { - let Some(watermark) = self.watermark_ms() else { + /// Builds this worker's part of every pane the global watermark has + /// passed: from its buffered samples, or empty. Every worker builds every + /// pane after the first, so a pane without samples is stored as empty, + /// not missing. + pub fn advance(&mut self) { + let Some(watermark) = self.shared.watermark_ms() else { return; }; + if self + .closed_through_ms + .is_some_and(|closed| watermark <= closed) + { + return; + } + self.closed_through_ms = Some(watermark); + let start_ms = self.shared.start_ms.load(Ordering::Acquire); for builder in &mut self.builders { - let Some(built) = builder.built_through else { - continue; - }; let pane_ms = builder.pane.definition.pane_ms; - // Every pane after the first through the watermark is built, so - // a pane without samples is stored as empty, not missing. + let built = *builder + .built_through + .get_or_insert(pane_end(start_ms, pane_ms)); let mut end = built + pane_ms; while end <= watermark { let samples = builder.open.remove(&end).unwrap_or_default(); @@ -222,31 +422,77 @@ impl Precompute { window_end_ms: end, revision: 1, }; - match builder.pane.run(&samples, end, scope) { - Ok(state) => { - self.store.insert(&builder.pane.definition.id, end, state); - self.stats.panes_built += 1; - } - Err(error) => { - tracing::error!(definition = ?builder.pane.definition.id, end, %error, "pane build failed"); - self.stats.panes_failed += 1; - } - } + let part = builder + .pane + .run(&samples, end, scope) + .map_err(|e| e.to_string()); + self.shared.add_part(&builder.pane.definition, end, part); builder.built_through = Some(end); end += pane_ms; } - // The newest evaluation that reads stored panes is at most the - // newest closed pane's end, so retention is measured from it, - // not from the watermark, which may already be past it. - let definition = &builder.pane.definition; - if let Some(newest) = builder.built_through { - self.store - .evict_before(&definition.id, definition.oldest_kept_end_ms(newest)); - } } } } +/// 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, + workers: Vec, +} + +impl Precompute { + /// One worker for every series. + pub fn new( + plan: &InstallablePlan, + sds: &Sds, + allowed_lateness_ms: i64, + ) -> Result { + Self::with_shards(plan, sds, allowed_lateness_ms, 1) + } + + /// `shards` workers, each owning the series [`shard_of`] assigns it. + pub fn with_shards( + plan: &InstallablePlan, + sds: &Sds, + allowed_lateness_ms: i64, + shards: usize, + ) -> Result { + let shared = Arc::new(Shared::new(shards, allowed_lateness_ms)); + let workers = (0..shared.workers()) + .map(|index| ShardWorker::new(plan, sds, index, shared.clone())) + .collect::>()?; + Ok(Self { shared, workers }) + } + + /// Routes `samples` to their workers, then lets every worker build the + /// panes the global watermark has passed. + pub fn ingest(&mut self, samples: &[RawSample]) { + let mut shards = vec![Vec::new(); self.workers.len()]; + for sample in samples { + shards[shard_of(&sample.labels, self.workers.len())].push(sample.clone()); + } + for (worker, samples) in self.workers.iter_mut().zip(&shards) { + worker.ingest(samples); + } + for worker in &mut self.workers { + worker.advance(); + } + } + + pub fn watermark_ms(&self) -> Option { + self.shared.watermark_ms() + } + + pub fn store(&self) -> &SummaryStore { + &self.shared.store + } + + pub fn stats(&self) -> PrecomputeStats { + self.shared.stats() + } +} + /// End of the tumbling pane of width `pane_ms` that contains `timestamp_ms`. /// Panes are `(end - pane_ms, end]`, like PromQL ranges: a sample at a pane /// boundary belongs to the pane it ends. diff --git a/data_plane/src/query.rs b/data_plane/src/query.rs index 8cb75c82..43247a52 100644 --- a/data_plane/src/query.rs +++ b/data_plane/src/query.rs @@ -25,17 +25,20 @@ //! forwards it to Prometheus. use std::collections::{BTreeMap, BTreeSet, HashMap}; -use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use asap_executor::operators::Operator; use asap_executor::physical_planner::{compile, CompiledPhysicalDAG, InputContract, Source}; +use asap_executor::plan::{Boundedness, Emission, PhysicalOperator, PlanProperties}; use asap_executor::planner::ir::physical_export::RawSampleWindow; use asap_executor::planner::ir::schema::{DataType, FieldDataType, Schema}; use asap_executor::planner::ir::{NonASAPOp, Operator as IrOperator}; -use asap_executor::runtime::Scope; -use asap_executor::values::{Batch, Value}; +use asap_executor::runtime::{Input, OutputStream, RunContext, Scope}; +use asap_executor::values::{Batch, SchemaRef, Value}; use asap_executor::{series_labels_of_output_row, RawSampleSelector, Sample}; use asap_planner::{InstallablePlan, PhysicalASAPDAG}; +use futures::StreamExt; use crate::precompute::{output_schema, run, PaneBuild, RawSample}; use crate::sds::{DefinitionId, Sds, SummaryStore}; @@ -109,14 +112,45 @@ pub struct KeptResultStats { pub reused: u64, } -pub struct QueryEngine { +/// The query-time results kept across evaluations, shared by every query +/// thread: a pane one evaluation builds is reused by the next, whichever +/// thread runs it. +#[derive(Default)] +pub struct KeptResults { /// Kept panes (`MatAtQueryKeptAcrossEvaluations`), by definition and /// pane end, in memory. - pub kept: SummaryStore, + pub panes: SummaryStore, + built: AtomicU64, + reused: AtomicU64, + /// Newest evaluation time, from which kept panes are evicted; `None` + /// before the first. + newest_evaluation_ms: Mutex>, +} + +impl KeptResults { /// How many kept panes evaluations built and reused. - pub kept_stats: KeptResultStats, - /// Newest evaluation time, from which kept panes are evicted. - newest_evaluation_ms: Option, + pub fn stats(&self) -> KeptResultStats { + KeptResultStats { + built: self.built.load(Ordering::Relaxed), + reused: self.reused.load(Ordering::Relaxed), + } + } + + /// The newest evaluation time, with `time_ms` evaluated. + fn newest_evaluation(&self, time_ms: i64) -> i64 { + let mut newest = self.newest_evaluation_ms.lock().expect("kept lock"); + let value = newest.map_or(time_ms, |n| n.max(time_ms)); + *newest = Some(value); + value + } +} + +/// One thread's compiled roots of the plan. The executor's compiled plans +/// are not `Send`, so each query thread builds its own; they share the +/// stored panes and the [`KeptResults`]. +pub struct QueryEngine { + /// Results kept across evaluations, shared with the other query threads. + pub kept: Arc, /// How each kept definition's panes are built, or why they cannot be. kept_builds: BTreeMap>, /// Each query, with its compiled root or why it cannot run. @@ -124,7 +158,13 @@ pub struct QueryEngine { } impl QueryEngine { + /// Compiles every root of `plan`, with its own kept results. pub fn new(plan: &InstallablePlan, sds: &Sds) -> Self { + Self::with_kept(plan, sds, Arc::default()) + } + + /// Compiles every root of `plan`, keeping query-time results in `kept`. + pub fn with_kept(plan: &InstallablePlan, sds: &Sds, kept: Arc) -> Self { let dag = &plan.dag.dag; let kept_builds: BTreeMap<_, _> = sds .kept_definitions @@ -149,9 +189,7 @@ impl QueryEngine { }) .collect(); Self { - kept: SummaryStore::new(), - kept_stats: KeptResultStats::default(), - newest_evaluation_ms: None, + kept, kept_builds, roots, } @@ -186,7 +224,7 @@ impl QueryEngine { .collect(); for input in &program.kept { let end_ms = time_ms - input.offset_ms; - if self.kept.get(&input.definition, end_ms).is_none() { + if self.kept.panes.get(&input.definition, end_ms).is_none() { let build = self.kept_build(&input.definition); ranges.push(RawRange { node: input.raw_input, @@ -206,7 +244,7 @@ impl QueryEngine { /// (`None` before any sample). Then kept panes no evaluation after the /// newest one reads are evicted. pub fn evaluate( - &mut self, + &self, root: usize, time_ms: i64, store: &SummaryStore, @@ -224,17 +262,17 @@ impl QueryEngine { }; let mut sources: BTreeMap> = BTreeMap::new(); for input in &program.stored { - let pane = pane(input, time_ms, store)?; - let source = - Operator::source(pane.schema().clone(), vec![pane.clone()]).map_err(failed)?; - sources.insert(input.node, Box::new(source)); + sources.insert( + input.node, + Box::new(PaneSource(pane(input, time_ms, store)?)), + ); } for input in &program.kept { let end_ms = time_ms - input.offset_ms; - let pane = match self.kept.get(&input.definition, end_ms) { + let pane = match self.kept.panes.get(&input.definition, end_ms) { Some(pane) => { - self.kept_stats.reused += 1; - pane.clone() + self.kept.reused.fetch_add(1, Ordering::Relaxed); + pane } None => { // `raw_ranges` asked for these samples unless the pane @@ -255,17 +293,17 @@ impl QueryEngine { .kept_build(&input.definition) .run(&samples_of(samples), end_ms, pane_scope) .map_err(|e| Unanswered::Execution(e.to_string()))?; - self.kept_stats.built += 1; + self.kept.built.fetch_add(1, Ordering::Relaxed); // Kept only once closed, as the precompute engine closes // panes: a sample may still arrive for an open one. if watermark_ms.is_some_and(|watermark| end_ms <= watermark) { - self.kept.insert(&input.definition, end_ms, pane.clone()); + self.kept.panes.insert(&input.definition, end_ms, pane) + } else { + Arc::new(pane) } - pane } }; - let source = Operator::source(pane.schema().clone(), vec![pane]).map_err(failed)?; - sources.insert(input.node, Box::new(source)); + sources.insert(input.node, Box::new(PaneSource(pane))); } for input in &program.raw { let samples = raw.get(&input.node).map(Vec::as_slice).unwrap_or_default(); @@ -279,12 +317,11 @@ impl QueryEngine { } // Evict the kept panes older than each definition's retention // before the newest evaluation, this one or a later one. - let newest = time_ms.max(self.newest_evaluation_ms.unwrap_or(time_ms)); - self.newest_evaluation_ms = Some(newest); + let newest = self.kept.newest_evaluation(time_ms); for (definition, build) in &self.kept_builds { if let Ok(build) = build { let oldest = build.definition.oldest_kept_end_ms(newest); - self.kept.evict_before(definition, oldest); + self.kept.panes.evict_before(definition, oldest); } } let dag = program.compiled.instantiate(sources).map_err(failed)?; @@ -328,7 +365,7 @@ impl QueryEngine { for &time_ms in steps { for input in &program.kept { let end_ms = time_ms - input.offset_ms; - if self.kept.get(&input.definition, end_ms).is_none() { + if self.kept.panes.get(&input.definition, end_ms).is_none() { let span = missing.entry(&input.definition).or_insert((end_ms, end_ms)); *span = (span.0.min(end_ms), span.1.max(end_ms)); } @@ -354,7 +391,7 @@ impl QueryEngine { /// keeps a pane only if it is closed at `watermark_ms`, so a pane still /// open is recomputed at each step that reads it. pub fn evaluate_steps( - &mut self, + &self, root: usize, steps: &[i64], store: &SummaryStore, @@ -408,7 +445,7 @@ impl QueryEngine { } for (definition, offset_ms, node, fetched, pane_ms) in &kept_inputs { let end_ms = time_ms - offset_ms; - if self.kept.get(definition, end_ms).is_none() { + if self.kept.panes.get(definition, end_ms).is_none() { if let Some(runs) = series.get(fetched) { step_raw.insert(*node, slice(runs, end_ms - pane_ms, end_ms)); } @@ -468,12 +505,51 @@ fn samples_of(samples: &[RawSample]) -> Vec { samples.iter().map(RawSample::to_sample).collect() } +/// A stored or kept pane bound as a source of a compiled root. The pane is +/// shared, not copied into the source; the executor's sources emit owned +/// batches, so each run takes one copy of its rows. +struct PaneSource(Arc); + +impl PhysicalOperator for PaneSource { + fn name(&self) -> &str { + "StoredPane" + } + + fn properties(&self, _inputs: &[PlanProperties]) -> PlanProperties { + PlanProperties { + boundedness: Boundedness::Bounded, + emission: Emission::Incremental, + } + } + + fn input_schemas(&self) -> Vec { + Vec::new() + } + + fn output_schema(&self) -> SchemaRef { + self.0.schema().clone() + } + + fn start<'a>( + &'a self, + _inputs: Vec>, + _context: RunContext, + ) -> Result, asap_executor::Error> { + let batch = (*self.0).clone(); + Ok(futures::stream::once(async move { Ok(batch) }).boxed_local()) + } + + fn output_bytes(&self, value: &Batch) -> usize { + value.bytes() + } +} + /// The stored pane `input` reads at `time_ms`. -fn pane<'a>( +fn pane( input: &StoredInput, time_ms: i64, - store: &'a SummaryStore, -) -> Result<&'a Batch, Unanswered> { + store: &SummaryStore, +) -> Result, Unanswered> { let end = time_ms - input.offset_ms; if end.rem_euclid(input.pane_ms) != 0 { return Err(Unanswered::MissingPane(format!( diff --git a/data_plane/src/sds.rs b/data_plane/src/sds.rs index 60a542cf..2a78b75e 100644 --- a/data_plane/src/sds.rs +++ b/data_plane/src/sds.rs @@ -15,6 +15,7 @@ //! node computes, not which nodes are equal. use std::collections::{BTreeMap, HashMap}; +use std::sync::{Arc, RwLock}; use asap_executor::planner::ir::physical_export::{ PhysicalASAPDAG, PhysicalASAPDAGNode, RawSampleWindow, @@ -188,39 +189,56 @@ fn hex(bytes: &[u8]) -> String { bytes.iter().map(|b| format!("{b:02x}")).collect() } -/// Built panes, by definition and pane end. A pane is stored once, when it -/// is complete, and never changes. -#[derive(Debug, Default)] +/// Built panes, by definition and pane end, shared by every thread that +/// reads or builds them. A pane is stored once, when it is complete, and +/// never changes: a read takes the read lock only to clone the pane'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 { - panes: BTreeMap<(DefinitionId, i64), S>, + panes: RwLock>>, +} + +impl Default for SummaryStore { + fn default() -> Self { + Self::new() + } } impl SummaryStore { pub fn new() -> Self { Self { - panes: BTreeMap::new(), + panes: RwLock::new(HashMap::new()), } } - pub fn insert(&mut self, definition: &DefinitionId, end_ms: i64, state: S) { - self.panes.insert((definition.clone(), end_ms), state); + /// Stores the pane ending at `end_ms` unless it is already stored, and + /// returns the stored pane: a pane is never rewritten. + pub fn insert(&self, definition: &DefinitionId, end_ms: i64, state: S) -> Arc { + let mut panes = self.panes.write().expect("store lock"); + Arc::clone( + panes + .entry((definition.clone(), end_ms)) + .or_insert_with(|| Arc::new(state)), + ) } - pub fn get(&self, definition: &DefinitionId, end_ms: i64) -> Option<&S> { - self.panes.get(&(definition.clone(), end_ms)) + pub fn get(&self, definition: &DefinitionId, end_ms: i64) -> Option> { + let panes = self.panes.read().expect("store lock"); + panes.get(&(definition.clone(), end_ms)).cloned() } /// Drops the panes of `definition` that end before `end_ms`. - pub fn evict_before(&mut self, definition: &DefinitionId, end_ms: i64) { - self.panes - .retain(|(id, end), _| id != definition || *end >= end_ms); + pub fn evict_before(&self, definition: &DefinitionId, end_ms: i64) { + let mut panes = self.panes.write().expect("store lock"); + panes.retain(|(id, end), _| id != definition || *end >= end_ms); } pub fn len(&self) -> usize { - self.panes.len() + self.panes.read().expect("store lock").len() } pub fn is_empty(&self) -> bool { - self.panes.is_empty() + self.len() == 0 } } diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs index 28f7dd63..8f03ce2f 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -401,3 +401,28 @@ async fn forwards_unplanned_and_invalid_range_queries_unchanged() { assert_eq!(response.body, json!({ "forwarded_range": expected })); } } + +/// Planned queries sent at once run in parallel on the query threads, each +/// reading the same stored panes, and every one answers as the same query +/// sent alone. +#[tokio::test] +async fn answers_concurrent_queries_like_one_at_a_time() { + let (_, app) = example4b(101 * MINUTE).await; + let times = [90 * MINUTE, 100 * MINUTE]; + let mut alone = Vec::new(); + for &t in × { + let response = instant(&app, P99, t).await; + assert_eq!(response.execution.as_deref(), Some("plan"), "{t}"); + alone.push(response.body); + } + // Polled together, the requests wait on the query threads at once. + let concurrent = futures::future::join_all((0..32).map(|i| { + let (app, t) = (app.clone(), times[i % times.len()]); + async move { (i, instant(&app, P99, t).await) } + })) + .await; + for (i, response) in concurrent { + assert_eq!(response.execution.as_deref(), Some("plan")); + assert_eq!(response.body, alone[i % times.len()]); + } +} diff --git a/data_plane/tests/e2e_examples.rs b/data_plane/tests/e2e_examples.rs index 6ea6c20f..92402644 100644 --- a/data_plane/tests/e2e_examples.rs +++ b/data_plane/tests/e2e_examples.rs @@ -80,7 +80,7 @@ fn reference( ) -> Vec { let sds = Sds::from_plan(plan).unwrap(); assert!(sds.definitions.is_empty()); - let mut engine = QueryEngine::new(plan, &sds); + let engine = QueryEngine::new(plan, &sds); let store = SummaryStore::new(); let (root, ranges) = engine.raw_ranges(query, time_ms, &store).unwrap(); let raw = ranges diff --git a/data_plane/tests/kept.rs b/data_plane/tests/kept.rs index e11281a8..94c2e44b 100644 --- a/data_plane/tests/kept.rs +++ b/data_plane/tests/kept.rs @@ -115,7 +115,7 @@ fn kept_panes_answer_like_the_reference_and_are_reused() { assert_eq!(first, answer(&mut reference, 10 * MINUTE, &all).0); assert_eq!(first.len(), 2); assert_eq!( - kept.kept_stats, + kept.kept.stats(), KeptResultStats { built: 5, reused: 0 @@ -134,7 +134,7 @@ fn kept_panes_answer_like_the_reference_and_are_reused() { assert_eq!(second, answer(&mut reference, 11 * MINUTE, &all).0); assert_ne!(second, first, "the answer moves with the window"); assert_eq!( - kept.kept_stats, + kept.kept.stats(), KeptResultStats { built: 6, reused: 4 @@ -146,14 +146,14 @@ fn kept_panes_answer_like_the_reference_and_are_reused() { assert!(ranges.is_empty()); assert_eq!(again, second); assert_eq!( - kept.kept_stats, + kept.kept.stats(), KeptResultStats { built: 6, reused: 9 } ); assert_eq!( - reference.kept_stats, + reference.kept.stats(), KeptResultStats::default(), "the reference keeps nothing" ); @@ -172,7 +172,7 @@ fn kept_panes_are_evicted_like_stored_panes() { let ends = |engine: &QueryEngine| { (0..=30) .map(|m| m * MINUTE) - .filter(|end| engine.kept.get(&definition, *end).is_some()) + .filter(|end| engine.kept.panes.get(&definition, *end).is_some()) .map(|end| end / MINUTE) .collect::>() }; @@ -195,7 +195,7 @@ fn kept_panes_are_evicted_like_stored_panes() { assert_eq!(ranges.len(), 5); assert_eq!(old.len(), 2); assert_eq!(ends(&kept), vec![16, 17, 18, 19, 20]); - assert_eq!(kept.kept.len(), 5); + assert_eq!(kept.kept.panes.len(), 5); } /// A pane is kept only once it is closed (its end at or before the @@ -213,7 +213,7 @@ fn panes_are_kept_only_once_closed() { let ends = |engine: &QueryEngine| { (0..=30) .map(|m| m * MINUTE) - .filter(|end| engine.kept.get(&definition, *end).is_some()) + .filter(|end| engine.kept.panes.get(&definition, *end).is_some()) .map(|end| end / MINUTE) .collect::>() }; @@ -254,5 +254,5 @@ fn panes_are_kept_only_once_closed() { // Without any sample there is no watermark, and nothing is kept. let mut fresh = engine(&plan); answer_at_watermark(&mut fresh, 10 * MINUTE, None, &all); - assert!(fresh.kept.is_empty()); + assert!(fresh.kept.panes.is_empty()); } diff --git a/data_plane/tests/precompute.rs b/data_plane/tests/precompute.rs index a09cc538..035b10ea 100644 --- a/data_plane/tests/precompute.rs +++ b/data_plane/tests/precompute.rs @@ -6,7 +6,8 @@ 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, RawSample}; +use data_plane::precompute::{pane_end, shard_of, Precompute, RawSample}; +use data_plane::query::QueryEngine; use data_plane::sds::Sds; const MINUTE: i64 = 60_000; @@ -44,17 +45,21 @@ fn builds_panes_after_the_first_as_the_watermark_passes() { let id = sds.definitions.keys().next().unwrap().clone(); // The first sample, at 1 min, opens pane (0, 10 min], which is skipped. engine.ingest(&samples("latency_ms", &["a", "b"], MINUTE, 35 * MINUTE)); - assert_eq!(engine.store.get(&id, PANE).map(|_| ()), None, "first pane"); + assert_eq!( + engine.store().get(&id, PANE).map(|_| ()), + None, + "first pane" + ); for end in [2 * PANE, 3 * PANE] { - let pane = engine.store.get(&id, end).expect("built pane"); + let pane = engine.store().get(&id, end).expect("built pane"); assert_eq!(pane.rows().len(), 2, "one state per series"); } assert_eq!( - engine.store.get(&id, 4 * PANE).map(|_| ()), + engine.store().get(&id, 4 * PANE).map(|_| ()), None, "open pane" ); - let stats = engine.stats; + let stats = engine.stats(); assert_eq!(stats.panes_built, 2); assert_eq!(stats.panes_failed, 0); // Samples in (0, 10 min] are dropped with the first pane. @@ -67,11 +72,11 @@ fn builds_panes_after_the_first_as_the_watermark_passes() { fn ignores_other_metrics_and_drops_late_samples() { let (mut engine, _) = precompute(common::plan("example4b", false)); engine.ingest(&samples("latency_ms", &["a"], 0, 25 * MINUTE)); - let before = engine.stats; + let before = engine.stats(); engine.ingest(&[sample("other", "a", 25 * MINUTE, 1.0)]); - assert_eq!(engine.stats, before); + assert_eq!(engine.stats(), before); engine.ingest(&[sample("latency_ms", "a", 15 * MINUTE, 1.0)]); - assert_eq!(engine.stats.dropped, before.dropped + 1); + assert_eq!(engine.stats().dropped, before.dropped + 1); } /// A pane without samples is stored empty, so a query can tell "no data" @@ -82,9 +87,9 @@ fn stores_empty_panes() { let id = sds.definitions.keys().next().unwrap().clone(); engine.ingest(&samples("latency_ms", &["a"], 0, 20 * MINUTE)); engine.ingest(&[sample("latency_ms", "a", 45 * MINUTE, 1.0)]); - assert_eq!(engine.store.get(&id, 2 * PANE).unwrap().rows().len(), 1); - assert_eq!(engine.store.get(&id, 3 * PANE).unwrap().rows().len(), 0); - assert_eq!(engine.store.get(&id, 4 * PANE).unwrap().rows().len(), 0); + assert_eq!(engine.store().get(&id, 2 * PANE).unwrap().rows().len(), 1); + assert_eq!(engine.store().get(&id, 3 * PANE).unwrap().rows().len(), 0); + assert_eq!(engine.store().get(&id, 4 * PANE).unwrap().rows().len(), 0); } /// The plan's label matchers select which series a pane summarizes. @@ -99,7 +104,7 @@ fn applies_label_matchers() { let (mut engine, sds) = precompute(plan_workload(&workload, &common::inputs(false)).unwrap()); let id = sds.definitions.keys().next().unwrap().clone(); engine.ingest(&samples("latency_ms", &["a", "b"], 0, 25 * MINUTE)); - assert_eq!(engine.store.get(&id, 2 * PANE).unwrap().rows().len(), 1); + assert_eq!(engine.store().get(&id, 2 * PANE).unwrap().rows().len(), 1); } /// Panes are `(end - width, end]`, like PromQL ranges: a sample at a @@ -126,11 +131,103 @@ fn keeps_panes_for_the_plans_retention() { let id = definition.id.clone(); engine.ingest(&samples("latency_ms", &["a"], MINUTE, 90 * MINUTE + 1)); assert!( - engine.store.get(&id, 4 * PANE).is_some(), + engine.store().get(&id, 4 * PANE).is_some(), "read at offset 50 min" ); assert!( - engine.store.get(&id, 3 * PANE).is_none(), + engine.store().get(&id, 3 * PANE).is_none(), "older than the retention" ); } + +/// Two job values whose series land on different shards of `shards`. +fn jobs_on_two_shards(shards: usize) -> (String, String) { + let shard = |job: &str| shard_of(&sample("latency_ms", job, 0, 0.0).labels, shards); + let first = "a".to_string(); + let other = (b'b'..=b'z') + .map(|c| (c as char).to_string()) + .find(|job| shard(job) != shard(&first)) + .expect("two shards"); + (first, other) +} + +/// With samples sharded by series, a pane closes only once every shard's +/// watermark has passed it: the minimum over the shards, not the newest +/// sample. +#[test] +fn a_pane_closes_once_every_shard_has_passed_it() { + let plan = common::plan("example4b", false); + let sds = Sds::from_plan(&plan).unwrap(); + let id = sds.definitions.keys().next().unwrap().clone(); + let mut engine = Precompute::with_shards(&plan, &sds, 0, 2).unwrap(); + let (ahead, behind) = jobs_on_two_shards(2); + engine.ingest(&samples( + "latency_ms", + &[&ahead, &behind], + MINUTE, + 15 * MINUTE, + )); + engine.ingest(&samples("latency_ms", &[&ahead], 15 * MINUTE, 35 * MINUTE)); + assert_eq!(engine.watermark_ms(), Some(15 * MINUTE - 15_000)); + assert!( + engine.store().get(&id, 2 * PANE).is_none(), + "one shard is at 14:45" + ); + engine.ingest(&samples("latency_ms", &[&behind], 15 * MINUTE, 35 * MINUTE)); + let pane = engine + .store() + .get(&id, 2 * PANE) + .expect("both shards passed 20 min"); + assert_eq!( + pane.rows().len(), + 2, + "one state per series, from both shards" + ); + assert!(engine.store().get(&id, 3 * PANE).is_some()); +} + +/// The planned queries answer the same from panes built on four shards as +/// from panes built on one, for per-series KLL panes (Example 4b) and for +/// counter rates and sums (Example 1), with groups both per series and +/// over all series. +#[test] +fn sharded_panes_answer_like_one_shard() { + for (example, time_ms) in [("example4b", 70 * MINUTE), ("example1", 30 * MINUTE)] { + let plan = common::plan(example, false); + let sds = Sds::from_plan(&plan).unwrap(); + assert!(!sds.definitions.is_empty(), "{example} stores panes"); + let metric = match example { + "example4b" => "latency_ms", + _ => "http_requests_total", + }; + 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, time_ms + 1); + let mut one = Precompute::new(&plan, &sds, 0).unwrap(); + let mut four = Precompute::with_shards(&plan, &sds, 0, 4).unwrap(); + one.ingest(&all); + four.ingest(&all); + assert_eq!(one.stats(), four.stats(), "{example}"); + let query = QueryEngine::new(&plan, &sds); + for planned in plan.queries.iter().map(|q| q.query.as_str()) { + let answer = |engine: &Precompute| { + let (root, ranges) = query.raw_ranges(planned, time_ms, engine.store()).unwrap(); + assert!(ranges.is_empty(), "{example}: answered from stored panes"); + let mut result = query + .evaluate( + root, + time_ms, + engine.store(), + engine.watermark_ms(), + &BTreeMap::new(), + ) + .unwrap(); + result.sort_by(|a, b| a.labels.cmp(&b.labels)); + result + }; + let single = answer(&one); + assert!(!single.is_empty(), "{example}: {planned}"); + assert_eq!(answer(&four), single, "{example}: {planned}"); + } + } +} diff --git a/data_plane/tests/query.rs b/data_plane/tests/query.rs index 8f80330b..900b6456 100644 --- a/data_plane/tests/query.rs +++ b/data_plane/tests/query.rs @@ -78,14 +78,14 @@ fn fetch(ranges: &[RawRange], all: &[RawSample]) -> BTreeMap fn answer(f: &mut Fixture, query: &str, time_ms: i64, all: &[RawSample]) -> Vec { let (root, ranges) = f .query - .raw_ranges(query, time_ms, &f.precompute.store) + .raw_ranges(query, time_ms, f.precompute.store()) .unwrap(); let mut result = f .query .evaluate( root, time_ms, - &f.precompute.store, + f.precompute.store(), f.precompute.watermark_ms(), &fetch(&ranges, all), ) @@ -112,14 +112,14 @@ fn stored_panes_answer_like_a_single_build() { assert!( panes .query - .raw_ranges(P99, 70 * MINUTE, &panes.precompute.store) + .raw_ranges(P99, 70 * MINUTE, panes.precompute.store()) .is_err(), "no panes before ingestion" ); panes.precompute.ingest(&all); let (_, ranges) = panes .query - .raw_ranges(P99, 70 * MINUTE, &panes.precompute.store) + .raw_ranges(P99, 70 * MINUTE, panes.precompute.store()) .unwrap(); assert!(ranges.is_empty(), "every input is a stored pane"); let from_panes = answer(&mut panes, P99, 70 * MINUTE, &all); @@ -127,7 +127,7 @@ fn stored_panes_answer_like_a_single_build() { let mut single = fixture(common::plan("example4b", true)); let (_, ranges) = single .query - .raw_ranges(P99, 70 * MINUTE, &single.precompute.store) + .raw_ranges(P99, 70 * MINUTE, single.precompute.store()) .unwrap(); assert_eq!( ranges @@ -194,7 +194,7 @@ fn query_time_plans_read_raw_ranges() { #[test] fn reports_what_it_cannot_answer() { let mut f = fixture(common::plan("example4b", false)); - let store = &f.precompute.store; + let store = f.precompute.store(); assert!(matches!( f.query.raw_ranges("up", 70 * MINUTE, store), Err(Unanswered::NotPlanned(_)) @@ -205,7 +205,7 @@ fn reports_what_it_cannot_answer() { )); let all = samples(&[("latency_ms", &[])], MINUTE, 75 * MINUTE, |_, t| t as f64); f.precompute.ingest(&all); - let store = &f.precompute.store; + let store = f.precompute.store(); assert!(f.query.raw_ranges(P99, 70 * MINUTE, store).is_ok()); assert!(matches!( f.query.raw_ranges(P99, 70 * MINUTE + 1, store), diff --git a/data_plane/tests/range.rs b/data_plane/tests/range.rs index 311fe0e3..ca1829f2 100644 --- a/data_plane/tests/range.rs +++ b/data_plane/tests/range.rs @@ -195,7 +195,7 @@ fn every_step_reads_stored_panes() { (i * 1000) as f64 + (t / 15_000) as f64 }); precompute.ingest(&all); - let store = &precompute.store; + let store = precompute.store(); let mut engine = QueryEngine::new(&plan, &sds); let steps = minutes(9..=10) .into_iter() @@ -272,7 +272,7 @@ fn kept_panes_are_shared_across_steps() { // Panes ending at 6..14 min are built once; each step after the first // reuses four of its five. assert_eq!( - engine.kept_stats, + engine.kept.stats(), KeptResultStats { built: 9, reused: 16 @@ -293,7 +293,7 @@ fn kept_panes_are_shared_across_steps() { // minus 5), so each step reading one rebuilds it: 4 at 10 min, 3 at 11, // 2 at 12, 1 at 13. assert_eq!( - engine.kept_stats, + engine.kept.stats(), KeptResultStats { built: 19, reused: 31 @@ -331,7 +331,7 @@ fn open_panes_are_recomputed_at_each_step() { // Panes 6..12 are built once and kept; 13 is built at 13 and 14 min, 14 // at 14 min. Reused: four panes at 11, 12 and 13 min, three at 14. assert_eq!( - engine.kept_stats, + engine.kept.stats(), KeptResultStats { built: 10, reused: 15 @@ -339,7 +339,7 @@ fn open_panes_are_recomputed_at_each_step() { ); let definition = sds(&plan).kept_definitions.keys().next().unwrap().clone(); let kept_ends: Vec = (0..=30) - .filter(|m| engine.kept.get(&definition, m * MINUTE).is_some()) + .filter(|m| engine.kept.panes.get(&definition, m * MINUTE).is_some()) .collect(); // 6..9 are evicted: they start before the newest step (14 min) minus // the plan's 5-minute retention. diff --git a/data_plane/tests/sds.rs b/data_plane/tests/sds.rs index 037e38a6..044f9466 100644 --- a/data_plane/tests/sds.rs +++ b/data_plane/tests/sds.rs @@ -149,16 +149,76 @@ fn kept_panes_share_one_definition() { #[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(&a, end, format!("a{end}")); } store.insert(&b, 60, "b60".to_string()); - assert_eq!(store.get(&a, 120).map(String::as_str), Some("a120")); + assert_eq!( + store.get(&a, 120).as_deref().map(String::as_str), + Some("a120") + ); assert_eq!(store.get(&a, 90), None); store.evict_before(&a, 180); assert_eq!(store.get(&a, 120), None); - assert_eq!(store.get(&a, 180).map(String::as_str), Some("a180")); - assert_eq!(store.get(&b, 60).map(String::as_str), Some("b60")); + assert_eq!( + store.get(&a, 180).as_deref().map(String::as_str), + Some("a180") + ); + assert_eq!( + store.get(&b, 60).as_deref().map(String::as_str), + Some("b60") + ); assert_eq!(store.len(), 2); } + +/// A read shares the stored pane instead of copying it, and a pane never +/// changes once stored: a second insert for the same pane keeps the first. +#[test] +fn store_shares_panes_and_never_rewrites_them() { + let a = DefinitionId("a".into()); + let store = SummaryStore::new(); + let stored = store.insert(&a, 60, vec![1.0, 2.0]); + let first = store.get(&a, 60).unwrap(); + assert!(std::sync::Arc::ptr_eq(&stored, &first)); + let again = store.insert(&a, 60, vec![9.0]); + assert!( + std::sync::Arc::ptr_eq(&again, &first), + "the first pane stays" + ); + assert_eq!(*store.get(&a, 60).unwrap(), vec![1.0, 2.0]); +} + +/// Readers on many threads read the store while a writer inserts and +/// evicts, each read seeing a whole pane. +#[test] +fn store_is_read_concurrently() { + let a = DefinitionId("a".into()); + let store = std::sync::Arc::new(SummaryStore::new()); + let writer = { + let (store, a) = (store.clone(), a.clone()); + std::thread::spawn(move || { + for end in 0..1_000 { + store.insert(&a, end, vec![end; 16]); + store.evict_before(&a, end - 10); + } + }) + }; + let readers: Vec<_> = (0..4) + .map(|_| { + let (store, a) = (store.clone(), a.clone()); + std::thread::spawn(move || { + for end in 0..1_000 { + if let Some(pane) = store.get(&a, end) { + assert_eq!(*pane, vec![end; 16]); + } + } + }) + }) + .collect(); + writer.join().unwrap(); + for reader in readers { + reader.join().unwrap(); + } + assert_eq!(store.len(), 11); +} From 1b5d7a9a2a257bd5e6075cf44cfce30dffc08fc6 Mon Sep 17 00:00:00 2001 From: zzylol Date: Mon, 19 Oct 2026 07:11:47 +0000 Subject: [PATCH 2/2] fix(data-plane): a query sees the writes before it; a worker closes the watermark it announces Two races of the sharded engine, both found by a test that writes a pane's samples and queries it at once: - A planned query ran on the query pool while the ingestion workers could still be ingesting the Remote Write requests answered before it, so the newest pane was missing and the query was forwarded. A query now waits until every worker has run what was queued before it (`sync`), as the single engine thread's queue ordered it. - A worker closed its panes at the watermark it read before another worker moved it, then announced the newer watermark to the other workers but not to itself, so its part of the newest pane never arrived. It now closes and announces the same watermark. --- data_plane/src/engine.rs | 14 +++++++-- data_plane/src/precompute.rs | 23 ++++++++------ data_plane/tests/answers.rs | 61 ++++++++++++++++++++++++++++++++++++ 3 files changed, 86 insertions(+), 12 deletions(-) diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 8150a63e..3a657b81 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -239,12 +239,14 @@ impl EngineHandle { } /// The plan root answering `query` at `time_ms` and the raw samples it - /// needs; see [`QueryEngine::raw_ranges`]. + /// needs, once the writes answered before are ingested; see + /// [`QueryEngine::raw_ranges`]. pub async fn raw_ranges( &self, query: String, time_ms: i64, ) -> Result<(usize, Vec), Unanswered> { + self.sync().await; self.ask(|reply| QueryJob::RawRanges { query, time_ms, @@ -276,6 +278,7 @@ impl EngineHandle { query: String, steps: Vec, ) -> Result<(usize, Vec), Unanswered> { + self.sync().await; self.ask(|reply| QueryJob::RawRangesForSteps { query, steps, @@ -311,6 +314,14 @@ impl EngineHandle { /// Precompute counters, once every sample queued before has been /// ingested and every pane it closed has been built. pub async fn stats(&self) -> PrecomputeStats { + self.sync().await; + self.precompute.stats() + } + + /// Returns once every sample queued before has been ingested and every + /// pane it closed has been built, so a query sees the writes answered + /// before it. + pub async fn sync(&self) { // The first round runs every queued ingest, which may tell other // workers to advance; the second runs those advances. for _ in 0..2 { @@ -325,7 +336,6 @@ impl EngineHandle { let _ = done.await; } } - self.precompute.stats() } } diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index 4ac4fd56..656df256 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -384,13 +384,13 @@ impl ShardWorker { .max() .expect("nonempty"); shared.newest_sample_ms[self.index].fetch_max(newest, Ordering::AcqRel); - self.advance(); - match shared.watermark_ms() { - Some(watermark) => { - shared.announced_ms.fetch_max(watermark, Ordering::AcqRel) < watermark - } - None => false, - } + // Close and announce the same watermark: announced, it is closed + // here, and the other workers close it on `advance`. + let Some(watermark) = shared.watermark_ms() else { + return false; + }; + self.advance_to(watermark); + shared.announced_ms.fetch_max(watermark, Ordering::AcqRel) < watermark } /// Builds this worker's part of every pane the global watermark has @@ -398,9 +398,12 @@ impl ShardWorker { /// pane after the first, so a pane without samples is stored as empty, /// not missing. pub fn advance(&mut self) { - let Some(watermark) = self.shared.watermark_ms() else { - return; - }; + if let Some(watermark) = self.shared.watermark_ms() { + self.advance_to(watermark); + } + } + + fn advance_to(&mut self, watermark: i64) { if self .closed_through_ms .is_some_and(|closed| watermark <= closed) diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs index 8f03ce2f..9b50c4e4 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -426,3 +426,64 @@ async fn answers_concurrent_queries_like_one_at_a_time() { assert_eq!(response.body, alone[i % times.len()]); } } + +/// A planned query sees every Remote Write request answered before it: the +/// pane its newest sample closes is stored by the time the query reads it, +/// although ingestion runs on other threads. +#[tokio::test] +async fn a_query_reads_the_panes_of_the_writes_before_it() { + let workload = common::with_range_evaluation(&common::workload("example4b"), 600_000, 600_000); + let plan = control_plane::plan_workload(&workload, &common::inputs(false)).unwrap(); + let state = Arc::new( + DataPlane::new( + serde_json::to_string(&plan).unwrap(), + fake_prometheus(vec![]).await, + 0, + ) + .unwrap(), + ); + let app = router(state.clone()); + // One request per 10-minute pane, then at once the query that reads it. + for pane in 1..=12 { + let timeseries = (0..50) + .map(|i| TimeSeries { + labels: vec![ + Label { + name: "__name__".into(), + value: "latency_ms".into(), + }, + Label { + name: "job".into(), + value: format!("job{i}"), + }, + ], + samples: ((pane - 1) * 10 * MINUTE + 15_000..=pane * 10 * MINUTE) + .step_by(15_000) + .map(|t| Sample { + timestamp: t, + value: (t / 15_000) 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); + if pane >= 8 { + let response = instant(&app, P99, pane * 10 * MINUTE).await; + assert_eq!( + response.execution.as_deref(), + Some("plan"), + "pane {pane}: {:?}", + response.reason + ); + } + } +}