From cf53170b51381ce26eef83bc063f1b5deb404276 Mon Sep 17 00:00:00 2001 From: zzylol Date: Mon, 19 Oct 2026 06:52:05 +0000 Subject: [PATCH 1/2] feat(data-plane): read query-time raw windows from a local raw buffer and Prometheus remote read - A local raw buffer keeps the recent samples Remote Write delivered, sharded like ingestion: each ingestion worker appends its shard. It keeps the plan's longest query-time raw read back from the newest sample: the largest `end_offset + width` of `query_time_raw_windows` and of the panes built at query time and kept (their retention, which covers range evaluations). `--raw-buffer-ms` overrides it; 0 keeps none. - A raw window is read from the buffer; the part before the buffer's cut (its first sample, or the trimmed horizon of any shard) is read with Prometheus remote read (protobuf, `/api/v1/read`). The JSON query API is no longer used for raw samples. Forwarding of unplanned queries is unchanged. - The data plane logs the buffer's series and samples every 30 s. --- README.md | 9 +- data_plane/src/engine.rs | 41 +++++-- data_plane/src/lib.rs | 44 +++++-- data_plane/src/main.rs | 25 +++- data_plane/src/precompute.rs | 15 +++ data_plane/src/prometheus.rs | 71 +----------- data_plane/src/raw_buffer.rs | 193 +++++++++++++++++++++++++++++++ data_plane/src/remote_read.rs | 129 +++++++++++++++++++++ data_plane/tests/answers.rs | 80 ++++++++----- data_plane/tests/e2e_examples.rs | 158 ++++++++++++++++++++++--- data_plane/tests/raw_buffer.rs | 103 +++++++++++++++++ 11 files changed, 731 insertions(+), 137 deletions(-) create mode 100644 data_plane/src/raw_buffer.rs create mode 100644 data_plane/src/remote_read.rs create mode 100644 data_plane/tests/raw_buffer.rs diff --git a/README.md b/README.md index 8c72d1a3..ab3d92f2 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,11 @@ buffered samples instead. 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 `T - offset`, so `T` must be on the pane grid and those panes stored; its -query-time raw ranges are read from Prometheus. An input the plan keeps +query-time raw windows are read from the local raw buffer: the recent +samples Remote Write delivered, sharded like ingestion and kept for the +plan's longest query-time raw read back from the newest sample +(`--raw-buffer-ms` overrides it; 0 keeps none). The part of a window older +than the buffer is read with Prometheus remote read (`/api/v1/read`). An input the plan keeps across evaluations (`MatAtQueryKeptAcrossEvaluations`) is computed from its pane's raw samples when an evaluation reads it and no evaluation has kept it. It is kept in memory, by definition and pane end, once the pane is @@ -75,8 +79,7 @@ be answered also says why in `x-asap-forward-reason`. A range query (`/api/v1/query_range`) whose text is one of the plan's queries is answered as one evaluation per step, `start + k * step` through `end`, -and returns a Prometheus matrix. Each query-time raw range is read from -Prometheus once for all steps and sliced per step; stored and kept panes are +and returns a Prometheus matrix. Each query-time raw range is read once for all steps and sliced per step; stored and kept panes are read per step, so a kept pane built at one step is reused by the later ones. If any step cannot be answered, for example because a step is off the pane grid or reads an evicted pane, the whole range query is forwarded with the diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 7353efcf..8fc2e487 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -24,6 +24,7 @@ use tokio::sync::{mpsc, oneshot}; use crate::precompute::{PrecomputeStats, RawSample, ShardWorker, Shared}; use crate::query::{KeptResults, QueryEngine, RawRange, ResultSample, Unanswered}; +use crate::raw_buffer::RawBuffer; use crate::sds::Sds; /// Messages queued for each ingestion worker, and queries for the pool, @@ -69,22 +70,28 @@ enum QueryJob { #[error("the engine is busy")] pub struct Busy; -/// How many threads the engine runs. +/// How the engine runs: its threads and its local raw buffer. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub struct Threads { +pub struct EngineOptions { /// Ingestion workers; series are sharded across them. pub ingestion_workers: usize, /// Query threads; each runs one query at a time. pub query_threads: usize, + /// How far back from the newest sample the local raw buffer keeps + /// samples, in ms; `None` for the plan's longest query-time raw read + /// ([`RawBuffer::horizon_ms`]). 0 keeps none: every raw read is a + /// Prometheus remote read. + pub raw_buffer_ms: Option, } -impl Default for Threads { - /// One of each per available core. +impl Default for EngineOptions { + /// One thread of each kind per available core, and the plan's buffer. fn default() -> Self { let cores = std::thread::available_parallelism().map_or(1, |n| n.get()); Self { ingestion_workers: cores, query_threads: cores, + raw_buffer_ms: None, } } } @@ -95,6 +102,7 @@ pub struct EngineHandle { workers: Vec>, queries: SyncSender, precompute: Arc, + raw_buffer: Option>, } /// Reads and validates the plan serialized in `plan_json`, and its stored @@ -125,8 +133,9 @@ impl EngineHandle { pub fn spawn( plan_json: String, allowed_lateness_ms: i64, - threads: Threads, + options: EngineOptions, ) -> anyhow::Result { + let threads = options; let (plan, sds) = install(&plan_json)?; tracing::info!( definitions = sds.definitions.len(), @@ -136,7 +145,13 @@ impl EngineHandle { ?threads, "stored summaries" ); + let raw_buffer_ms = options + .raw_buffer_ms + .unwrap_or_else(|| RawBuffer::horizon_ms(&plan)); + tracing::info!(raw_buffer_ms, "local raw buffer"); drop(plan); + let raw_buffer = (raw_buffer_ms > 0) + .then(|| Arc::new(RawBuffer::new(raw_buffer_ms, threads.ingestion_workers))); 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()) @@ -146,10 +161,16 @@ impl EngineHandle { // 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()); + let (plan_json, shared, peers, raw_buffer) = ( + plan_json.clone(), + precompute.clone(), + peers.clone(), + raw_buffer.clone(), + ); spawn(format!("ingest-{index}"), move |ready| { let started = install(&plan_json).and_then(|(plan, sds)| { - Ok(ShardWorker::new(&plan, &sds, index, shared, true)?) + Ok(ShardWorker::new(&plan, &sds, index, shared, true)? + .with_raw_buffer(raw_buffer)) }); let mut worker = match started { Ok(worker) => { @@ -210,9 +231,15 @@ impl EngineHandle { workers: senders, queries, precompute, + raw_buffer, }) } + /// The local raw buffer, if the plan reads raw samples at query time. + pub fn raw_buffer(&self) -> Option<&RawBuffer> { + self.raw_buffer.as_deref() + } + /// How many ingestion workers series are sharded across. pub fn ingestion_workers(&self) -> usize { self.workers.len() diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 4eb8ba32..3a008b9b 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -13,6 +13,8 @@ mod per_sample; pub mod precompute; pub mod prometheus; pub mod query; +pub mod raw_buffer; +pub mod remote_read; pub mod remote_write; pub mod sds; @@ -27,7 +29,7 @@ use axum::http::{header, HeaderMap, HeaderValue, Method, StatusCode}; use axum::response::{IntoResponse, Response}; use axum::routing::{get, post}; use axum::Router; -use engine::{EngineHandle, Threads}; +use engine::{EngineHandle, EngineOptions}; use precompute::{shard_of, RawSample}; use prometheus::{InstantQuery, RangeQuery}; use query::Unanswered; @@ -54,7 +56,7 @@ pub struct DataPlane { impl DataPlane { /// Installs the plan serialized in `plan_json` (an [`InstallablePlan`]) - /// with [`Threads::default`]. Panes close `allowed_lateness_ms` after + /// with [`EngineOptions::default`]. Panes close `allowed_lateness_ms` after /// every ingestion worker's newest sample passes them. pub fn new( plan_json: String, @@ -65,22 +67,22 @@ impl DataPlane { plan_json, prometheus_url, allowed_lateness_ms, - Threads::default(), + EngineOptions::default(), ) } - /// [`Self::new`] with `threads` ingestion workers and query threads. + /// [`Self::new`] with `options`' threads and raw buffer. pub fn with_threads( plan_json: String, prometheus_url: String, allowed_lateness_ms: i64, - threads: Threads, + options: EngineOptions, ) -> 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, threads)?, + engine: EngineHandle::spawn(plan_json, allowed_lateness_ms, options)?, prometheus_url: prometheus_url.trim_end_matches('/').to_string(), limits: remote_write::Limits::default(), samples: AtomicU64::new(0), @@ -235,16 +237,40 @@ async fn answer_range( Ok(prometheus::matrix_response(&steps, &results)) } -/// The samples of each of `ranges` from Prometheus, by node. +/// The samples of each of `ranges`, by node: from the local raw buffer, and +/// before it from Prometheus remote read. async fn fetch( state: &DataPlane, ranges: &[query::RawRange], ) -> Result>, Unanswered> { let mut raw = BTreeMap::new(); for range in ranges { - let samples = prometheus::fetch(&state.client, &state.prometheus_url, range) + // The local raw buffer holds the recent part; anything before it + // comes from Prometheus remote read. + let (mut samples, local) = match state.engine.raw_buffer() { + Some(buffer) => { + let read = buffer.read(&range.metric, range.start_ms, range.end_ms); + (Vec::new(), Some(read)) + } + None => (Vec::new(), None), + }; + let older_than_ms = local + .as_ref() + .map_or(range.end_ms, |read| read.older_than_ms); + if older_than_ms > range.start_ms { + samples = remote_read::read( + &state.client, + &state.prometheus_url, + &range.metric, + range.start_ms, + older_than_ms, + ) .await - .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))?; + .map_err(|e| Unanswered::Execution(format!("reading raw samples: {e}")))?; + } + if let Some(read) = local { + samples.extend(read.samples); + } raw.insert(range.node, samples); } Ok(raw) diff --git a/data_plane/src/main.rs b/data_plane/src/main.rs index 990f4283..24ed4838 100644 --- a/data_plane/src/main.rs +++ b/data_plane/src/main.rs @@ -2,7 +2,7 @@ use std::path::PathBuf; use std::sync::Arc; use clap::Parser; -use data_plane::engine::Threads; +use data_plane::engine::EngineOptions; use data_plane::{router, DataPlane}; // jemalloc, as ASAPQuery's engine uses: the executor allocates per row, and @@ -34,6 +34,11 @@ struct Args { /// Query threads, each running one query at a time; default one per core. #[arg(long)] query_threads: Option, + /// How far back from the newest sample to keep raw samples for + /// query-time raw windows, in ms; default the plan's longest such read. + /// 0 reads every raw window from Prometheus remote read. + #[arg(long)] + raw_buffer_ms: Option, } #[tokio::main] @@ -44,18 +49,30 @@ 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 defaults = Threads::default(); - let threads = Threads { + let defaults = EngineOptions::default(); + let options = EngineOptions { ingestion_workers: args.ingestion_workers.unwrap_or(defaults.ingestion_workers), query_threads: args.query_threads.unwrap_or(defaults.query_threads), + raw_buffer_ms: args.raw_buffer_ms, }; let state = Arc::new(DataPlane::with_threads( plan, args.prometheus_url, args.allowed_lateness_ms, - threads, + options, )?); tracing::info!(queries = state.queries.len(), "installed plan"); + let reporting = state.clone(); + tokio::spawn(async move { + let mut every = tokio::time::interval(std::time::Duration::from_secs(30)); + loop { + every.tick().await; + if let Some(buffer) = reporting.engine.raw_buffer() { + let (series, samples) = buffer.size(); + tracing::info!(series, samples, "local raw buffer"); + } + } + }); let listener = tokio::net::TcpListener::bind(&args.listen).await?; tracing::info!(address = %listener.local_addr()?, "listening"); axum::serve(listener, router(state)).await?; diff --git a/data_plane/src/precompute.rs b/data_plane/src/precompute.rs index fe171eb3..82d28123 100644 --- a/data_plane/src/precompute.rs +++ b/data_plane/src/precompute.rs @@ -30,6 +30,7 @@ use futures::executor::block_on; use futures::StreamExt; use crate::per_sample::PerSampleBuild; +use crate::raw_buffer::RawBuffer; use crate::sds::{DefinitionId, Sds, SummaryDefinition, SummaryStore}; /// One raw sample: its series' full label set (including `__name__`), time @@ -312,6 +313,9 @@ pub struct ShardWorker { shared: Arc, /// The global watermark this worker has built its parts through. closed_through_ms: Option, + /// The local raw buffer, which keeps this worker's recent samples for + /// query-time raw windows; `None` when the plan reads none. + raw_buffer: Option>, } impl ShardWorker { @@ -350,9 +354,17 @@ impl ShardWorker { builders, shared, closed_through_ms: None, + raw_buffer: None, }) } + /// Also keeps this worker's samples in `raw_buffer`, its shard + /// [`Self::new`]'s `index`. + pub fn with_raw_buffer(mut self, raw_buffer: Option>) -> Self { + self.raw_buffer = raw_buffer; + self + } + /// 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 @@ -362,6 +374,9 @@ impl ShardWorker { let Some(first) = samples.iter().map(|s| s.timestamp_ms).min() else { return false; }; + if let Some(raw_buffer) = &self.raw_buffer { + raw_buffer.append(self.index, samples); + } let shared = self.shared.clone(); let _ = shared .start_ms diff --git a/data_plane/src/prometheus.rs b/data_plane/src/prometheus.rs index 31819026..33a9a449 100644 --- a/data_plane/src/prometheus.rs +++ b/data_plane/src/prometheus.rs @@ -1,14 +1,12 @@ //! The Prometheus HTTP API as the data plane speaks it: instant- and -//! range-query parameters and results, and fetching raw samples for -//! query-time ranges. +//! range-query parameters and results. Raw samples are read with remote +//! read ([`crate::remote_read`]). use std::collections::BTreeMap; -use std::sync::Arc; use serde_json::{json, Value}; -use crate::precompute::RawSample; -use crate::query::{RawRange, ResultSample}; +use crate::query::ResultSample; /// Parameters of an instant query, from its URL query and form body. pub struct InstantQuery { @@ -142,10 +140,6 @@ fn seconds_to_ms(seconds: f64) -> i64 { (seconds * 1000.0).round() as i64 } -fn ms_to_seconds(ms: i64) -> String { - format!("{:.3}", ms as f64 / 1000.0) -} - /// A Prometheus instant-query response for `samples` at `time_ms`. pub fn vector_response(time_ms: i64, samples: &[ResultSample]) -> Value { let time = time_ms as f64 / 1000.0; @@ -186,64 +180,6 @@ fn format_value(value: f64) -> String { } } -/// The samples of `range`'s metric in `(start, end]`, read from Prometheus -/// with a range selector evaluated at the range's end. -pub async fn fetch( - client: &reqwest::Client, - prometheus_url: &str, - range: &RawRange, -) -> anyhow::Result> { - let selector = format!( - "{{__name__=\"{}\"}}[{}ms]", - range.metric, - range.end_ms - range.start_ms - ); - let response: Value = client - .get(format!("{prometheus_url}/api/v1/query")) - .query(&[ - ("query", selector.as_str()), - ("time", &ms_to_seconds(range.end_ms)), - ]) - .send() - .await? - .error_for_status()? - .json() - .await?; - let malformed = || anyhow::anyhow!("unexpected Prometheus response: {response}"); - let series = response["data"]["result"] - .as_array() - .ok_or_else(malformed)?; - let mut samples = Vec::new(); - for series in series { - let labels: BTreeMap = serde_json::from_value(series["metric"].clone())?; - let labels = Arc::new(labels); - for point in series["values"].as_array().ok_or_else(malformed)? { - let seconds = point[0].as_f64().ok_or_else(malformed)?; - let value = point[1].as_str().ok_or_else(malformed)?; - let timestamp_ms = seconds_to_ms(seconds); - // Prometheus 2 ranges include their start; the plan's do not. - if timestamp_ms <= range.start_ms { - continue; - } - samples.push(RawSample { - labels: labels.clone(), - timestamp_ms, - value: parse_value(value).ok_or_else(malformed)?, - }); - } - } - Ok(samples) -} - -fn parse_value(value: &str) -> Option { - match value { - "+Inf" | "Inf" => Some(f64::INFINITY), - "-Inf" => Some(f64::NEG_INFINITY), - "NaN" => Some(f64::NAN), - v => v.parse().ok(), - } -} - #[cfg(test)] mod tests { use super::*; @@ -285,7 +221,6 @@ mod tests { {"metric": {}, "value": [1.5, "+Inf"]}, ]}}) ); - assert_eq!(ms_to_seconds(4_200_001), "4200.001"); } fn range(params: &str) -> Option<(i64, i64, i64)> { diff --git a/data_plane/src/raw_buffer.rs b/data_plane/src/raw_buffer.rs new file mode 100644 index 00000000..09035bbe --- /dev/null +++ b/data_plane/src/raw_buffer.rs @@ -0,0 +1,193 @@ +//! The local raw buffer: recent raw samples, from which queries read their +//! query-time raw windows instead of asking Prometheus. +//! +//! It is fed by Remote Write through the ingestion workers and sharded like +//! them: each worker appends its shard's samples. It keeps the samples of +//! the plan's longest query-time read back from the newest sample (its +//! horizon); a read further back takes the older part from Prometheus +//! remote read. + +use std::collections::{BTreeMap, HashMap, VecDeque}; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, RwLock}; + +use asap_planner::{InstallablePlan, StoredOutputBuild}; + +use crate::precompute::RawSample; + +type Labels = Arc>; + +/// No value yet in an `AtomicI64` slot. +const UNSET: i64 = i64::MIN; + +#[derive(Default)] +struct Shard { + /// Each series' samples in time order, by metric and label set. + series: HashMap>>, + /// Samples at or before this time may have been dropped. + trimmed_through_ms: i64, + samples: usize, +} + +pub struct RawBuffer { + horizon_ms: i64, + shards: Vec>, + /// The first sample's time: the buffer holds every sample after it + /// that Remote Write delivered, until it is trimmed. + start_ms: AtomicI64, + newest_ms: AtomicI64, +} + +/// What a buffer read gives: the samples it holds of the range, and the +/// time before which the range must be read from Prometheus instead. +pub struct BufferRead { + /// Samples in `(older_than_ms, end]`, each series' together in order. + pub samples: Vec, + /// The buffer holds every sample after this time (at least the range's + /// start). + pub older_than_ms: i64, +} + +impl RawBuffer { + /// The longest a query of `plan` reads raw samples back from its + /// evaluation time: over its query-time raw windows and the panes it + /// builds at query time and keeps, `end_offset + width`, and for the + /// latter their retention, which covers their range evaluations. + pub fn horizon_ms(plan: &InstallablePlan) -> i64 { + let windows = plan + .query_time_raw_windows + .iter() + .map(|w| w.end_offset_ms + w.width_ms); + let kept = plan + .stored_outputs + .iter() + .filter(|o| o.built_at == StoredOutputBuild::QueryRun) + .map(|o| (o.pane.reach_ms().max(o.retention_ms)) as i64); + windows.chain(kept).max().unwrap_or(0) + } + + pub fn new(horizon_ms: i64, shards: usize) -> Self { + Self { + horizon_ms, + shards: (0..shards.max(1)) + .map(|_| { + RwLock::new(Shard { + trimmed_through_ms: UNSET, + ..Shard::default() + }) + }) + .collect(), + start_ms: AtomicI64::new(UNSET), + newest_ms: AtomicI64::new(UNSET), + } + } + + /// Appends `samples`, all of shard `shard`, then drops its samples older + /// than the horizon before the newest sample. + pub fn append(&self, shard: usize, samples: &[RawSample]) { + let (Some(first), Some(newest)) = ( + samples.iter().map(|s| s.timestamp_ms).min(), + samples.iter().map(|s| s.timestamp_ms).max(), + ) else { + return; + }; + let _ = self + .start_ms + .compare_exchange(UNSET, first, Ordering::AcqRel, Ordering::Acquire); + let newest = self + .newest_ms + .fetch_max(newest, Ordering::AcqRel) + .max(newest); + let trim_through = newest - self.horizon_ms; + let mut shard = self.shards[shard].write().expect("raw buffer lock"); + let shard = &mut *shard; + let mut added = 0usize; + // Samples of one series arrive together: look each series up once. + for run in samples.chunk_by(|a, b| Arc::ptr_eq(&a.labels, &b.labels)) { + let labels = &run[0].labels; + let metric = labels.get("__name__").cloned().unwrap_or_default(); + let series = shard + .series + .entry(metric) + .or_default() + .entry(labels.clone()) + .or_default(); + for sample in run.iter().filter(|s| s.timestamp_ms > trim_through) { + let point = (sample.timestamp_ms, sample.value); + match series.back() { + Some(&(time, _)) if time >= point.0 => { + let at = series.partition_point(|&(t, _)| t < point.0); + if series.get(at).is_some_and(|&(t, _)| t == point.0) { + continue; + } + series.insert(at, point); + } + _ => series.push_back(point), + } + added += 1; + } + } + shard.samples += added; + if trim_through > shard.trimmed_through_ms { + shard.trimmed_through_ms = trim_through; + let mut dropped = 0; + for series in shard.series.values_mut() { + series.retain(|_, points| { + let keep = points.partition_point(|&(t, _)| t <= trim_through); + dropped += keep; + points.drain(..keep); + !points.is_empty() + }); + } + shard.samples -= dropped; + } + } + + /// `metric`'s samples in `(start_ms, end_ms]` that the buffer holds, + /// and from when it holds all of them. + pub fn read(&self, metric: &str, start_ms: i64, end_ms: i64) -> BufferRead { + let first = self.start_ms.load(Ordering::Acquire); + let mut older_than_ms = start_ms.max(if first == UNSET { end_ms } else { first }); + let mut parts = Vec::with_capacity(self.shards.len()); + for shard in &self.shards { + let shard = shard.read().expect("raw buffer lock"); + older_than_ms = older_than_ms.max(shard.trimmed_through_ms); + let mut samples = Vec::new(); + for (labels, points) in shard.series.get(metric).into_iter().flatten() { + let from = points.partition_point(|&(t, _)| t <= start_ms); + let to = points.partition_point(|&(t, _)| t <= end_ms); + samples.extend( + points + .range(from..to) + .map(|&(timestamp_ms, value)| RawSample { + labels: labels.clone(), + timestamp_ms, + value, + }), + ); + } + parts.push(samples); + } + // A shard trimmed past another's cut may miss samples before it, so + // every shard's samples before the latest cut come from Prometheus. + let older_than_ms = older_than_ms.min(end_ms); + let samples = parts + .into_iter() + .flatten() + .filter(|s| s.timestamp_ms > older_than_ms) + .collect(); + BufferRead { + samples, + older_than_ms, + } + } + + /// Series and samples held, for reporting. + pub fn size(&self) -> (usize, usize) { + self.shards.iter().fold((0, 0), |(series, samples), shard| { + let shard = shard.read().expect("raw buffer lock"); + let n: usize = shard.series.values().map(HashMap::len).sum(); + (series + n, samples + shard.samples) + }) + } +} diff --git a/data_plane/src/remote_read.rs b/data_plane/src/remote_read.rs new file mode 100644 index 00000000..670458a6 --- /dev/null +++ b/data_plane/src/remote_read.rs @@ -0,0 +1,129 @@ +//! Prometheus remote read (protobuf, sampled responses): raw samples older +//! than the local raw buffer holds. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use prost::Message; + +use crate::precompute::RawSample; +use crate::remote_write::TimeSeries; + +#[derive(Clone, PartialEq, Message)] +pub struct ReadRequest { + #[prost(message, repeated, tag = "1")] + pub queries: Vec, + /// Response types the client accepts; `SAMPLES` (0) only. + #[prost(int32, repeated, tag = "2")] + pub accepted_response_types: Vec, +} + +#[derive(Clone, PartialEq, Message)] +pub struct Query { + /// Inclusive bounds, in milliseconds. + #[prost(int64, tag = "1")] + pub start_timestamp_ms: i64, + #[prost(int64, tag = "2")] + pub end_timestamp_ms: i64, + #[prost(message, repeated, tag = "3")] + pub matchers: Vec, +} + +#[derive(Clone, PartialEq, Message)] +pub struct LabelMatcher { + /// `EQ` is 0. + #[prost(int32, tag = "1")] + pub r#type: i32, + #[prost(string, tag = "2")] + pub name: String, + #[prost(string, tag = "3")] + pub value: String, +} + +#[derive(Clone, PartialEq, Message)] +pub struct ReadResponse { + #[prost(message, repeated, tag = "1")] + pub results: Vec, +} + +#[derive(Clone, PartialEq, Message)] +pub struct QueryResult { + #[prost(message, repeated, tag = "1")] + pub timeseries: Vec, +} + +/// The request for `metric`'s samples in `(start_ms, end_ms]`. +pub fn request(metric: &str, start_ms: i64, end_ms: i64) -> ReadRequest { + ReadRequest { + queries: vec![Query { + // Remote read bounds are inclusive; the plan's ranges are + // left-open. + start_timestamp_ms: start_ms + 1, + end_timestamp_ms: end_ms, + matchers: vec![LabelMatcher { + r#type: 0, + name: "__name__".into(), + value: metric.into(), + }], + }], + accepted_response_types: vec![0], + } +} + +/// Snappy-compressed protobuf, as remote read sends both ways. +pub fn encode(message: &impl Message) -> Vec { + snap::raw::Encoder::new() + .compress_vec(&message.encode_to_vec()) + .expect("snappy compresses any input") +} + +/// The message in a snappy-compressed protobuf body. +pub fn decode(body: &[u8]) -> anyhow::Result { + let decoded = snap::raw::Decoder::new().decompress_vec(body)?; + Ok(M::decode(decoded.as_slice())?) +} + +/// `metric`'s samples in `(start_ms, end_ms]` from Prometheus remote read, +/// each series' samples together and in time order. +pub async fn read( + client: &reqwest::Client, + prometheus_url: &str, + metric: &str, + start_ms: i64, + end_ms: i64, +) -> anyhow::Result> { + let body = client + .post(format!("{prometheus_url}/api/v1/read")) + .header("Content-Encoding", "snappy") + .header("Content-Type", "application/x-protobuf") + .header("X-Prometheus-Remote-Read-Version", "0.1.0") + .body(encode(&request(metric, start_ms, end_ms))) + .send() + .await? + .error_for_status()? + .bytes() + .await?; + let response: ReadResponse = decode(&body)?; + let mut samples = Vec::new(); + for series in response.results.into_iter().flat_map(|r| r.timeseries) { + let labels: Arc> = Arc::new( + series + .labels + .into_iter() + .map(|label| (label.name, label.value)) + .collect(), + ); + samples.extend( + series + .samples + .into_iter() + .filter(|s| start_ms < s.timestamp && s.timestamp <= end_ms) + .map(|s| RawSample { + labels: labels.clone(), + timestamp_ms: s.timestamp, + value: s.value, + }), + ); + } + Ok(samples) +} diff --git a/data_plane/tests/answers.rs b/data_plane/tests/answers.rs index 9b50c4e4..4b21f6d8 100644 --- a/data_plane/tests/answers.rs +++ b/data_plane/tests/answers.rs @@ -7,10 +7,12 @@ use std::collections::BTreeMap; use std::sync::Arc; use axum::body::Body; +use axum::body::Bytes; use axum::extract::Query; use axum::http::{Request, StatusCode}; use axum::routing::get; use axum::{Json, Router}; +use data_plane::remote_read::{self, QueryResult, ReadRequest, ReadResponse}; use data_plane::remote_write::{encode, Label, Sample, TimeSeries, WriteRequest}; use data_plane::{router, DataPlane, EXECUTION_HEADER, FORWARD_REASON_HEADER}; use http_body_util::BodyExt; @@ -28,49 +30,65 @@ fn series(metric: &str, labels: &[(&str, &str)]) -> BTreeMap { map } -/// A stand-in Prometheus holding `points`. It answers -/// `{__name__="m"}[Nms]` at `time` with the samples in `[time - N, time]`, -/// both ends included as Prometheus 2 does, and anything else with -/// `{"forwarded": query}`. Range queries are answered with +/// A stand-in Prometheus holding `points`. It answers remote read with each +/// metric's samples within the request's inclusive bounds, and any instant +/// query with `{"forwarded": query}`. Range queries are answered with /// `{"forwarded_range": parameters}`. async fn fake_prometheus(points: Vec) -> String { let points = Arc::new(points); - let handler = move |Query(params): Query>| { - let points = points.clone(); + let read_points = points.clone(); + let read = move |body: Bytes| { + let points = read_points.clone(); async move { - let query = ¶ms["query"]; - let parsed = query - .strip_prefix("{__name__=\"") - .and_then(|rest| rest.split_once("\"}[")) - .and_then(|(metric, rest)| { - Some((metric, rest.strip_suffix("ms]")?.parse::().ok()?)) - }); - let Some((metric, range)) = parsed else { - return Json(json!({ "forwarded": query })); - }; - let time = (params["time"].parse::().unwrap() * 1000.0).round() as i64; - let mut by_series: BTreeMap<&BTreeMap, Vec> = BTreeMap::new(); - for (labels, t, v) in points.iter() { - if labels["__name__"] == metric && time - range <= *t && *t <= time { - by_series - .entry(labels) - .or_default() - .push(json!([*t as f64 / 1000.0, v.to_string()])); - } - } - let result: Vec = by_series - .into_iter() - .map(|(labels, values)| json!({"metric": labels, "values": values})) + let request: ReadRequest = remote_read::decode(&body).unwrap(); + let results = request + .queries + .iter() + .map(|query| { + let metric = &query.matchers[0].value; + let mut by_series: BTreeMap<&BTreeMap, Vec> = + BTreeMap::new(); + for (labels, t, v) in points.iter() { + if &labels["__name__"] == metric + && query.start_timestamp_ms <= *t + && *t <= query.end_timestamp_ms + { + by_series.entry(labels).or_default().push(Sample { + value: *v, + timestamp: *t, + }); + } + } + QueryResult { + timeseries: by_series + .into_iter() + .map(|(labels, samples)| TimeSeries { + labels: labels + .iter() + .map(|(name, value)| Label { + name: name.clone(), + value: value.clone(), + }) + .collect(), + samples, + }) + .collect(), + } + }) .collect(); - Json(json!({"status": "success", "data": {"resultType": "matrix", "result": result}})) + remote_read::encode(&ReadResponse { results }) } }; + let handler = |Query(params): Query>| async move { + Json(json!({ "forwarded": params["query"] })) + }; let range = |Query(params): Query>| async move { Json(json!({ "forwarded_range": params })) }; let app = Router::new() .route("/api/v1/query", get(handler)) - .route("/api/v1/query_range", get(range)); + .route("/api/v1/query_range", get(range)) + .route("/api/v1/read", axum::routing::post(read)); let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let url = format!("http://{}", listener.local_addr().unwrap()); tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); diff --git a/data_plane/tests/e2e_examples.rs b/data_plane/tests/e2e_examples.rs index 92402644..600ce290 100644 --- a/data_plane/tests/e2e_examples.rs +++ b/data_plane/tests/e2e_examples.rs @@ -21,6 +21,7 @@ use axum::{Json, Router}; use control_plane::{plan_workload, save_plan, DeploymentInputs}; use data_plane::precompute::RawSample; use data_plane::query::{QueryEngine, ResultSample}; +use data_plane::remote_read::{self, QueryResult, ReadRequest, ReadResponse}; use data_plane::remote_write::{encode, Label, Sample, TimeSeries, WriteRequest}; use data_plane::sds::{Sds, SummaryStore}; use serde_json::{json, Value}; @@ -106,14 +107,25 @@ fn sorted(mut samples: Vec) -> Vec { samples } +/// A Prometheus holding `all`, and how many raw reads it has answered. +struct Prometheus { + url: String, + /// Range selectors read through the JSON query API. + json_reads: Arc, + /// Remote read requests. + remote_reads: Arc, +} + /// A Prometheus holding `all`: range selectors `{__name__="m"}[Nms]` at /// `time` return the samples in `[time - N, time]`, as Prometheus 2 does; +/// remote read returns each metric's samples in its inclusive bounds; /// every other query returns an empty vector, and every range query an -/// empty matrix. Returns its URL and how many range selectors it has -/// answered. -async fn prometheus(all: Arc>) -> (String, Arc) { +/// empty matrix. +async fn prometheus(all: Arc>) -> Prometheus { let range_reads = Arc::new(AtomicUsize::new(0)); + let remote_reads = Arc::new(AtomicUsize::new(0)); let reads = range_reads.clone(); + let all_for_read = all.clone(); let handler = move |Query(params): Query>| { let all = all.clone(); let reads = reads.clone(); @@ -157,11 +169,64 @@ async fn prometheus(all: Arc>) -> (String, Arc) { let empty_matrix = || async { Json(json!({"status": "success", "data": {"resultType": "matrix", "result": []}})) }; + let read = { + let (all, reads) = (all_for_read, remote_reads.clone()); + move |body: axum::body::Bytes| { + let (all, reads) = (all.clone(), reads.clone()); + async move { + reads.fetch_add(1, Ordering::SeqCst); + let request: ReadRequest = remote_read::decode(&body).unwrap(); + let results = request + .queries + .iter() + .map(|query| { + let [matcher] = query.matchers.as_slice() else { + panic!("one matcher: {query:?}") + }; + assert_eq!((matcher.r#type, matcher.name.as_str()), (0, "__name__")); + let mut by_series: BTreeMap<&BTreeMap, Vec> = + BTreeMap::new(); + for s in all.iter().filter(|s| { + s.labels["__name__"] == matcher.value + && query.start_timestamp_ms <= s.timestamp_ms + && s.timestamp_ms <= query.end_timestamp_ms + }) { + by_series.entry(&s.labels).or_default().push(Sample { + value: s.value, + timestamp: s.timestamp_ms, + }); + } + QueryResult { + timeseries: by_series + .into_iter() + .map(|(labels, samples)| TimeSeries { + labels: labels + .iter() + .map(|(name, value)| Label { + name: name.clone(), + value: value.clone(), + }) + .collect(), + samples, + }) + .collect(), + } + }) + .collect(); + remote_read::encode(&ReadResponse { results }) + } + } + }; let app = Router::new() .route("/api/v1/query", get(handler)) - .route("/api/v1/query_range", get(empty_matrix)); + .route("/api/v1/query_range", get(empty_matrix)) + .route("/api/v1/read", axum::routing::post(read)); tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); - (url, range_reads) + Prometheus { + url, + json_reads: range_reads, + remote_reads, + } } /// The `data_plane` binary serving `plan`, killed on drop. @@ -186,7 +251,8 @@ impl Drop for TempDir { } } -async fn start(plan: &InstallablePlan, prometheus_url: &str) -> DataPlaneProcess { +/// The `data_plane` binary serving `plan` with `args` added. +async fn start(plan: &InstallablePlan, prometheus_url: &str, args: &[&str]) -> DataPlaneProcess { let dir = std::env::temp_dir().join(format!( "data-plane-e2e-{}-{}", std::process::id(), @@ -205,6 +271,7 @@ async fn start(plan: &InstallablePlan, prometheus_url: &str) -> DataPlaneProcess .args(["--prometheus-url", prometheus_url]) .args(["--listen", &format!("127.0.0.1:{port}")]) .args(["--allowed-lateness-ms", "0"]) + .args(args) .spawn() .unwrap(); let process = DataPlaneProcess { @@ -426,7 +493,8 @@ async fn check_example( let plan = plan_workload(&workload, inputs).unwrap(); let at_query_time = plan.query_time_reference().unwrap(); let all = Arc::new(all); - let process = start(&plan, &prometheus(all.clone()).await.0).await; + let prometheus = prometheus(all.clone()).await; + let process = start(&plan, &prometheus.url, &[]).await; remote_write(&process.url, &all).await; for query in plan.queries.iter().map(|q| &q.query) { let expected = reference(&at_query_time, query, time_ms, &all); @@ -557,8 +625,10 @@ async fn example4_b3() { let sds = Sds::from_plan(&plan).unwrap(); assert_eq!((sds.definitions.len(), sds.kept_definitions.len()), (0, 1)); let all = Arc::new(all); - let (prometheus_url, range_reads) = prometheus(all.clone()).await; - let process = start(&plan, &prometheus_url).await; + // Without the local raw buffer, every raw read is a remote read. + let prometheus = prometheus(all.clone()).await; + let range_reads = prometheus.remote_reads.clone(); + let process = start(&plan, &prometheus.url, &["--raw-buffer-ms", "0"]).await; remote_write(&process.url, &all).await; let query = &plan.queries[0].query; let mut reads = Vec::new(); @@ -592,8 +662,10 @@ async fn example1_range() { let plan = plan_workload(&workload, &inputs).unwrap(); let at_query_time = plan.query_time_reference().unwrap(); let all = Arc::new(all); - let (prometheus_url, reads) = prometheus(all.clone()).await; - let process = start(&plan, &prometheus_url).await; + // Without the local raw buffer, every raw read is a remote read. + let prometheus = prometheus(all.clone()).await; + let reads = prometheus.remote_reads.clone(); + let process = start(&plan, &prometheus.url, &["--raw-buffer-ms", "0"]).await; remote_write(&process.url, &all).await; for query in plan.queries.iter().map(|q| &q.query) { let range = (3 * MINUTE, 6 * MINUTE, "30s"); @@ -633,8 +705,10 @@ async fn example4b_range() { let plan = plan_workload(&workload, &inputs).unwrap(); let at_query_time = plan.query_time_reference().unwrap(); let all = Arc::new(all); - let (prometheus_url, reads) = prometheus(all.clone()).await; - let process = start(&plan, &prometheus_url).await; + // Without the local raw buffer, every raw read is a remote read. + let prometheus = prometheus(all.clone()).await; + let reads = prometheus.remote_reads.clone(); + let process = start(&plan, &prometheus.url, &["--raw-buffer-ms", "0"]).await; remote_write(&process.url, &all).await; let query = &plan.queries[0].query; let range = (90 * MINUTE, 100 * MINUTE, "10m"); @@ -672,8 +746,10 @@ async fn example4_b3_range() { let plan = plan_workload(&workload, &inputs).unwrap(); let at_query_time = plan.query_time_reference().unwrap(); let all = Arc::new(all); - let (prometheus_url, reads) = prometheus(all.clone()).await; - let process = start(&plan, &prometheus_url).await; + // Without the local raw buffer, every raw read is a remote read. + let prometheus = prometheus(all.clone()).await; + let reads = prometheus.remote_reads.clone(); + let process = start(&plan, &prometheus.url, &["--raw-buffer-ms", "0"]).await; remote_write(&process.url, &all).await; let query = &plan.queries[0].query; let range = (10 * MINUTE, 14 * MINUTE, "60"); @@ -700,3 +776,55 @@ async fn example4_b3_range() { .await; assert_eq!(read, 1, "the evicted panes"); } + +/// Example 1's raw windows read from the local raw buffer and, before its +/// first sample, from Prometheus remote read; never from the JSON query +/// API. The data plane receives samples from 4:45 on while Prometheus holds +/// them from 0: at 5 min the 1-minute window takes 4:15..4:45 from remote +/// read and 5:00 from the buffer; at 10 min, the whole window is the +/// buffer's. Every answer equals the Planner executor's over all samples. +#[tokio::test] +async fn example1_raw_windows_from_the_buffer_and_remote_read() { + let series = [ + series("http_requests_total", &[("job", "api"), ("instance", "1")]), + series("http_requests_total", &[("job", "api"), ("instance", "2")]), + series("http_requests_total", &[("job", "web"), ("instance", "1")]), + ]; + let all = samples(&series, 0, 10 * MINUTE + 1, 15_000, |i, t| { + (i as f64 + 1.0) * (t as f64 / 1000.0).powi(2) + }); + let plan = plan_workload(&workload("example1"), &DeploymentInputs::default()).unwrap(); + assert!(!plan.query_time_raw_windows.is_empty()); + let at_query_time = plan.query_time_reference().unwrap(); + let all = Arc::new(all); + let prometheus = prometheus(all.clone()).await; + let process = start(&plan, &prometheus.url, &[]).await; + let written = |from: i64, to: i64| -> Vec { + all.iter() + .filter(|s| from <= s.timestamp_ms && s.timestamp_ms <= to) + .cloned() + .collect() + }; + remote_write(&process.url, &written(4 * MINUTE + 45_000, 5 * MINUTE)).await; + for query in plan.queries.iter().map(|q| &q.query) { + let before = prometheus.remote_reads.load(Ordering::SeqCst); + let expected = reference(&at_query_time, query, 5 * MINUTE, &all); + let actual = planned_answer(&process.url, query, 5 * MINUTE).await; + assert_same(query, &actual, &expected, 0.0); + let reads = prometheus.remote_reads.load(Ordering::SeqCst) - before; + assert!( + reads >= 1, + "{query}: the part before the buffer's first sample" + ); + } + remote_write(&process.url, &written(5 * MINUTE + 1, 10 * MINUTE)).await; + for query in plan.queries.iter().map(|q| &q.query) { + let before = prometheus.remote_reads.load(Ordering::SeqCst); + let expected = reference(&at_query_time, query, 10 * MINUTE, &all); + let actual = planned_answer(&process.url, query, 10 * MINUTE).await; + assert_same(query, &actual, &expected, 0.0); + let reads = prometheus.remote_reads.load(Ordering::SeqCst) - before; + assert_eq!(reads, 0, "{query}: the window is the buffer's"); + } + assert_eq!(prometheus.json_reads.load(Ordering::SeqCst), 0); +} diff --git a/data_plane/tests/raw_buffer.rs b/data_plane/tests/raw_buffer.rs new file mode 100644 index 00000000..6485c69c --- /dev/null +++ b/data_plane/tests/raw_buffer.rs @@ -0,0 +1,103 @@ +//! The local raw buffer holds the recent samples query-time raw windows +//! read, and says from when a read must come from Prometheus instead. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use data_plane::precompute::{shard_of, RawSample}; +use data_plane::raw_buffer::RawBuffer; + +const MINUTE: i64 = 60_000; + +/// Samples every 15 s in `[from, to)` of `metric` for each of `series`. +fn samples(metric: &str, series: usize, from: i64, to: i64) -> Vec { + let labels: Vec<_> = (0..series) + .map(|i| { + Arc::new(BTreeMap::from([ + ("__name__".to_string(), metric.to_string()), + ("instance".to_string(), i.to_string()), + ])) + }) + .collect(); + labels + .iter() + .flat_map(|labels| { + (from..to).step_by(15_000).map(move |t| RawSample { + labels: labels.clone(), + timestamp_ms: t, + value: t as f64, + }) + }) + .collect() +} + +/// `samples` appended as the ingestion workers would: each series to its +/// shard, in request order. +fn append(buffer: &RawBuffer, shards: usize, samples: &[RawSample]) { + let mut by_shard = vec![Vec::new(); shards]; + for sample in samples { + by_shard[shard_of(&sample.labels, shards)].push(sample.clone()); + } + for (shard, samples) in by_shard.iter().enumerate() { + buffer.append(shard, samples); + } +} + +fn sorted(mut samples: Vec) -> Vec<(BTreeMap, i64)> { + let mut points: Vec<_> = samples + .drain(..) + .map(|s| ((*s.labels).clone(), s.timestamp_ms)) + .collect(); + points.sort(); + points +} + +/// A read inside the horizon is the buffer's alone and equals the samples +/// of the metric in `(start, end]`; another metric's are not returned. +#[test] +fn reads_recent_samples_of_one_metric() { + let buffer = RawBuffer::new(10 * MINUTE, 3); + let all = samples("m", 5, 0, 20 * MINUTE + 1); + append(&buffer, 3, &all); + append(&buffer, 3, &samples("other", 2, 0, 20 * MINUTE + 1)); + let read = buffer.read("m", 15 * MINUTE, 20 * MINUTE); + assert_eq!(read.older_than_ms, 15 * MINUTE, "nothing from Prometheus"); + let expected: Vec<_> = all + .iter() + .filter(|s| 15 * MINUTE < s.timestamp_ms && s.timestamp_ms <= 20 * MINUTE) + .cloned() + .collect(); + assert_eq!(sorted(read.samples), sorted(expected)); +} + +/// The buffer keeps its horizon back from the newest sample: a read that +/// starts earlier gets the samples after the cut from the buffer, and must +/// read the rest, up to the cut, from Prometheus. +#[test] +fn reads_past_the_horizon_split_at_the_cut() { + let buffer = RawBuffer::new(5 * MINUTE, 2); + let all = samples("m", 4, 0, 20 * MINUTE + 1); + append(&buffer, 2, &all); + let read = buffer.read("m", 10 * MINUTE, 20 * MINUTE); + assert_eq!(read.older_than_ms, 15 * MINUTE); + assert!(read.samples.iter().all(|s| s.timestamp_ms > 15 * MINUTE)); + assert_eq!(read.samples.len(), 4 * 20); + let (series, held) = buffer.size(); + assert_eq!((series, held), (4, 4 * 20), "older samples are dropped"); +} + +/// Before its first sample the buffer holds nothing: a read of earlier +/// times comes from Prometheus up to that sample. +#[test] +fn samples_before_startup_come_from_prometheus() { + let buffer = RawBuffer::new(30 * MINUTE, 1); + assert_eq!(buffer.read("m", 0, MINUTE).older_than_ms, MINUTE); + append( + &buffer, + 1, + &samples("m", 1, 4 * MINUTE + 30_000, 5 * MINUTE + 1), + ); + let read = buffer.read("m", 4 * MINUTE, 5 * MINUTE); + assert_eq!(read.older_than_ms, 4 * MINUTE + 30_000); + assert_eq!(read.samples.len(), 2, "4:45 and 5:00"); +} From 66488dccc97cd2a9db0cf10246925c557d54cf2b Mon Sep 17 00:00:00 2001 From: zzylol Date: Mon, 19 Oct 2026 08:20:41 +0000 Subject: [PATCH 2/2] feat(data-plane): keep each metric in the raw buffer for its own longest query-time read The buffer kept every metric for the plan's single longest raw read: in the benefit plan, 60 minutes of the counter whose raw reads span one minute, and of metrics no raw window reads. It now keeps only the metrics the plan reads at query time, each for its own longest read (query-time raw windows, and panes built at query time and kept). --- README.md | 6 +- data_plane/src/engine.rs | 25 +++-- data_plane/src/main.rs | 7 +- data_plane/src/raw_buffer.rs | 166 +++++++++++++++++++++++---------- data_plane/tests/raw_buffer.rs | 42 +++++++++ 5 files changed, 183 insertions(+), 63 deletions(-) diff --git a/README.md b/README.md index ab3d92f2..9ed8b0ed 100644 --- a/README.md +++ b/README.md @@ -56,9 +56,9 @@ An instant query whose text is one of the plan's queries is answered by the plan (`x-asap-execution: plan`). Its stored inputs read the panes ending at `T - offset`, so `T` must be on the pane grid and those panes stored; its query-time raw windows are read from the local raw buffer: the recent -samples Remote Write delivered, sharded like ingestion and kept for the -plan's longest query-time raw read back from the newest sample -(`--raw-buffer-ms` overrides it; 0 keeps none). The part of a window older +samples Remote Write delivered, sharded like ingestion. Each metric the +plan reads at query time is kept for its longest query-time read back from +the newest sample (`--raw-buffer-ms` overrides it; 0 keeps none). The part of a window older than the buffer is read with Prometheus remote read (`/api/v1/read`). An input the plan keeps across evaluations (`MatAtQueryKeptAcrossEvaluations`) is computed from its pane's raw samples when an evaluation reads it and no evaluation has kept diff --git a/data_plane/src/engine.rs b/data_plane/src/engine.rs index 8fc2e487..88cf86d4 100644 --- a/data_plane/src/engine.rs +++ b/data_plane/src/engine.rs @@ -77,9 +77,10 @@ pub struct EngineOptions { pub ingestion_workers: usize, /// Query threads; each runs one query at a time. pub query_threads: usize, - /// How far back from the newest sample the local raw buffer keeps - /// samples, in ms; `None` for the plan's longest query-time raw read - /// ([`RawBuffer::horizon_ms`]). 0 keeps none: every raw read is a + /// How far back from the newest sample the local raw buffer keeps the + /// samples of each metric the plan reads at query time, in ms; `None` + /// for each metric's longest query-time raw read + /// ([`RawBuffer::horizons`]). 0 keeps none: every raw read is a /// Prometheus remote read. pub raw_buffer_ms: Option, } @@ -145,13 +146,19 @@ impl EngineHandle { ?threads, "stored summaries" ); - let raw_buffer_ms = options - .raw_buffer_ms - .unwrap_or_else(|| RawBuffer::horizon_ms(&plan)); - tracing::info!(raw_buffer_ms, "local raw buffer"); + let mut horizons = RawBuffer::horizons(&plan); + if let Some(override_ms) = options.raw_buffer_ms { + horizons.retain(|_, _| override_ms > 0); + horizons.values_mut().for_each(|ms| *ms = override_ms); + } + tracing::info!(?horizons, "local raw buffer"); drop(plan); - let raw_buffer = (raw_buffer_ms > 0) - .then(|| Arc::new(RawBuffer::new(raw_buffer_ms, threads.ingestion_workers))); + let raw_buffer = (!horizons.is_empty()).then(|| { + Arc::new(RawBuffer::with_horizons( + horizons, + threads.ingestion_workers, + )) + }); 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()) diff --git a/data_plane/src/main.rs b/data_plane/src/main.rs index 24ed4838..d1f33755 100644 --- a/data_plane/src/main.rs +++ b/data_plane/src/main.rs @@ -34,9 +34,10 @@ struct Args { /// Query threads, each running one query at a time; default one per core. #[arg(long)] query_threads: Option, - /// How far back from the newest sample to keep raw samples for - /// query-time raw windows, in ms; default the plan's longest such read. - /// 0 reads every raw window from Prometheus remote read. + /// How far back from the newest sample to keep the raw samples of each + /// metric read at query time, in ms; default each metric's longest + /// query-time read in the plan. 0 reads every raw window from + /// Prometheus remote read. #[arg(long)] raw_buffer_ms: Option, } diff --git a/data_plane/src/raw_buffer.rs b/data_plane/src/raw_buffer.rs index 09035bbe..b831d74f 100644 --- a/data_plane/src/raw_buffer.rs +++ b/data_plane/src/raw_buffer.rs @@ -2,15 +2,17 @@ //! query-time raw windows instead of asking Prometheus. //! //! It is fed by Remote Write through the ingestion workers and sharded like -//! them: each worker appends its shard's samples. It keeps the samples of -//! the plan's longest query-time read back from the newest sample (its -//! horizon); a read further back takes the older part from Prometheus -//! remote read. +//! them: each worker appends its shard's samples. It keeps each metric the +//! plan reads at query time for that metric's longest query-time read back +//! from the newest sample (its horizon), and no other metric; a read +//! further back takes the older part from Prometheus remote read. use std::collections::{BTreeMap, HashMap, VecDeque}; use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::{Arc, RwLock}; +use asap_executor::planner::ir::operator::Source; +use asap_executor::planner::ir::{NonASAPOp, Operator}; use asap_planner::{InstallablePlan, StoredOutputBuild}; use crate::precompute::RawSample; @@ -24,13 +26,30 @@ const UNSET: i64 = i64::MIN; struct Shard { /// Each series' samples in time order, by metric and label set. series: HashMap>>, - /// Samples at or before this time may have been dropped. - trimmed_through_ms: i64, + /// By metric: samples at or before this time may have been dropped. + trimmed_through_ms: HashMap, samples: usize, } +/// How long each metric is kept. +enum Horizons { + /// Every metric, for the same time. + All(i64), + /// Only these metrics, each for its own time. + PerMetric(BTreeMap), +} + +impl Horizons { + fn of(&self, metric: &str) -> Option { + match self { + Self::All(ms) => Some(*ms), + Self::PerMetric(by_metric) => by_metric.get(metric).copied(), + } + } +} + pub struct RawBuffer { - horizon_ms: i64, + horizons: Horizons, shards: Vec>, /// The first sample's time: the buffer holds every sample after it /// that Remote Write delivered, until it is trimmed. @@ -49,41 +68,73 @@ pub struct BufferRead { } impl RawBuffer { - /// The longest a query of `plan` reads raw samples back from its - /// evaluation time: over its query-time raw windows and the panes it - /// builds at query time and keeps, `end_offset + width`, and for the - /// latter their retention, which covers their range evaluations. - pub fn horizon_ms(plan: &InstallablePlan) -> i64 { + /// The longest a query of `plan` reads each metric's raw samples back + /// from its evaluation time: over its query-time raw windows and the + /// panes it builds at query time and keeps, `end_offset + width`, and + /// for the latter their retention, which covers their range evaluations. + pub fn horizons(plan: &InstallablePlan) -> BTreeMap { + let dag = &plan.dag.dag; + let metric = |scan| { + dag.nodes + .iter() + .find(|n| n.id == scan) + .and_then(|n| match &n.payload { + Operator::NonASAP(NonASAPOp::Scan { + source: Source::TimeSeries { metric }, + .. + }) => Some(metric.clone()), + _ => None, + }) + }; let windows = plan .query_time_raw_windows .iter() - .map(|w| w.end_offset_ms + w.width_ms); + .map(|w| (w.scan_node_id, w.end_offset_ms + w.width_ms)); let kept = plan .stored_outputs .iter() .filter(|o| o.built_at == StoredOutputBuild::QueryRun) - .map(|o| (o.pane.reach_ms().max(o.retention_ms)) as i64); - windows.chain(kept).max().unwrap_or(0) + .map(|o| { + ( + o.pane.scan_node_id, + o.pane.reach_ms().max(o.retention_ms) as i64, + ) + }); + let mut horizons = BTreeMap::new(); + for (scan, reach_ms) in windows.chain(kept) { + if let Some(metric) = metric(scan) { + let horizon = horizons.entry(metric).or_insert(0); + *horizon = reach_ms.max(*horizon); + } + } + horizons } + /// A buffer that keeps every metric for `horizon_ms`. pub fn new(horizon_ms: i64, shards: usize) -> Self { + Self::with(Horizons::All(horizon_ms), shards) + } + + /// A buffer that keeps only the metrics of `horizons`, each for its own + /// time ([`Self::horizons`]). + pub fn with_horizons(horizons: BTreeMap, shards: usize) -> Self { + Self::with(Horizons::PerMetric(horizons), shards) + } + + fn with(horizons: Horizons, shards: usize) -> Self { Self { - horizon_ms, + horizons, shards: (0..shards.max(1)) - .map(|_| { - RwLock::new(Shard { - trimmed_through_ms: UNSET, - ..Shard::default() - }) - }) + .map(|_| RwLock::new(Shard::default())) .collect(), start_ms: AtomicI64::new(UNSET), newest_ms: AtomicI64::new(UNSET), } } - /// Appends `samples`, all of shard `shard`, then drops its samples older - /// than the horizon before the newest sample. + /// Appends the samples of `samples`, all of shard `shard`, whose metric + /// the buffer keeps, then drops the shard's samples older than each + /// metric's horizon before the newest sample. pub fn append(&self, shard: usize, samples: &[RawSample]) { let (Some(first), Some(newest)) = ( samples.iter().map(|s| s.timestamp_ms).min(), @@ -98,60 +149,79 @@ impl RawBuffer { .newest_ms .fetch_max(newest, Ordering::AcqRel) .max(newest); - let trim_through = newest - self.horizon_ms; let mut shard = self.shards[shard].write().expect("raw buffer lock"); - let shard = &mut *shard; - let mut added = 0usize; + let Shard { + series, + trimmed_through_ms, + samples: held, + } = &mut *shard; // Samples of one series arrive together: look each series up once. for run in samples.chunk_by(|a, b| Arc::ptr_eq(&a.labels, &b.labels)) { let labels = &run[0].labels; let metric = labels.get("__name__").cloned().unwrap_or_default(); - let series = shard - .series + let Some(horizon_ms) = self.horizons.of(&metric) else { + continue; + }; + let trim_through = newest - horizon_ms; + if run.iter().all(|s| s.timestamp_ms <= trim_through) { + continue; + } + let points = series .entry(metric) .or_default() .entry(labels.clone()) .or_default(); for sample in run.iter().filter(|s| s.timestamp_ms > trim_through) { let point = (sample.timestamp_ms, sample.value); - match series.back() { + match points.back() { Some(&(time, _)) if time >= point.0 => { - let at = series.partition_point(|&(t, _)| t < point.0); - if series.get(at).is_some_and(|&(t, _)| t == point.0) { + let at = points.partition_point(|&(t, _)| t < point.0); + if points.get(at).is_some_and(|&(t, _)| t == point.0) { continue; } - series.insert(at, point); + points.insert(at, point); } - _ => series.push_back(point), + _ => points.push_back(point), } - added += 1; + *held += 1; } } - shard.samples += added; - if trim_through > shard.trimmed_through_ms { - shard.trimmed_through_ms = trim_through; - let mut dropped = 0; - for series in shard.series.values_mut() { - series.retain(|_, points| { - let keep = points.partition_point(|&(t, _)| t <= trim_through); - dropped += keep; - points.drain(..keep); - !points.is_empty() - }); + for (metric, by_labels) in series.iter_mut() { + let Some(horizon_ms) = self.horizons.of(metric) else { + continue; + }; + let trim_through = newest - horizon_ms; + let trimmed = trimmed_through_ms.entry(metric.clone()).or_insert(UNSET); + if trim_through <= *trimmed { + continue; } - shard.samples -= dropped; + *trimmed = trim_through; + by_labels.retain(|_, points| { + let old = points.partition_point(|&(t, _)| t <= trim_through); + *held -= old; + points.drain(..old); + !points.is_empty() + }); } } /// `metric`'s samples in `(start_ms, end_ms]` that the buffer holds, /// and from when it holds all of them. pub fn read(&self, metric: &str, start_ms: i64, end_ms: i64) -> BufferRead { + if self.horizons.of(metric).is_none() { + return BufferRead { + samples: Vec::new(), + older_than_ms: end_ms, + }; + } let first = self.start_ms.load(Ordering::Acquire); let mut older_than_ms = start_ms.max(if first == UNSET { end_ms } else { first }); let mut parts = Vec::with_capacity(self.shards.len()); for shard in &self.shards { let shard = shard.read().expect("raw buffer lock"); - older_than_ms = older_than_ms.max(shard.trimmed_through_ms); + if let Some(&trimmed) = shard.trimmed_through_ms.get(metric) { + older_than_ms = older_than_ms.max(trimmed); + } let mut samples = Vec::new(); for (labels, points) in shard.series.get(metric).into_iter().flatten() { let from = points.partition_point(|&(t, _)| t <= start_ms); diff --git a/data_plane/tests/raw_buffer.rs b/data_plane/tests/raw_buffer.rs index 6485c69c..e410339a 100644 --- a/data_plane/tests/raw_buffer.rs +++ b/data_plane/tests/raw_buffer.rs @@ -101,3 +101,45 @@ fn samples_before_startup_come_from_prometheus() { assert_eq!(read.older_than_ms, 4 * MINUTE + 30_000); assert_eq!(read.samples.len(), 2, "4:45 and 5:00"); } + +/// Each metric is kept for its own longest read, and a metric no raw +/// window reads is not kept at all. +#[test] +fn keeps_each_metric_for_its_own_horizon() { + let horizons = BTreeMap::from([ + ("short".to_string(), MINUTE), + ("long".to_string(), 10 * MINUTE), + ]); + let buffer = RawBuffer::with_horizons(horizons, 2); + for metric in ["short", "long", "unread"] { + append(&buffer, 2, &samples(metric, 3, 0, 20 * MINUTE + 1)); + } + assert_eq!( + buffer.read("short", 0, 20 * MINUTE).older_than_ms, + 19 * MINUTE + ); + assert_eq!( + buffer.read("long", 0, 20 * MINUTE).older_than_ms, + 10 * MINUTE + ); + assert_eq!( + buffer.size(), + (6, 3 * 4 + 3 * 40), + "unread samples are not kept" + ); +} + +/// The plan's horizons: Example 1 reads its counter at query time over +/// 1-minute windows, nothing else. +#[test] +fn horizons_come_from_the_plans_query_time_reads() { + let path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) + .join("../examples/planner-layering-example1.json"); + let workload = serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap(); + let plan = control_plane::plan_workload(&workload, &control_plane::DeploymentInputs::default()) + .unwrap(); + assert_eq!( + RawBuffer::horizons(&plan), + BTreeMap::from([("http_requests_total".to_string(), MINUTE)]) + ); +}