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..3a657b81 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,49 +69,185 @@ 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 - /// 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.ask(|reply| Command::RawRanges { + self.sync().await; + self.ask(|reply| QueryJob::RawRanges { query, time_ms, reply, @@ -105,7 +262,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 +278,8 @@ impl EngineHandle { query: String, steps: Vec, ) -> Result<(usize, Vec), Unanswered> { - self.ask(|reply| Command::RawRangesForSteps { + self.sync().await; + self.ask(|reply| QueryJob::RawRangesForSteps { query, steps, reply, @@ -129,14 +287,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 +304,82 @@ 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(); - } - stats.await.unwrap_or_default() + self.sync().await; + 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, - }) + /// 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 { + 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; + } + } } +} - 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..656df256 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,51 @@ impl Precompute { .entry(end) .or_default() .push(sample.to_sample()); - accepted = true; + buffered = true; } - if accepted { - self.stats.accepted += 1; - } - self.newest_sample_ms = self.newest_sample_ms.max(Some(sample.timestamp_ms)); + 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); + // 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 + /// 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) { + if let Some(watermark) = self.shared.watermark_ms() { + self.advance_to(watermark); } - self.close_panes(); } - fn close_panes(&mut self) { - let Some(watermark) = self.watermark_ms() else { + fn advance_to(&mut self, watermark: i64) { + 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 +425,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..9b50c4e4 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -401,3 +401,89 @@ 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()]); + } +} + +/// 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 + ); + } + } +} 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); +}