feat(data-plane): read query-time raw windows from a local raw buffer and Prometheus remote read

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. 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
 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

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,29 @@ 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 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<i64>,
 }
 
-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 +103,7 @@ pub struct EngineHandle {
     workers: Vec<mpsc::Sender<WorkerCommand>>,
     queries: SyncSender<QueryJob>,
     precompute: Arc<Shared>,
+    raw_buffer: Option<Arc<RawBuffer>>,
 }
 
 /// Reads and validates the plan serialized in `plan_json`, and its stored
@@ -125,8 +134,9 @@ impl EngineHandle {
     pub fn spawn(
         plan_json: String,
         allowed_lateness_ms: i64,
-        threads: Threads,
+        options: EngineOptions,
     ) -> anyhow::Result<Self> {
+        let threads = options;
         let (plan, sds) = install(&plan_json)?;
         tracing::info!(
             definitions = sds.definitions.len(),
@@ -136,7 +146,19 @@ impl EngineHandle {
             ?threads,
             "stored summaries"
         );
+        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 = (!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())
@@ -146,10 +168,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 +238,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()

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<Self> {
         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<BTreeMap<u64, Vec<RawSample>>, 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)

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,12 @@ struct Args {
     /// Query threads, each running one query at a time; default one per core.
     #[arg(long)]
     query_threads: Option<usize>,
+    /// 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<i64>,
 }
 
 #[tokio::main]
@@ -44,18 +50,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?;

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<Shared>,
     /// The global watermark this worker has built its parts through.
     closed_through_ms: Option<i64>,
+    /// 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<Arc<RawBuffer>>,
 }
 
 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<Arc<RawBuffer>>) -> 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