From 42ccf6d583d17692d5c2950785df3ff146a0072c Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Tue, 21 Apr 2026 08:29:13 -0400 Subject: [PATCH] =?UTF-8?q?feat(sketchdb):=20cold-query=20fallback=20with?= =?UTF-8?q?=20local-FS=20raw=20store=20(=C2=A75.2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses TODO.md blocker #1. Unblocks the paper's "sketches for hot, exact for cold" story: capability-miss queries now serve exact answers from raw observability samples in the cold tier instead of hitting Prometheus (or failing) unconditionally. ## What's new * `drivers/query/fallback/cold_store/` — storage-agnostic `ColdStore` trait + `LocalFsColdStore` impl. On-disk layout (`raw//YYYY/MM/DD/HH/part-NNNNNN.jsonl`) is identical to what a future S3 cold store will use, so the adapter stays source-compatible when we swap backends. * `drivers/query/fallback/s3_adapter.rs` — `ColdFallback` implements `FallbackClient`. Parses PromQL, extracts `(metric, predicates, op)`, scans the cold store, computes the answer. Supported shapes for v1: - bare instant vector selector (`metric{labels}`) - no-grouping scalar aggregation (`sum|count|avg|min|max(metric{...})`) Label matchers: `=`, `!=` — regex delegated upstream. Anything outside this surface falls through to the optional inner `FallbackClient` (chain-of-responsibility, typically the existing Prometheus proxy). * `drivers/query/fallback/metrics.rs` — hot/cold telemetry counters (`queryengine_hot_queries_total`, `queryengine_cold_queries_total`, `queryengine_cold_bytes_served_total`), keyed by `(metric, shape)`. Mirrors the PR #51 schema-barrier-counter pattern. * `AdapterConfig::prometheus_promql_with_cold` convenience constructor that composes the cold adapter in front of a Prometheus proxy. ## Routing No engine changes needed. The existing `process_query_request` already falls through to `FallbackClient` on engine-miss — configuring the fallback as a `ColdFallback` naturally routes Purged-segment / capability-miss queries through the cold tier with the Prometheus proxy as the tail-of-chain for unsupported shapes. ## Tests 752 → 777 tests (+25): * 5 unit tests on the JSONL format + hour-prefix helpers * 4 on `LocalFsColdStore::scan` (range filter, missing prefix, hour-boundary span) * 10 on `ColdFallback` plan extraction + aggregation + float formatting * 6 end-to-end HTTP integration tests in `tests/cold_fallback_tests.rs` covering: bare selector, sum aggregation, label filtering, unsupported-shape delegation, telemetry-counter increments, and the "Purged range served from raw" paper story. --- TODO.md | 39 +- .../src/drivers/query/adapters/config.rs | 42 ++ .../query/fallback/cold_store/format.rs | 150 +++++ .../query/fallback/cold_store/local_fs.rs | 188 ++++++ .../drivers/query/fallback/cold_store/mod.rs | 78 +++ .../src/drivers/query/fallback/metrics.rs | 57 ++ .../src/drivers/query/fallback/mod.rs | 6 + .../src/drivers/query/fallback/s3_adapter.rs | 572 ++++++++++++++++++ .../src/tests/cold_fallback_tests.rs | 402 ++++++++++++ asap-query-engine/src/tests/mod.rs | 1 + 10 files changed, 1516 insertions(+), 19 deletions(-) create mode 100644 asap-query-engine/src/drivers/query/fallback/cold_store/format.rs create mode 100644 asap-query-engine/src/drivers/query/fallback/cold_store/local_fs.rs create mode 100644 asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs create mode 100644 asap-query-engine/src/drivers/query/fallback/metrics.rs create mode 100644 asap-query-engine/src/drivers/query/fallback/s3_adapter.rs create mode 100644 asap-query-engine/src/tests/cold_fallback_tests.rs diff --git a/TODO.md b/TODO.md index bcb1e082..fb86f763 100644 --- a/TODO.md +++ b/TODO.md @@ -8,25 +8,26 @@ See the design source at [`docs/design-sketch-db.md`](docs/design-sketch-db.md). ## For paper submission (blocker) -### 1. Cold-query fallback — §5.2 of the sketch-DB design - -Capability-miss at query time today falls through to the §5.2 -forwarding adapter which hits Prometheus (the raw source). For -the paper's "hot sketch + cold exact" story we need: - -- **Adapter that reads from S3-resident raw exports on cold - miss.** New module `drivers/query/fallback/s3_adapter.rs` in - parallel with the existing Prometheus fallback. Input: a - `(metric, time_range, labels)` triple; output: a computed - answer using exact raw records. -- **Cost model aware of hot/cold split.** When the schema - timeline says a time range is `Purged`, the query routes - through the S3 adapter instead of failing. -- **Telemetry.** Counters for bytes-served-from-sketch vs - bytes-served-from-S3 per query, keyed by query shape. - Mirror the PR #47 pattern. - -Scale target: ≤2× P99 latency degradation vs. warm-hot queries. +### 1. Cold-query fallback — §5.2 of the sketch-DB design — **done (local-FS cold store)** + +Initial v1 landed: [`drivers/query/fallback/s3_adapter.rs`](asap-query-engine/src/drivers/query/fallback/s3_adapter.rs) +is a `FallbackClient` that serves capability-misses from a +hour-bucketed JSONL raw store. The format (`raw//YYYY/MM/DD/HH/part-NNNNNN.jsonl`) +is byte-identical to the S3 layout, so a future +`S3ColdStore: ColdStore` drops in with no adapter changes. + +Follow-ups (not paper-blocking): + +- **S3-backed `ColdStore` impl** next to the local-FS one; same + trait, `aws-sdk-s3` list-objects-v2 for prefix pruning. +- **Richer query surface.** Today we compute `metric{...}`, + `sum|count|avg|min|max(...)`. Regex matchers, `by (...)` + grouping, and `rate/increase` over raw samples delegate to + the chained inner fallback (typically Prometheus). Adding + grouping + regex is ~200 LOC when needed. +- **Latency target.** Paper claim is ≤2× P99 vs. warm-hot — + unverified until the multi-agent harness lands (blocker #6 + of `DataCollector/TODO.md`). ### 2. Accuracy-profile library per sketch type diff --git a/asap-query-engine/src/drivers/query/adapters/config.rs b/asap-query-engine/src/drivers/query/adapters/config.rs index 948237ae..f38e2b21 100644 --- a/asap-query-engine/src/drivers/query/adapters/config.rs +++ b/asap-query-engine/src/drivers/query/adapters/config.rs @@ -60,6 +60,48 @@ impl AdapterConfig { ) } + /// Prometheus + cold-tier fallback chain (§5.2 of the sketch-DB design). + /// + /// Composes a [`ColdFallback`](crate::drivers::query::fallback::ColdFallback) + /// in front of a [`PrometheusHttpFallback`](crate::drivers::query::fallback::PrometheusHttpFallback) + /// so capability-misses first try the raw cold tier (exact + /// answers for supported query shapes) and only hit the live + /// Prometheus if the cold adapter can't handle the shape. + /// + /// * `cold_root` — local-FS root that mirrors the S3 key + /// layout documented in + /// [`cold_store::format`](crate::drivers::query::fallback::cold_store::format). + /// Swap in an S3-backed [`ColdStore`](crate::drivers::query::fallback::ColdStore) + /// impl later without touching this config. + /// * `prom_fallback_url` — upstream Prometheus used for the + /// tail of the fallback chain; set to `None` to short-circuit + /// unsupported shapes with an empty vector instead of + /// forwarding. + pub fn prometheus_promql_with_cold( + cold_root: std::path::PathBuf, + prom_fallback_url: Option, + ) -> Self { + use crate::drivers::query::fallback::{ + ColdFallback, LocalFsColdStore, PrometheusHttpFallback, + }; + + let cold_store = Arc::new(LocalFsColdStore::new(cold_root)); + let cold = ColdFallback::new(cold_store); + let cold: Arc = match prom_fallback_url { + Some(url) => { + let prom: Arc = Arc::new(PrometheusHttpFallback::new(url)); + Arc::new(cold.with_inner(prom)) + } + None => Arc::new(cold), + }; + + Self::new( + QueryProtocol::PrometheusHttp, + QueryLanguage::promql, + Some(cold), + ) + } + /// Create a configuration for ClickHouse HTTP with SQL /// Convenience constructor for ClickHouse adapter pub fn clickhouse_sql(base_url: String, database: String, forward_unsupported: bool) -> Self { diff --git a/asap-query-engine/src/drivers/query/fallback/cold_store/format.rs b/asap-query-engine/src/drivers/query/fallback/cold_store/format.rs new file mode 100644 index 00000000..0359ab5c --- /dev/null +++ b/asap-query-engine/src/drivers/query/fallback/cold_store/format.rs @@ -0,0 +1,150 @@ +//! JSONL raw-sample format + key-prefix helpers. +//! +//! The format is intentionally boring: one JSON object per line, +//! one sample per object. That makes it cheap for an OTel exporter +//! or a test fixture to produce and for any reader (Python, jq, +//! ClickHouse external table, etc.) to consume. +//! +//! Key layout is hour-bucketed so a range scan that spans `N` +//! hours touches at most `N` key-prefixes regardless of ingest +//! rate. The same layout works on S3 (list-objects-v2 with +//! `Prefix`) without modification. + +use chrono::{DateTime, Datelike, Timelike, Utc}; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeMap; + +use super::ColdStoreError; + +/// A single raw observability sample as written by the cold +/// exporter. `labels` is a `BTreeMap` so the on-disk JSON is +/// deterministic per sample (useful for golden tests). +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct RawSample { + pub ts_ms: i64, + pub labels: BTreeMap, + pub value: f64, +} + +/// Key-prefix for the hour-bucket containing `ts_ms`, relative to +/// the cold-store root. Identical shape for local-FS and S3. +/// +/// Example: `raw/http_requests_total/2026/04/21/08/` +pub fn part_path_prefix(metric: &str, ts_ms: i64) -> String { + let dt: DateTime = DateTime::::from_timestamp_millis(ts_ms) + .unwrap_or_else(|| DateTime::::from_timestamp(0, 0).unwrap()); + format!( + "raw/{}/{:04}/{:02}/{:02}/{:02}/", + metric, + dt.year(), + dt.month(), + dt.day(), + dt.hour(), + ) +} + +/// Enumerate the hour-bucket prefixes covering the half-open +/// range `[start_ms, end_ms)`. Always returns at least one bucket +/// (the one containing `start_ms`). Used by `ColdStore` impls to +/// drive object-listing / directory-walk. +pub fn hour_prefixes(metric: &str, start_ms: i64, end_ms: i64) -> Vec { + if end_ms <= start_ms { + return vec![part_path_prefix(metric, start_ms)]; + } + const HOUR_MS: i64 = 3_600_000; + let first_hour = (start_ms / HOUR_MS) * HOUR_MS; + // Align `end` up to the next hour boundary; we scan strictly + // *less than* `end_ms` so the last included bucket is the one + // containing `end_ms - 1`. + let last_hour = ((end_ms - 1) / HOUR_MS) * HOUR_MS; + let mut out = Vec::new(); + let mut cur = first_hour; + while cur <= last_hour { + out.push(part_path_prefix(metric, cur)); + cur += HOUR_MS; + } + out +} + +/// Parse a `.jsonl` blob into `RawSample`s, filtering to the +/// half-open range `[start_ms, end_ms)`. Malformed lines fail the +/// whole parse — partial results from a corrupted part are worse +/// than an error a caller can route around. +pub fn parse_jsonl( + bytes: &[u8], + start_ms: i64, + end_ms: i64, +) -> Result, ColdStoreError> { + let text = std::str::from_utf8(bytes) + .map_err(|e| ColdStoreError::Malformed(format!("non-utf8: {e}")))?; + let mut out = Vec::new(); + for (lineno, line) in text.lines().enumerate() { + let line = line.trim(); + if line.is_empty() { + continue; + } + let sample: RawSample = serde_json::from_str(line) + .map_err(|e| ColdStoreError::Malformed(format!("line {}: {}", lineno + 1, e)))?; + if sample.ts_ms >= start_ms && sample.ts_ms < end_ms { + out.push(sample); + } + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + + use chrono::TimeZone; + + fn ts_ms(year: i32, month: u32, day: u32, hour: u32, minute: u32) -> i64 { + Utc.with_ymd_and_hms(year, month, day, hour, minute, 0) + .unwrap() + .timestamp_millis() + } + + #[test] + fn prefix_shape() { + let ts = ts_ms(2026, 4, 21, 8, 15); + assert_eq!( + part_path_prefix("http_requests_total", ts), + "raw/http_requests_total/2026/04/21/08/" + ); + } + + #[test] + fn hour_prefixes_single_bucket() { + let start = ts_ms(2026, 4, 21, 8, 15); + let end = start + 60_000; + let p = hour_prefixes("m", start, end); + assert_eq!(p.len(), 1); + } + + #[test] + fn hour_prefixes_span_two_hours() { + let start = ts_ms(2026, 4, 21, 8, 15); + let end = start + 3_600_000 + 1; + let p = hour_prefixes("m", start, end); + assert_eq!(p.len(), 2); + assert_ne!(p[0], p[1]); + } + + #[test] + fn parse_jsonl_filters_range() { + let blob = r#"{"ts_ms":100,"labels":{"a":"1"},"value":1.0} +{"ts_ms":200,"labels":{"a":"2"},"value":2.0} +{"ts_ms":300,"labels":{"a":"3"},"value":3.0} +"#; + let out = parse_jsonl(blob.as_bytes(), 150, 300).unwrap(); + // end is exclusive -> only ts=200 matches + assert_eq!(out.len(), 1); + assert_eq!(out[0].ts_ms, 200); + } + + #[test] + fn parse_jsonl_malformed_errs() { + let blob = "not-json\n"; + assert!(parse_jsonl(blob.as_bytes(), 0, i64::MAX).is_err()); + } +} diff --git a/asap-query-engine/src/drivers/query/fallback/cold_store/local_fs.rs b/asap-query-engine/src/drivers/query/fallback/cold_store/local_fs.rs new file mode 100644 index 00000000..781dc75c --- /dev/null +++ b/asap-query-engine/src/drivers/query/fallback/cold_store/local_fs.rs @@ -0,0 +1,188 @@ +//! Local-filesystem [`ColdStore`] impl. +//! +//! Walks the same directory layout an S3 bucket would hold, so a +//! future `S3ColdStore` can drop in without the adapter caring. +//! Used today by tests and by the single-node evaluation +//! deployment. +//! +//! Concurrency: scans read each file via `tokio::fs::read`, so +//! multiple overlapping scans can progress in parallel without +//! serialization. + +use async_trait::async_trait; +use std::path::{Path, PathBuf}; + +use super::format::{hour_prefixes, parse_jsonl}; +use super::{ColdStore, ColdStoreError, RawSample}; + +/// Cold store backed by a local directory tree. +pub struct LocalFsColdStore { + root: PathBuf, +} + +impl LocalFsColdStore { + /// Create a store rooted at `root`. The directory must exist; + /// producing raw dumps is the exporter's job, not the reader's. + pub fn new(root: impl Into) -> Self { + Self { root: root.into() } + } + + pub fn root(&self) -> &Path { + &self.root + } + + /// List the JSONL parts inside `prefix_dir`, sorted by file name + /// so scans over the same input are deterministic. + async fn list_parts(&self, prefix_dir: &Path) -> Result, ColdStoreError> { + let mut out = Vec::new(); + let mut rd = match tokio::fs::read_dir(prefix_dir).await { + Ok(rd) => rd, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(out), + Err(e) => return Err(e.into()), + }; + while let Some(entry) = rd.next_entry().await? { + let path = entry.path(); + if path + .extension() + .and_then(|s| s.to_str()) + .is_some_and(|ext| ext == "jsonl") + { + out.push(path); + } + } + out.sort(); + Ok(out) + } +} + +#[async_trait] +impl ColdStore for LocalFsColdStore { + async fn scan( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError> { + let mut out = Vec::new(); + for prefix in hour_prefixes(metric, start_ms, end_ms) { + let dir = self.root.join(&prefix); + for part in self.list_parts(&dir).await? { + let bytes = tokio::fs::read(&part).await?; + let mut samples = parse_jsonl(&bytes, start_ms, end_ms)?; + out.append(&mut samples); + } + } + Ok(out) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use chrono::{TimeZone, Utc}; + use std::collections::BTreeMap; + use tempfile::TempDir; + + fn hour_ms(year: i32, month: u32, day: u32, hour: u32) -> i64 { + Utc.with_ymd_and_hms(year, month, day, hour, 0, 0) + .unwrap() + .timestamp_millis() + } + + async fn write_part(root: &Path, rel: &str, lines: &[RawSample]) { + let dir = root.join(rel); + tokio::fs::create_dir_all(&dir).await.unwrap(); + let mut buf = String::new(); + for s in lines { + buf.push_str(&serde_json::to_string(s).unwrap()); + buf.push('\n'); + } + tokio::fs::write(dir.join("part-000001.jsonl"), buf) + .await + .unwrap(); + } + + fn sample(ts_ms: i64, labels: &[(&str, &str)], value: f64) -> RawSample { + RawSample { + ts_ms, + labels: labels + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect::>(), + value, + } + } + + #[tokio::test] + async fn scan_returns_samples_in_range() { + let tmp = TempDir::new().unwrap(); + let base = hour_ms(2026, 4, 21, 8); + write_part( + tmp.path(), + "raw/http_requests_total/2026/04/21/08/", + &[ + sample(base + 1_000, &[("zone", "a")], 1.0), + sample(base + 2_000, &[("zone", "b")], 2.0), + ], + ) + .await; + + let s = LocalFsColdStore::new(tmp.path()); + let got = s + .scan("http_requests_total", base, base + 60_000) + .await + .unwrap(); + assert_eq!(got.len(), 2); + } + + #[tokio::test] + async fn scan_prunes_by_time() { + let tmp = TempDir::new().unwrap(); + let base = hour_ms(2026, 4, 21, 8); + write_part( + tmp.path(), + "raw/m/2026/04/21/08/", + &[ + sample(base + 1_000, &[], 1.0), + sample(base + 60_000, &[], 2.0), + ], + ) + .await; + + let s = LocalFsColdStore::new(tmp.path()); + let got = s.scan("m", base, base + 30_000).await.unwrap(); + assert_eq!(got.len(), 1); + assert_eq!(got[0].value, 1.0); + } + + #[tokio::test] + async fn scan_missing_prefix_is_empty_not_error() { + let tmp = TempDir::new().unwrap(); + let s = LocalFsColdStore::new(tmp.path()); + let got = s.scan("never_written", 0, 1).await.unwrap(); + assert!(got.is_empty()); + } + + #[tokio::test] + async fn scan_spans_hour_boundary() { + let tmp = TempDir::new().unwrap(); + let h8 = hour_ms(2026, 4, 21, 8); + let h9 = hour_ms(2026, 4, 21, 9); + write_part( + tmp.path(), + "raw/m/2026/04/21/08/", + &[sample(h8 + 3_599_000, &[], 1.0)], + ) + .await; + write_part( + tmp.path(), + "raw/m/2026/04/21/09/", + &[sample(h9 + 1_000, &[], 2.0)], + ) + .await; + + let s = LocalFsColdStore::new(tmp.path()); + let got = s.scan("m", h8 + 3_598_000, h9 + 2_000).await.unwrap(); + assert_eq!(got.len(), 2); + } +} diff --git a/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs b/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs new file mode 100644 index 00000000..c9a8a679 --- /dev/null +++ b/asap-query-engine/src/drivers/query/fallback/cold_store/mod.rs @@ -0,0 +1,78 @@ +//! Cold raw-sample store used by the §5.2 cold-query fallback. +//! +//! In the paper architecture, the edge OTel collector dumps raw +//! observability data to a cheap cold tier (S3) in parallel with +//! the sketch path. When a query hits a capability-miss — most +//! notably a `TimelineCoverage::Purged` segment whose sketch was +//! aged out — the engine falls through to this store to recover +//! an exact answer from raw records. +//! +//! This module exposes a **storage-agnostic** `ColdStore` trait so +//! the same `s3_adapter` fallback can point at either a local +//! filesystem root (used today + for tests) or a real S3 bucket +//! (future swap, identical object key layout — see +//! [`format::part_path_prefix`]). +//! +//! # Format +//! +//! Raw samples live under a deterministic key tree: +//! +//! ```text +//! /raw//YYYY/MM/DD/HH/part-NNNNNN.jsonl +//! ``` +//! +//! Each line is one sample encoded as JSON: +//! +//! ```json +//! {"ts_ms": 1713657600000, "labels": {"zone": "a"}, "value": 42.5} +//! ``` +//! +//! See [`format`] for serialization and path helpers. + +use async_trait::async_trait; +use std::collections::BTreeMap; +use thiserror::Error; + +pub mod format; +pub mod local_fs; + +pub use format::{part_path_prefix, RawSample}; +pub use local_fs::LocalFsColdStore; + +/// Error surface for cold-store scans. +#[derive(Debug, Error)] +pub enum ColdStoreError { + #[error("I/O error: {0}")] + Io(#[from] std::io::Error), + #[error("malformed raw record: {0}")] + Malformed(String), +} + +/// Read-only view over a cold raw-sample store. +/// +/// Scans are `(metric, [start_ms, end_ms))` — inclusive start, +/// exclusive end — matching the half-open range convention used by +/// the rest of the engine. Implementations are expected to: +/// +/// * prune by metric via the `/` key prefix, +/// * prune by hour via the `YYYY/MM/DD/HH/` key prefix, +/// * scan inside candidate parts and emit only samples whose +/// `ts_ms` falls in the requested range. +/// +/// Label matching is **not** pushed down here — callers filter +/// samples client-side. This keeps the trait small and makes the +/// local-FS / S3 impls trivially swappable. +#[async_trait] +pub trait ColdStore: Send + Sync { + /// Return all samples for `metric` whose timestamp lies in + /// `[start_ms, end_ms)`. Ordering is not guaranteed. + async fn scan( + &self, + metric: &str, + start_ms: i64, + end_ms: i64, + ) -> Result, ColdStoreError>; +} + +/// Convenience alias: a label set as stored in a [`RawSample`]. +pub type LabelSet = BTreeMap; diff --git a/asap-query-engine/src/drivers/query/fallback/metrics.rs b/asap-query-engine/src/drivers/query/fallback/metrics.rs new file mode 100644 index 00000000..c7fc1257 --- /dev/null +++ b/asap-query-engine/src/drivers/query/fallback/metrics.rs @@ -0,0 +1,57 @@ +//! Query-engine hot/cold telemetry counters. +//! +//! The paper's "sketches for hot, exact for cold" story needs a +//! hot-vs-cold breakdown at query time. We mirror the PR #51 +//! `queryengine_ingest_samples_blocked_by_schema_barrier_total` +//! pattern: `lazy_static!` registration, one counter per signal, +//! keyed by query shape so the dashboard can split by metric. +//! +//! * **Hot** = the `SimpleEngine` handled the query from live +//! sketch-backed state. +//! * **Cold** = the query fell through to a +//! [`ColdFallback`](super::s3_adapter::ColdFallback) and was +//! answered from the raw cold tier. +//! +//! The "shape" label is the parsed query's root op (`sum`, +//! `count`, `avg`, `selector`, ...) — low-cardinality by design, +//! so the `CounterVec` doesn't explode on the backend. A +//! companion `metric` label carries the first selector's +//! `__name__` so dashboards can slice by metric without the label +//! explosion that a raw-query label would cause. + +use lazy_static::lazy_static; +use prometheus::{register_counter_vec, CounterVec}; + +lazy_static! { + /// Queries served by the hot (sketch) path, keyed by + /// `(metric, shape)`. Incremented once per successful + /// `SimpleEngine::handle_query` that returned `Some(_)`. + pub static ref QUERIES_HOT_TOTAL: CounterVec = register_counter_vec!( + "queryengine_hot_queries_total", + "Queries answered from the sketch (hot) path, keyed by metric + query shape", + &["metric", "shape"] + ) + .unwrap(); + + /// Queries served by the cold (raw S3 / local-FS) path, keyed + /// by `(metric, shape)`. Incremented once per successful + /// `ColdFallback::execute_query`. + pub static ref QUERIES_COLD_TOTAL: CounterVec = register_counter_vec!( + "queryengine_cold_queries_total", + "Queries answered from the cold raw-sample path, keyed by metric + query shape", + &["metric", "shape"] + ) + .unwrap(); + + /// Raw sample bytes served out of the cold store in response + /// to queries. Bytes here = the JSON payload of the raw + /// records scanned (pre-filter), which is the cheapest stable + /// proxy for "how much cold data the query had to materialise". + /// Keyed by `(metric, shape)` to match the counter above. + pub static ref BYTES_SERVED_COLD_TOTAL: CounterVec = register_counter_vec!( + "queryengine_cold_bytes_served_total", + "Raw sample bytes scanned from the cold tier in answering queries", + &["metric", "shape"] + ) + .unwrap(); +} diff --git a/asap-query-engine/src/drivers/query/fallback/mod.rs b/asap-query-engine/src/drivers/query/fallback/mod.rs index 6b42f516..4950b31a 100644 --- a/asap-query-engine/src/drivers/query/fallback/mod.rs +++ b/asap-query-engine/src/drivers/query/fallback/mod.rs @@ -83,6 +83,12 @@ mod clickhouse; mod elastic; mod prometheus; +pub mod cold_store; +pub mod metrics; +pub mod s3_adapter; + pub use clickhouse::ClickHouseHttpFallback; +pub use cold_store::{ColdStore, ColdStoreError, LocalFsColdStore, RawSample}; pub use elastic::ElasticHttpFallback; pub use prometheus::PrometheusHttpFallback; +pub use s3_adapter::ColdFallback; diff --git a/asap-query-engine/src/drivers/query/fallback/s3_adapter.rs b/asap-query-engine/src/drivers/query/fallback/s3_adapter.rs new file mode 100644 index 00000000..64aa9c9c --- /dev/null +++ b/asap-query-engine/src/drivers/query/fallback/s3_adapter.rs @@ -0,0 +1,572 @@ +//! Cold-tier fallback adapter (§5.2 of the sketch-DB design). +//! +//! When a query hits a capability-miss — most notably a +//! `TimelineCoverage::Purged` segment whose sketch was aged out — +//! the server falls through to a [`FallbackClient`] impl. This +//! module provides [`ColdFallback`], which answers the query from +//! raw observability samples in a [`ColdStore`] (local FS today, +//! S3 tomorrow — identical key layout, see +//! [`super::cold_store::format`]). +//! +//! Supported query shapes (v1 paper scope): +//! +//! * bare instant vector selector: `metric_name{label="val",...}` +//! at time `t` — per-series latest value within `[t - 5m, t]` +//! (the Prometheus default lookback delta) +//! * scalar aggregation without grouping: +//! `sum|count|avg|min|max ( metric_name{...} )` over the same +//! instant vector +//! +//! Anything outside that surface delegates to the optional +//! `inner` fallback (typically a Prometheus proxy) — the cold +//! adapter chains with the existing §5.2 forwarding adapter +//! rather than replacing it. +//! +//! Accuracy: results computed here are exact on the set of +//! samples in the cold tier. The adapter does not attempt to +//! reconcile against missing/late samples — the raw tier is the +//! canonical source of truth per the paper architecture. + +use async_trait::async_trait; +use axum::http::StatusCode; +use promql_parser::label::MatchOp; +use promql_parser::parser::{Expr, VectorSelector}; +use serde_json::{json, Value}; +use std::collections::{BTreeMap, HashMap}; +use std::sync::Arc; +use tracing::{debug, warn}; + +use crate::drivers::query::adapters::{ParsedQueryRequest, PrometheusResponse}; + +use super::cold_store::{ColdStore, RawSample}; +use super::metrics::{BYTES_SERVED_COLD_TOTAL, QUERIES_COLD_TOTAL}; +use super::{FallbackClient, FallbackResponse}; + +/// Prometheus's default instant-query lookback delta (5 minutes). +/// Controls how far back the adapter scans the cold store to find +/// the most-recent sample per series. +const INSTANT_LOOKBACK_MS: i64 = 5 * 60 * 1_000; + +/// Cold-tier fallback client. Parametrised by the `ColdStore` +/// impl so the same adapter runs over `LocalFsColdStore` in tests +/// and over a (future) S3-backed store in production. +pub struct ColdFallback { + store: Arc, + /// Chain-of-responsibility: unsupported query shapes fall + /// through to this inner client. Typically a + /// [`PrometheusHttpFallback`](super::PrometheusHttpFallback) + /// pointing at the live Prometheus so development / demo + /// environments stay operational. + inner: Option>, +} + +impl ColdFallback { + pub fn new(store: Arc) -> Self { + Self { store, inner: None } + } + + pub fn with_inner(mut self, inner: Arc) -> Self { + self.inner = Some(inner); + self + } +} + +#[async_trait] +impl FallbackClient for ColdFallback { + async fn execute_query( + &self, + request: &ParsedQueryRequest, + ) -> Result { + let ast = match promql_parser::parser::parse(&request.query) { + Ok(a) => a, + Err(e) => { + warn!( + "cold fallback: PromQL parse failed ({}); delegating to inner", + e + ); + return self.delegate(request).await; + } + }; + + let plan = match plan_query(&ast) { + Some(p) => p, + None => { + debug!( + "cold fallback: unsupported query shape for '{}'; delegating to inner", + request.query + ); + return self.delegate(request).await; + } + }; + + let query_time_ms = (request.time * 1_000.0) as i64; + let start_ms = query_time_ms - INSTANT_LOOKBACK_MS; + let end_ms = query_time_ms + 1; + + let samples = match self.store.scan(&plan.metric, start_ms, end_ms).await { + Ok(s) => s, + Err(e) => { + warn!( + "cold fallback: scan failed for metric '{}': {}", + plan.metric, e + ); + return Err(StatusCode::INTERNAL_SERVER_ERROR); + } + }; + + // Telemetry — size is a stable proxy for "cold work" even + // when the post-filter result has zero samples. + let bytes = approx_wire_bytes(&samples); + let shape = plan.op.as_label(); + BYTES_SERVED_COLD_TOTAL + .with_label_values(&[plan.metric.as_str(), shape]) + .inc_by(bytes as f64); + QUERIES_COLD_TOTAL + .with_label_values(&[plan.metric.as_str(), shape]) + .inc(); + + let filtered = samples + .into_iter() + .filter(|s| plan.matches_labels(&s.labels)) + .collect::>(); + + let latest = latest_per_series(&filtered); + let data = compute_result(&plan, &latest, &request.query, query_time_ms); + let resp = PrometheusResponse::success(data); + let value = serde_json::to_value(resp).unwrap_or_else(|_| json!({"status":"error"})); + Ok(FallbackResponse::Json(value)) + } + + async fn execute_query_with_headers( + &self, + request: &ParsedQueryRequest, + _headers: HashMap, + ) -> Result { + self.execute_query(request).await + } + + async fn get_runtime_info(&self) -> Result { + // Runtime info is meaningless for a cold store; defer to + // inner if configured, else return empty. + match &self.inner { + Some(inner) => inner.get_runtime_info().await, + None => Ok(json!({})), + } + } +} + +impl ColdFallback { + async fn delegate(&self, request: &ParsedQueryRequest) -> Result { + match &self.inner { + Some(inner) => inner.execute_query(request).await, + None => { + // No chain — return an empty Prometheus success + // response rather than a 5xx; the adapter is + // advisory for query shapes it can't handle. + let resp = PrometheusResponse::success(json!({ + "resultType": "vector", + "result": [] + })); + Ok(FallbackResponse::Json( + serde_json::to_value(resp).unwrap_or(json!({"status":"error"})), + )) + } + } + } +} + +/// The query op we recognise for cold evaluation. Kept narrow so +/// the adapter's behaviour is obvious from the outside — anything +/// else delegates. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum QueryOp { + /// Bare instant vector (`metric{labels}`). + Selector, + Sum, + Count, + Avg, + Min, + Max, +} + +impl QueryOp { + fn as_label(&self) -> &'static str { + match self { + QueryOp::Selector => "selector", + QueryOp::Sum => "sum", + QueryOp::Count => "count", + QueryOp::Avg => "avg", + QueryOp::Min => "min", + QueryOp::Max => "max", + } + } + + fn from_agg_str(s: &str) -> Option { + match s.to_ascii_lowercase().as_str() { + "sum" => Some(QueryOp::Sum), + "count" => Some(QueryOp::Count), + "avg" => Some(QueryOp::Avg), + "min" => Some(QueryOp::Min), + "max" => Some(QueryOp::Max), + _ => None, + } + } +} + +/// Matcher set we understand. Regex matchers (`=~`, `!~`) are +/// deliberately out of scope for v1 — passing a regex matcher +/// causes the adapter to delegate upstream. +#[derive(Debug, Clone)] +struct LabelPredicate { + name: String, + value: String, + equals: bool, +} + +struct QueryPlan { + metric: String, + predicates: Vec, + op: QueryOp, +} + +impl QueryPlan { + fn matches_labels(&self, labels: &BTreeMap) -> bool { + for p in &self.predicates { + let hit = labels.get(&p.name).map(|v| v == &p.value).unwrap_or(false); + if p.equals && !hit { + return false; + } + if !p.equals && hit { + return false; + } + } + true + } +} + +/// Inspect the PromQL AST and, if it's a shape we support, return +/// the extracted `(metric, predicates, op)`. Returns `None` +/// otherwise — caller delegates to inner fallback. +fn plan_query(ast: &Expr) -> Option { + match ast { + Expr::VectorSelector(vs) => { + let (metric, predicates) = selector_to_plan(vs)?; + Some(QueryPlan { + metric, + predicates, + op: QueryOp::Selector, + }) + } + Expr::Paren(p) => plan_query(&p.expr), + Expr::Aggregate(agg) => { + // Only recognise no-grouping-modifier, no-param aggs — + // `sum by (...)` / `topk(k, expr)` exceed v1 scope. + if agg.modifier.is_some() { + return None; + } + if agg.param.is_some() { + return None; + } + let op = QueryOp::from_agg_str(&agg.op.to_string())?; + let vs = match agg.expr.as_ref() { + Expr::VectorSelector(vs) => vs, + Expr::Paren(p) => match p.expr.as_ref() { + Expr::VectorSelector(vs) => vs, + _ => return None, + }, + _ => return None, + }; + let (metric, predicates) = selector_to_plan(vs)?; + Some(QueryPlan { + metric, + predicates, + op, + }) + } + _ => None, + } +} + +fn selector_to_plan(vs: &VectorSelector) -> Option<(String, Vec)> { + let metric = vs.name.clone()?; + let mut predicates = Vec::new(); + for m in &vs.matchers.matchers { + // __name__ is already captured via vs.name — skip any + // explicit __name__ matcher that duplicates it. + if m.name == "__name__" { + continue; + } + let equals = match m.op { + MatchOp::Equal => true, + MatchOp::NotEqual => false, + _ => return None, // regex matchers out of scope + }; + predicates.push(LabelPredicate { + name: m.name.clone(), + value: m.value.clone(), + equals, + }); + } + Some((metric, predicates)) +} + +/// Reduce a pool of samples to one-per-unique-label-set, taking +/// the sample with the largest `ts_ms`. Mirrors Prometheus's +/// instant-vector semantics. +fn latest_per_series(samples: &[RawSample]) -> Vec { + let mut seen: HashMap, RawSample> = HashMap::new(); + for s in samples { + let key: Vec<_> = s + .labels + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + seen.entry(key) + .and_modify(|existing| { + if s.ts_ms > existing.ts_ms { + *existing = s.clone(); + } + }) + .or_insert_with(|| s.clone()); + } + seen.into_values().collect() +} + +/// Build the Prometheus HTTP `data` payload for the query result. +/// * `Selector` → vector with one entry per series. +/// * `sum|count|avg|min|max` → single-element vector with no +/// `__name__` label (matches `sum(foo)` Prometheus output). +fn compute_result( + plan: &QueryPlan, + latest: &[RawSample], + _query: &str, + query_time_ms: i64, +) -> Value { + let ts_s = (query_time_ms as f64) / 1_000.0; + match plan.op { + QueryOp::Selector => { + let items = latest + .iter() + .map(|s| { + let mut m = serde_json::Map::new(); + m.insert("__name__".to_string(), Value::String(plan.metric.clone())); + for (k, v) in &s.labels { + m.insert(k.clone(), Value::String(v.clone())); + } + json!({ + "metric": Value::Object(m), + "value": [ts_s, format_prom_float(s.value)], + }) + }) + .collect::>(); + json!({ + "resultType": "vector", + "result": items, + }) + } + op => { + let agg = aggregate(op, latest); + let value = match agg { + Some(v) => json!([ts_s, format_prom_float(v)]), + // Empty-input semantics: Prometheus returns empty + // result for sum/min/max/avg over no samples. + None => { + return json!({ + "resultType": "vector", + "result": [], + }) + } + }; + json!({ + "resultType": "vector", + "result": [ + { + "metric": {}, + "value": value, + } + ], + }) + } + } +} + +fn aggregate(op: QueryOp, latest: &[RawSample]) -> Option { + if latest.is_empty() { + // count over empty = 0 is what Prometheus does; other ops + // return empty vector. + return match op { + QueryOp::Count => Some(0.0), + _ => None, + }; + } + let xs = latest.iter().map(|s| s.value); + Some(match op { + QueryOp::Sum => xs.sum(), + QueryOp::Count => latest.len() as f64, + QueryOp::Avg => latest.iter().map(|s| s.value).sum::() / (latest.len() as f64), + QueryOp::Min => xs.fold(f64::INFINITY, f64::min), + QueryOp::Max => xs.fold(f64::NEG_INFINITY, f64::max), + QueryOp::Selector => unreachable!("Selector handled in compute_result"), + }) +} + +/// Stable, Prometheus-style float formatting for the HTTP JSON +/// `value` field — integers render without a trailing `.0`. +fn format_prom_float(v: f64) -> String { + if v.is_nan() { + return "NaN".to_string(); + } + if v.is_infinite() { + return if v > 0.0 { + "+Inf".into() + } else { + "-Inf".into() + }; + } + if v == v.trunc() && v.abs() < 1e15 { + format!("{}", v as i64) + } else { + format!("{v}") + } +} + +fn approx_wire_bytes(samples: &[RawSample]) -> usize { + // Rough proxy: one JSON line per sample. Avoids re-serialising + // every scan into memory just to count. Good enough for + // observability dashboards. + samples + .iter() + .map(|s| { + let label_bytes: usize = s + .labels + .iter() + .map(|(k, v)| k.len() + v.len() + 6) // "k":"v", + .sum(); + // {"ts_ms":<13>,"labels":{...},"value":<~20>} + newline + 13 + 12 + label_bytes + 12 + 20 + 2 + }) + .sum() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rs(ts_ms: i64, labels: &[(&str, &str)], value: f64) -> RawSample { + RawSample { + ts_ms, + labels: labels + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect(), + value, + } + } + + #[test] + fn plan_bare_selector() { + let ast = promql_parser::parser::parse("up{zone=\"a\"}").unwrap(); + let plan = plan_query(&ast).unwrap(); + assert_eq!(plan.metric, "up"); + assert_eq!(plan.predicates.len(), 1); + assert_eq!(plan.op, QueryOp::Selector); + } + + #[test] + fn plan_sum_over_selector() { + let ast = promql_parser::parser::parse("sum(up)").unwrap(); + let plan = plan_query(&ast).unwrap(); + assert_eq!(plan.op, QueryOp::Sum); + } + + #[test] + fn plan_sum_by_is_unsupported() { + let ast = promql_parser::parser::parse("sum by (zone) (up)").unwrap(); + assert!(plan_query(&ast).is_none()); + } + + #[test] + fn plan_rate_is_unsupported() { + let ast = promql_parser::parser::parse("rate(up[1m])").unwrap(); + assert!(plan_query(&ast).is_none()); + } + + #[test] + fn plan_regex_matcher_is_unsupported() { + let ast = promql_parser::parser::parse("up{zone=~\"a.*\"}").unwrap(); + assert!(plan_query(&ast).is_none()); + } + + #[test] + fn latest_per_series_picks_newest() { + let ss = vec![ + rs(100, &[("zone", "a")], 1.0), + rs(200, &[("zone", "a")], 2.0), + rs(150, &[("zone", "b")], 9.0), + ]; + let got = latest_per_series(&ss); + assert_eq!(got.len(), 2); + let a = got.iter().find(|s| s.labels["zone"] == "a").unwrap(); + let b = got.iter().find(|s| s.labels["zone"] == "b").unwrap(); + assert_eq!(a.value, 2.0); + assert_eq!(b.value, 9.0); + } + + #[test] + fn aggregate_ops() { + let samples = vec![ + rs(1, &[("z", "a")], 1.0), + rs(1, &[("z", "b")], 2.0), + rs(1, &[("z", "c")], 3.0), + ]; + assert_eq!(aggregate(QueryOp::Sum, &samples), Some(6.0)); + assert_eq!(aggregate(QueryOp::Count, &samples), Some(3.0)); + assert_eq!(aggregate(QueryOp::Avg, &samples), Some(2.0)); + assert_eq!(aggregate(QueryOp::Min, &samples), Some(1.0)); + assert_eq!(aggregate(QueryOp::Max, &samples), Some(3.0)); + } + + #[test] + fn aggregate_empty_count_is_zero() { + assert_eq!(aggregate(QueryOp::Count, &[]), Some(0.0)); + assert_eq!(aggregate(QueryOp::Sum, &[]), None); + } + + #[test] + fn format_prom_float_matches_prometheus_conventions() { + assert_eq!(format_prom_float(1.0), "1"); + assert_eq!(format_prom_float(1.5), "1.5"); + assert_eq!(format_prom_float(f64::INFINITY), "+Inf"); + assert_eq!(format_prom_float(f64::NAN), "NaN"); + } + + #[test] + fn label_predicate_equal_and_not_equal() { + let plan = QueryPlan { + metric: "m".into(), + predicates: vec![ + LabelPredicate { + name: "zone".into(), + value: "a".into(), + equals: true, + }, + LabelPredicate { + name: "env".into(), + value: "prod".into(), + equals: false, + }, + ], + op: QueryOp::Selector, + }; + let match_a: BTreeMap = [("zone", "a"), ("env", "dev")] + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + let no_match: BTreeMap = [("zone", "a"), ("env", "prod")] + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + assert!(plan.matches_labels(&match_a)); + assert!(!plan.matches_labels(&no_match)); + } +} diff --git a/asap-query-engine/src/tests/cold_fallback_tests.rs b/asap-query-engine/src/tests/cold_fallback_tests.rs new file mode 100644 index 00000000..2c23f10a --- /dev/null +++ b/asap-query-engine/src/tests/cold_fallback_tests.rs @@ -0,0 +1,402 @@ +//! End-to-end tests for the §5.2 cold-query fallback. +//! +//! Exercises the full HTTP → engine-miss → [`ColdFallback`] path +//! with raw sample fixtures on a local-FS [`ColdStore`]. The +//! format used here (hour-bucketed JSONL) is byte-identical to +//! what a future S3 cold adapter will read, so these tests +//! double as format-lock tests for the on-disk / on-object layout. +//! +//! Coverage: +//! * bare instant vector → cold path, correct per-series latest +//! * `sum(...)` → cold path, correct scalar +//! * unsupported query shape (e.g. `rate(...)`) falls through to +//! the inner Prometheus chain +//! * telemetry counters increment on cold hits +//! * purged-time-range semantics: sketch-absent data served from cold + +#[cfg(test)] +use crate::data_model::{CleanupPolicy, InferenceConfig, QueryLanguage, StreamingConfig}; +use crate::drivers::query::adapters::AdapterConfig; +use crate::drivers::query::fallback::cold_store::format::RawSample; +use crate::drivers::query::fallback::metrics::{BYTES_SERVED_COLD_TOTAL, QUERIES_COLD_TOTAL}; +use crate::drivers::query::fallback::{ + ColdFallback, FallbackClient, LocalFsColdStore, PrometheusHttpFallback, +}; +use crate::drivers::query::servers::http::{HttpServer, HttpServerConfig}; +use crate::engines::SimpleEngine; +use crate::stores::sketch_db::simple_map_store::SimpleMapStore; +use chrono::{TimeZone, Utc}; +use reqwest::Client; +use serde_json::Value; +use std::collections::BTreeMap; +use std::path::Path; +use std::sync::Arc; +use tempfile::TempDir; +use tokio::net::TcpListener; +use tokio::time::{sleep, Duration}; + +/// Build a [`RawSample`] with ergonomic literal labels. +fn sample(ts_ms: i64, labels: &[(&str, &str)], value: f64) -> RawSample { + RawSample { + ts_ms, + labels: labels + .iter() + .map(|(k, v)| ((*k).to_string(), (*v).to_string())) + .collect::>(), + value, + } +} + +/// Write a JSONL part under the hour-bucket prefix (same key +/// layout as S3). +async fn write_part(root: &Path, rel: &str, lines: &[RawSample]) { + let dir = root.join(rel); + tokio::fs::create_dir_all(&dir).await.unwrap(); + let mut buf = String::new(); + for s in lines { + buf.push_str(&serde_json::to_string(s).unwrap()); + buf.push('\n'); + } + tokio::fs::write(dir.join("part-000001.jsonl"), buf) + .await + .unwrap(); +} + +fn make_engine_and_store() -> (Arc, Arc) { + let inference_config = InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::NoCleanup); + let streaming_config = Arc::new(StreamingConfig::default()); + let store = Arc::new(SimpleMapStore::new( + streaming_config.clone(), + CleanupPolicy::NoCleanup, + )); + let engine = Arc::new(SimpleEngine::new( + store.clone(), + inference_config, + streaming_config.clone(), + 15000, + QueryLanguage::promql, + )); + (engine, store) +} + +/// Start an HTTP server whose fallback chain is +/// `ColdFallback(LocalFsColdStore) → None`. Returns the server +/// port. No sketch data is ingested, so every query +/// capability-misses and is served from cold. +async fn start_cold_only_server(cold_root: &Path) -> u16 { + let cold_store = Arc::new(LocalFsColdStore::new(cold_root)); + let cold = Arc::new(ColdFallback::new(cold_store)) as Arc; + + let adapter_config = AdapterConfig::new( + crate::data_model::enums::QueryProtocol::PrometheusHttp, + QueryLanguage::promql, + Some(cold), + ); + let config = HttpServerConfig { + port: 0, + handle_http_requests: true, + adapter_config, + }; + let (engine, store) = make_engine_and_store(); + let server = HttpServer::new(config, engine, store, None); + server + .start_test_server() + .await + .expect("Failed to start test server") +} + +/// Mock upstream Prometheus that returns a fixed marker body, used +/// to check that unsupported shapes fall through the cold adapter +/// to the inner chain. +async fn start_mock_prometheus(port: u16, marker: &'static str) { + use axum::{routing::get, Json, Router}; + use serde_json::json; + async fn h(marker: &'static str) -> Json { + Json(json!({ + "status": "success", + "data": {"resultType":"scalar", "result":[0, marker]} + })) + } + let app = Router::new().route("/api/v1/query", get(move || h(marker))); + let listener = TcpListener::bind(format!("127.0.0.1:{port}")) + .await + .unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + sleep(Duration::from_millis(50)).await; +} + +async fn start_chained_server(cold_root: &Path, prom_url: String) -> u16 { + let cold_store = Arc::new(LocalFsColdStore::new(cold_root)); + let prom: Arc = Arc::new(PrometheusHttpFallback::new(prom_url)); + let cold = Arc::new(ColdFallback::new(cold_store).with_inner(prom)) as Arc; + let adapter_config = AdapterConfig::new( + crate::data_model::enums::QueryProtocol::PrometheusHttp, + QueryLanguage::promql, + Some(cold), + ); + let config = HttpServerConfig { + port: 0, + handle_http_requests: true, + adapter_config, + }; + let (engine, store) = make_engine_and_store(); + let server = HttpServer::new(config, engine, store, None); + server + .start_test_server() + .await + .expect("Failed to start test server") +} + +#[tokio::test] +async fn cold_fallback_serves_bare_selector_from_raw_samples() { + let tmp = TempDir::new().unwrap(); + // 2026-04-21 08:00:00 UTC + let base = Utc + .with_ymd_and_hms(2026, 4, 21, 8, 0, 0) + .unwrap() + .timestamp_millis(); + write_part( + tmp.path(), + "raw/http_requests_total/2026/04/21/08/", + &[ + sample(base + 10_000, &[("zone", "a")], 1.0), + sample(base + 20_000, &[("zone", "a")], 2.0), // latest for zone=a + sample(base + 15_000, &[("zone", "b")], 9.0), + ], + ) + .await; + + let server_port = start_cold_only_server(tmp.path()).await; + let client = Client::new(); + + // Query time = base + 30s. Lookback window covers all three samples. + let query_time = (base + 30_000) as f64 / 1_000.0; + let resp = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) + .query(&[ + ("query", "http_requests_total".to_string()), + ("time", query_time.to_string()), + ]) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), reqwest::StatusCode::OK); + let body: Value = resp.json().await.unwrap(); + + assert_eq!(body["status"], "success"); + assert_eq!(body["data"]["resultType"], "vector"); + let items = body["data"]["result"].as_array().unwrap(); + assert_eq!(items.len(), 2); + // Find zone=a entry and verify it got the latest value. + let za = items + .iter() + .find(|v| v["metric"]["zone"] == "a") + .expect("zone=a series"); + assert_eq!(za["metric"]["__name__"], "http_requests_total"); + assert_eq!(za["value"][1], "2"); +} + +#[tokio::test] +async fn cold_fallback_serves_sum_aggregation() { + let tmp = TempDir::new().unwrap(); + let base = Utc + .with_ymd_and_hms(2026, 4, 21, 8, 0, 0) + .unwrap() + .timestamp_millis(); + write_part( + tmp.path(), + "raw/cpu_seconds_total/2026/04/21/08/", + &[ + sample(base + 1_000, &[("pod", "a")], 10.0), + sample(base + 2_000, &[("pod", "b")], 20.0), + sample(base + 3_000, &[("pod", "c")], 30.0), + ], + ) + .await; + + let server_port = start_cold_only_server(tmp.path()).await; + let client = Client::new(); + + let query_time = (base + 10_000) as f64 / 1_000.0; + let resp = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) + .query(&[ + ("query", "sum(cpu_seconds_total)".to_string()), + ("time", query_time.to_string()), + ]) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), reqwest::StatusCode::OK); + let body: Value = resp.json().await.unwrap(); + + assert_eq!(body["status"], "success"); + assert_eq!(body["data"]["resultType"], "vector"); + let items = body["data"]["result"].as_array().unwrap(); + assert_eq!(items.len(), 1); + assert_eq!(items[0]["value"][1], "60"); + // Aggregation without grouping → empty metric labels. + let metric = items[0]["metric"].as_object().unwrap(); + assert!(metric.is_empty()); +} + +#[tokio::test] +async fn cold_fallback_label_matcher_filters() { + let tmp = TempDir::new().unwrap(); + let base = Utc + .with_ymd_and_hms(2026, 4, 21, 8, 0, 0) + .unwrap() + .timestamp_millis(); + write_part( + tmp.path(), + "raw/requests_total/2026/04/21/08/", + &[ + sample(base + 1_000, &[("zone", "a")], 1.0), + sample(base + 2_000, &[("zone", "b")], 2.0), + sample(base + 3_000, &[("zone", "c")], 3.0), + ], + ) + .await; + + let server_port = start_cold_only_server(tmp.path()).await; + let client = Client::new(); + + let query_time = (base + 10_000) as f64 / 1_000.0; + let resp = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) + .query(&[ + ("query", "sum(requests_total{zone=\"b\"})".to_string()), + ("time", query_time.to_string()), + ]) + .send() + .await + .unwrap(); + let body: Value = resp.json().await.unwrap(); + let items = body["data"]["result"].as_array().unwrap(); + assert_eq!(items.len(), 1); + assert_eq!(items[0]["value"][1], "2"); +} + +#[tokio::test] +async fn cold_fallback_unsupported_query_delegates_to_inner() { + let tmp = TempDir::new().unwrap(); + // Pick a deterministic port for the mock Prom — low risk of + // collision with the rest of the suite since each test picks + // a different one. + let prom_port = 19_201; + start_mock_prometheus(prom_port, "MARKER_INNER").await; + + let server_port = + start_chained_server(tmp.path(), format!("http://127.0.0.1:{prom_port}")).await; + let client = Client::new(); + + // rate(...) is not a shape we handle in cold — expect the + // inner Prometheus mock to answer. + let resp = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) + .query(&[("query", "rate(foo[1m])"), ("time", "1000")]) + .send() + .await + .unwrap(); + let body: Value = resp.json().await.unwrap(); + assert_eq!(body["status"], "success"); + // Mock Prom returns a scalar with "MARKER_INNER" at index [1]. + assert_eq!(body["data"]["result"][1], "MARKER_INNER"); +} + +#[tokio::test] +async fn cold_fallback_increments_telemetry_counters() { + let tmp = TempDir::new().unwrap(); + let base = Utc + .with_ymd_and_hms(2026, 4, 21, 8, 0, 0) + .unwrap() + .timestamp_millis(); + write_part( + tmp.path(), + "raw/telemetry_test_metric/2026/04/21/08/", + &[sample(base + 1_000, &[("zone", "a")], 1.0)], + ) + .await; + + let before_q = QUERIES_COLD_TOTAL + .with_label_values(&["telemetry_test_metric", "sum"]) + .get(); + let before_b = BYTES_SERVED_COLD_TOTAL + .with_label_values(&["telemetry_test_metric", "sum"]) + .get(); + + let server_port = start_cold_only_server(tmp.path()).await; + let client = Client::new(); + let query_time = (base + 10_000) as f64 / 1_000.0; + let _ = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) + .query(&[ + ("query", "sum(telemetry_test_metric)".to_string()), + ("time", query_time.to_string()), + ]) + .send() + .await + .unwrap(); + + let after_q = QUERIES_COLD_TOTAL + .with_label_values(&["telemetry_test_metric", "sum"]) + .get(); + let after_b = BYTES_SERVED_COLD_TOTAL + .with_label_values(&["telemetry_test_metric", "sum"]) + .get(); + assert!( + after_q >= before_q + 1.0, + "expected cold queries counter to advance: before={before_q} after={after_q}" + ); + assert!( + after_b > before_b, + "expected cold bytes-served counter to advance" + ); +} + +#[tokio::test] +async fn cold_fallback_purged_time_range_served_from_raw() { + // Simulates the §5.2 "Purged segment" path: no sketch exists + // for the queried range (nothing ingested into the engine), + // but the raw tier has samples, and the cold adapter recovers + // the exact answer. This is the canonical paper claim. + let tmp = TempDir::new().unwrap(); + let base = Utc + .with_ymd_and_hms(2026, 4, 21, 8, 0, 0) + .unwrap() + .timestamp_millis(); + // Three distinct series so `avg` runs over three latest-per-series + // values — matches Prometheus instant-vector semantics. + write_part( + tmp.path(), + "raw/purged_metric/2026/04/21/08/", + &[ + sample(base + 1_000, &[("pod", "a")], 10.0), + sample(base + 2_000, &[("pod", "b")], 20.0), + sample(base + 3_000, &[("pod", "c")], 30.0), + ], + ) + .await; + + let server_port = start_cold_only_server(tmp.path()).await; + let client = Client::new(); + let query_time = (base + 10_000) as f64 / 1_000.0; + + // avg over the cold raw samples. + let resp = client + .get(format!("http://127.0.0.1:{server_port}/api/v1/query")) + .query(&[ + ("query", "avg(purged_metric)".to_string()), + ("time", query_time.to_string()), + ]) + .send() + .await + .unwrap(); + let body: Value = resp.json().await.unwrap(); + let items = body["data"]["result"].as_array().unwrap(); + assert_eq!(items.len(), 1); + // avg(10, 20, 30) = 20. + assert_eq!(items[0]["value"][1], "20"); +} diff --git a/asap-query-engine/src/tests/mod.rs b/asap-query-engine/src/tests/mod.rs index beed0b52..3abb8d0d 100644 --- a/asap-query-engine/src/tests/mod.rs +++ b/asap-query-engine/src/tests/mod.rs @@ -1,5 +1,6 @@ pub mod capability_matching_tests; pub mod clickhouse_forwarding_tests; +pub mod cold_fallback_tests; pub mod datafusion; pub mod elastic_dsl_query_tests; pub mod elastic_forwarding_tests;