diff --git a/asap-common/dependencies/rs/asap_types/src/streaming_config.rs b/asap-common/dependencies/rs/asap_types/src/streaming_config.rs index 88d4a8f9..4ab02498 100644 --- a/asap-common/dependencies/rs/asap_types/src/streaming_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/streaming_config.rs @@ -50,20 +50,33 @@ impl StreamingConfig { data: &Value, inference_config: Option<&InferenceConfig>, ) -> Result { - let mut retention_map: HashMap = HashMap::new(); + // None = unbounded, which wins over any bounded count from another reference. + let mut retention_map: HashMap> = HashMap::new(); let mut read_count_threshold_map: HashMap = HashMap::new(); if let Some(inference_config) = inference_config { for query_config in &inference_config.query_configs { for aggregation in &query_config.aggregations { let aggregation_id = aggregation.aggregation_id; - if let Some(num_aggregates) = aggregation.num_aggregates_to_retain { - // OLD: Keep last value only (for backwards compatibility) - retention_map.insert(aggregation_id, num_aggregates); - - // NEW: Sum up num_aggregates_to_retain across all queries - *read_count_threshold_map.entry(aggregation_id).or_insert(0) += - num_aggregates; + let num_aggregates = aggregation.num_aggregates_to_retain; + retention_map + .entry(aggregation_id) + .and_modify(|retain| { + *retain = match (*retain, num_aggregates) { + (Some(a), Some(b)) => Some(a.max(b)), + _ => None, + } + }) + .or_insert(num_aggregates); + + if let Some(read_count_threshold) = aggregation.read_count_threshold { + let threshold = read_count_threshold_map.entry(aggregation_id).or_insert(0); + *threshold = + threshold.checked_add(read_count_threshold).ok_or_else(|| { + anyhow::anyhow!( + "read-count threshold for aggregation {aggregation_id} overflowed" + ) + })?; } } } @@ -90,11 +103,12 @@ impl StreamingConfig { aggregation_id ) })?; - let num_aggregates_to_retain = retention_map.get(&aggregation_id_u64); + let num_aggregates_to_retain = + retention_map.get(&aggregation_id_u64).copied().flatten(); let read_count_threshold = read_count_threshold_map.get(&aggregation_id_u64); let config = AggregationConfig::from_yaml_data( aggregation_data, - num_aggregates_to_retain.copied(), + num_aggregates_to_retain, read_count_threshold.copied(), query_language, )?; @@ -136,6 +150,10 @@ impl Default for StreamingConfig { mod tests { use super::*; use crate::aggregation_config::AggregationConfigError; + use crate::aggregation_reference::AggregationReference; + use crate::enums::{CleanupPolicy, QueryLanguage}; + use crate::inference_config::InferenceConfig; + use crate::query_config::QueryConfig; use promql_utilities::query_logics::enums::AggregationType; #[test] @@ -422,4 +440,87 @@ aggregations: ); } } + + #[test] + fn shared_references_keep_the_largest_circular_buffer_retention() { + // A later short-range query must not reduce a shared aggregation's retention. + let yaml = minmax_yaml("MinMax", "max"); + for references in [[7, 2], [2, 7]] { + let mut inference = + InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::CircularBuffer); + inference.query_configs = references + .into_iter() + .map(|retain| { + QueryConfig::new(format!("max(metric[{retain}m])")) + .add_aggregation(AggregationReference::new(1, Some(retain))) + }) + .collect(); + + let config = StreamingConfig::from_yaml_data(&yaml, Some(&inference)) + .expect("shared aggregation references should load"); + + assert_eq!(config[1].num_aggregates_to_retain, Some(7)); + assert_eq!(config[1].read_count_threshold, None); + } + } + + #[test] + fn shared_references_keep_unbounded_retention() { + // A reference without a retained count needs every window; a bounded + // reference to the same aggregation must not cap it, in either order. + let yaml = minmax_yaml("MinMax", "max"); + for references in [[None, Some(2)], [Some(2), None]] { + let mut inference = + InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::CircularBuffer); + inference.query_configs = references + .into_iter() + .enumerate() + .map(|(i, retain)| { + QueryConfig::new(format!("max(metric[{i}m])")) + .add_aggregation(AggregationReference::new(1, retain)) + }) + .collect(); + + let config = StreamingConfig::from_yaml_data(&yaml, Some(&inference)) + .expect("shared aggregation references should load"); + + assert_eq!(config[1].num_aggregates_to_retain, None, "{references:?}"); + } + } + + #[test] + fn shared_references_sum_read_based_thresholds() { + // Each shared query consumes a read before the aggregate can be cleaned up. + let yaml = minmax_yaml("MinMax", "max"); + let mut inference = InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::ReadBased); + inference.query_configs = vec![ + QueryConfig::new("max(metric[10m])".into()) + .add_aggregation(AggregationReference::with_read_count_threshold(1, Some(7))), + QueryConfig::new("max(metric[1m])".into()) + .add_aggregation(AggregationReference::with_read_count_threshold(1, Some(2))), + ]; + + let config = StreamingConfig::from_yaml_data(&yaml, Some(&inference)) + .expect("shared aggregation references should load"); + + assert_eq!(config[1].num_aggregates_to_retain, None); + assert_eq!(config[1].read_count_threshold, Some(9)); + } + + #[test] + fn overflowing_read_based_threshold_is_an_error() { + let yaml = minmax_yaml("MinMax", "max"); + let mut inference = InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::ReadBased); + inference.query_configs = vec![ + QueryConfig::new("max(metric[10m])".into()).add_aggregation( + AggregationReference::with_read_count_threshold(1, Some(u64::MAX)), + ), + QueryConfig::new("max(metric[1m])".into()) + .add_aggregation(AggregationReference::with_read_count_threshold(1, Some(1))), + ]; + + let error = StreamingConfig::from_yaml_data(&yaml, Some(&inference)) + .expect_err("an overflowing threshold must not load"); + assert!(error.to_string().contains("overflowed"), "{error}"); + } } diff --git a/asap-planner-rs/src/optimizer/milp_output.rs b/asap-planner-rs/src/optimizer/milp_output.rs index 4b6653f4..5e33ce98 100644 --- a/asap-planner-rs/src/optimizer/milp_output.rs +++ b/asap-planner-rs/src/optimizer/milp_output.rs @@ -33,8 +33,8 @@ pub enum MilpOutputError { #[error("query {0:?} uses avg, which the engine can't answer from sum and count yet")] AvgQuery(String), #[error( - "query groups set step_ms or range_duration_ms, whose retention MILP configs \ - don't set yet (#800): {0:?}" + "query groups set step_ms or range_duration_ms, whose extra retention MILP \ + configs don't size: {0:?}" )] RangeQueryOverrides(Vec), #[error("sketch-bench variant {variant} ({capability:?}) has no ASAPQuery aggregation type")] @@ -73,6 +73,16 @@ pub enum MilpOutputError { }, #[error("query {0:?} is a roll-up but not an aggregation whose grouping can be rewritten")] UnregroupableQuery(String), + #[error( + "query {query:?} looks back {lookback_ms}ms, which isn't whole windows of its \ + deployment (window {window_ms}ms, slide {slide_ms}ms)" + )] + PartialWindowLookback { + query: String, + lookback_ms: u64, + window_ms: u64, + slide_ms: u64, + }, #[error(transparent)] Generator(#[from] ControllerError), } @@ -119,8 +129,9 @@ pub fn reject_unwritable_queries(config: MilpConfig) -> Result<(), MilpOutputErr } /// Streaming and inference YAML for `solution`. Aggregation ids follow the -/// plan's deployment order, starting at 1. No cleanup policy yet, so the -/// retained instance counts are not emitted. +/// plan's deployment order, starting at 1. Cleanup is CircularBuffer: each +/// query retains the closed windows its lookback needs, and a shared +/// aggregation keeps the largest of its queries' counts. pub fn plan_to_planner_output( config: MilpConfig, workload: &MilpWorkload, @@ -227,12 +238,12 @@ pub fn deployments_to_planner_output( let schema = config.schema_from_hints(); ( build_streaming_yaml(&aggregations, &id_map, &schema)?, - build_inference_yaml(CleanupPolicy::NoCleanup, &queries, &id_map, &schema)?, + build_inference_yaml(CleanupPolicy::CircularBuffer, &queries, &id_map, &schema)?, ) } MilpConfig::Sql(config) => ( build_sql_streaming_yaml(config, &aggregations, &id_map)?, - build_sql_inference_yaml(config, CleanupPolicy::NoCleanup, &queries, &id_map)?, + build_sql_inference_yaml(config, CleanupPolicy::CircularBuffer, &queries, &id_map)?, ), }; Ok(PlannerOutput::from_output(GeneratorOutput { @@ -277,7 +288,17 @@ fn query_plan_entry( aggregation_type: AggregationType, stage_blocker: Option, ) -> Result { - let aggregation_keys = vec![(key.to_string(), None)]; + // The store holds closed windows only; open ones live in the precompute + // engine. + let retained = deployment + .closed_instance_count(raqe.lookback_ms) + .ok_or_else(|| MilpOutputError::PartialWindowLookback { + query: query.to_string(), + lookback_ms: raqe.lookback_ms, + window_ms: deployment.window_ms, + slide_ms: deployment.slide_ms, + })?; + let aggregation_keys = vec![(key.to_string(), Some(retained))]; let (fine, coarse) = (&deployment.grouping_labels, &raqe.grouping_labels); if coarse == fine { return Ok(QueryPlanEntry::fully_planned( @@ -650,11 +671,58 @@ mod tests { for query in [sum, count, quantile, topk_value, topk_count] { let a = aggregation_for(&output, query); assert_eq!(a.metric, "http_requests_total"); - // No cleanup policy yet: retention is unbounded. - assert_eq!(a.num_aggregates_to_retain, None); + assert_eq!( + output + .to_inference_config(QueryLanguage::promql) + .unwrap() + .cleanup_policy, + CleanupPolicy::CircularBuffer + ); + assert!(a.num_aggregates_to_retain.is_some_and(|n| n >= 1)); + } + } + + /// Each query retains the closed windows wholly inside its lookback, + /// `(L − x) / y + 1`; open windows live in the precompute engine. + #[test] + fn retention_is_the_closed_windows_in_the_lookback() { + let query = "quantile_over_time(0.99, http_requests_total[10m])"; + for (window_ms, slide_ms, retained) in [ + (60_000, 60_000, 10), + (600_000, 600_000, 1), + (300_000, 60_000, 6), + ] { + let (config, workload, mut solution) = plan(&group(query, 0.99)); + let deployment = &mut solution.deployments[0].deployment; + deployment.window_ms = window_ms; + deployment.slide_ms = slide_ms; + let output = + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution).unwrap(); + assert_eq!( + aggregation_for(&output, query).num_aggregates_to_retain, + Some(retained), + "window {window_ms}ms, slide {slide_ms}ms" + ); } } + #[test] + fn lookback_not_in_whole_windows_is_an_error() { + let query = "quantile_over_time(0.99, http_requests_total[10m])"; + let (config, workload, mut solution) = plan(&group(query, 0.99)); + let deployment = &mut solution.deployments[0].deployment; + deployment.window_ms = 420_000; + deployment.slide_ms = 420_000; + assert!(matches!( + plan_to_planner_output(MilpConfig::Promql(&config), &workload, &solution), + Err(MilpOutputError::PartialWindowLookback { + lookback_ms: 600_000, + window_ms: 420_000, + .. + }) + )); + } + /// What the engine builds from an aggregation, without its id, metric /// and SQL columns. fn layout(a: &AggregationConfig) -> String { diff --git a/asap-query-engine/examples/engine_config.yaml b/asap-query-engine/examples/engine_config.yaml index c6a99d6a..5cee373a 100644 --- a/asap-query-engine/examples/engine_config.yaml +++ b/asap-query-engine/examples/engine_config.yaml @@ -65,6 +65,7 @@ backend: store: lock_strategy: "per-key" # global | per-key + # spill_dir: "./spill" # append evicted windows here; unset drops them # --------------------------------------------------------------------------- # Ingest source — choose exactly one type. diff --git a/asap-query-engine/src/engine_config.rs b/asap-query-engine/src/engine_config.rs index 9388dc25..4c7a2cbd 100644 --- a/asap-query-engine/src/engine_config.rs +++ b/asap-query-engine/src/engine_config.rs @@ -228,12 +228,19 @@ fn default_elastic_url() -> String { #[serde(default)] pub struct StoreSettings { pub lock_strategy: LockStrategy, + /// Directory evicted windows are appended to; unset drops them. + pub spill_dir: Option, + /// How far in the past queries may evaluate; CircularBuffer keeps this much + /// beyond the planned windows. 0 keeps exactly the planned windows. + pub retention_slack_ms: u64, } impl Default for StoreSettings { fn default() -> Self { Self { lock_strategy: LockStrategy::PerKey, + spill_dir: None, + retention_slack_ms: 0, } } } @@ -469,6 +476,38 @@ output_dir: "./output" ); } + #[test] + fn spill_dir_is_off_unless_overridden() { + let config: EngineConfig = Figment::new() + .merge(Yaml::string(MINIMAL_YAML)) + .extract() + .expect("config should deserialize"); + assert_eq!(config.store.spill_dir, None); + + let config: EngineConfig = Figment::new() + .merge(Yaml::string(MINIMAL_YAML)) + .merge(("store.spill_dir", "/data/spill")) + .extract() + .expect("config should deserialize"); + assert_eq!(config.store.spill_dir.as_deref(), Some("/data/spill")); + } + + #[test] + fn retention_slack_defaults_to_zero() { + let config: EngineConfig = Figment::new() + .merge(Yaml::string(MINIMAL_YAML)) + .extract() + .expect("config should deserialize"); + assert_eq!(config.store.retention_slack_ms, 0); + + let config: EngineConfig = Figment::new() + .merge(Yaml::string(MINIMAL_YAML)) + .merge(("store.retention_slack_ms", 11_000)) + .extract() + .expect("config should deserialize"); + assert_eq!(config.store.retention_slack_ms, 11_000); + } + #[test] fn check_config_rejects_zero_scrape_interval() { let yaml = r#" diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index cda96f63..8a903dbf 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -20,6 +20,7 @@ use query_engine_rust::precompute_engine::json_ingest::{ JsonFileIngestConfig, JsonFileIngestSource, TimestampUnit, }; use query_engine_rust::precompute_engine::PrecomputeWorkerDiagnostics; +use query_engine_rust::stores::simple_map_store::spill::SpillWriter; use query_engine_rust::utils::file_io::{read_inference_config, read_streaming_config}; use query_engine_rust::InferenceConfig; use query_engine_rust::{ @@ -106,11 +107,17 @@ async fn main() -> Result<()> { let cleanup_policy = inference_config.cleanup_policy; info!("Using cleanup policy: {:?}", cleanup_policy); - let store = Arc::new(SimpleMapStore::new_with_strategy( + let mut store = SimpleMapStore::new_with_strategy( streaming_config.clone(), cleanup_policy, config.store.lock_strategy, - )); + ) + .with_retention_slack_ms(config.store.retention_slack_ms); + if let Some(dir) = &config.store.spill_dir { + info!("Spilling evicted windows to {}", dir); + store = store.with_spill(Arc::new(SpillWriter::new(dir)?)); + } + let store = Arc::new(store); let engine = Arc::new(SimpleEngine::new( store.clone(), diff --git a/asap-query-engine/src/stores/simple_map_store/common.rs b/asap-query-engine/src/stores/simple_map_store/common.rs index bab287c5..df2f6598 100644 --- a/asap-query-engine/src/stores/simple_map_store/common.rs +++ b/asap-query-engine/src/stores/simple_map_store/common.rs @@ -7,7 +7,33 @@ use tracing::debug; pub type MetricID = u32; pub type EpochID = u64; + +/// Windows a CircularBuffer store keeps: the planned count plus enough windows to +/// cover `retention_slack_ms` of query delay, e.g. a query evaluated 10s in the past +/// needs 10 more 1s windows. `None` (no planned count) stays unlimited. A slide of 0 +/// means tumbling, so windows step by `window_size_ms`; a zero-width window has no +/// time step, so it gets no slack. +pub fn retained_windows( + num_aggregates_to_retain: Option, + window_size_ms: u64, + slide_interval_ms: u64, + retention_slack_ms: u64, +) -> Option { + let step_ms = if slide_interval_ms == 0 { + window_size_ms + } else { + slide_interval_ms + }; + let slack_windows = match step_ms { + 0 => 0, + step_ms => retention_slack_ms.div_ceil(step_ms), + }; + num_aggregates_to_retain.map(|n| n + slack_windows) +} + pub type MetricBucketMap = HashMap)>>; +/// An entry removed from an epoch by cleanup. +pub type EvictedEntry = (TimestampRange, MetricID, Arc); /// Sorts one key's buckets into chronological (ascending start) order. /// @@ -142,6 +168,14 @@ impl MutableEpoch { } } + /// The `n` oldest distinct windows, oldest first. + pub fn oldest_windows(&self, n: usize) -> Vec { + let mut windows: Vec = self.windows_set.iter().copied().collect(); + windows.sort_unstable(); + windows.truncate(n); + windows + } + pub fn window_count(&self) -> usize { self.windows_set.len() } @@ -195,6 +229,8 @@ impl MutableEpoch { /// Consume this epoch and produce an immutable SealedEpoch by sorting in-place. /// Zips the three columns into tuples and sorts — moves Arcs without cloning. /// O(M log M) paid once at rotation time, not at query time. + /// Unused: no cleanup policy seals epochs. + #[allow(dead_code)] pub fn seal(self) -> SealedEpoch { let min_start = self.min_start; let max_end = self.max_end; @@ -286,17 +322,20 @@ impl MutableEpoch { ) } - /// Remove specific windows (ReadBased cleanup). + /// Remove specific windows (ReadBased cleanup), returning the removed entries. /// Drains all three columns in lockstep — moves Arcs without cloning. - pub fn remove_windows(&mut self, windows: &[TimestampRange]) { + pub fn remove_windows(&mut self, windows: &[TimestampRange]) -> Vec { let window_set: HashSet = windows.iter().copied().collect(); let old_windows = std::mem::take(&mut self.windows_col); let old_metrics = std::mem::take(&mut self.metric_ids_col); let old_aggs = std::mem::take(&mut self.aggregates_col); + let mut removed = Vec::new(); for ((tr, mid), agg) in old_windows.into_iter().zip(old_metrics).zip(old_aggs) { - if !window_set.contains(&tr) { + if window_set.contains(&tr) { + removed.push((tr, mid, agg)); + } else { self.windows_col.push(tr); self.metric_ids_col.push(mid); self.aggregates_col.push(agg); @@ -314,6 +353,7 @@ impl MutableEpoch { // Recompute bounds (cleanup is rare; linear scan is fine). self.min_start = self.windows_col.iter().map(|tr| tr.0).min(); self.max_end = self.windows_col.iter().map(|tr| tr.1).max(); + removed } } @@ -393,11 +433,14 @@ impl SealedEpoch { } } - /// Remove specific windows (ReadBased / CircularBuffer cleanup). Rebuilds Vec in one pass. - /// Also updates `distinct_window_count`. - pub fn remove_windows(&mut self, windows: &[TimestampRange]) { + /// Remove specific windows (ReadBased / CircularBuffer cleanup), returning the removed + /// entries. Rebuilds Vec in one pass. Also updates `distinct_window_count`. + pub fn remove_windows(&mut self, windows: &[TimestampRange]) -> Vec { let window_set: HashSet = windows.iter().copied().collect(); - self.entries.retain(|(tr, _, _)| !window_set.contains(tr)); + let (removed, kept) = std::mem::take(&mut self.entries) + .into_iter() + .partition(|(tr, _, _)| window_set.contains(tr)); + self.entries = kept; self.min_start = self.entries.iter().map(|(tr, _, _)| tr.0).min(); self.max_end = self.entries.iter().map(|(tr, _, _)| tr.1).max(); // Recount distinct windows (entries remain sorted; dedup in one pass). @@ -410,10 +453,12 @@ impl SealedEpoch { } } self.distinct_window_count = count; + removed } /// Deduplicated windows (entries sorted, so consecutive dupes are adjacent). /// Used to purge `read_counts` when this epoch is dropped. + #[allow(dead_code)] pub fn unique_windows(&self) -> Vec { let mut windows: Vec = self.entries.iter().map(|(tr, _, _)| *tr).collect(); windows.dedup(); @@ -467,3 +512,30 @@ pub fn resolve_exact_windows( (results, matched_windows, total_entries) } + +#[cfg(test)] +mod tests { + use super::retained_windows; + + #[test] + fn retention_slack_adds_whole_windows_rounded_up() { + // 11s of slack: 11 more 1s windows, but one more 60s window. + assert_eq!(retained_windows(Some(1), 1_000, 1_000, 11_000), Some(12)); + assert_eq!(retained_windows(Some(15), 60_000, 60_000, 11_000), Some(16)); + assert_eq!(retained_windows(Some(15), 60_000, 60_000, 0), Some(15)); + assert_eq!(retained_windows(None, 1_000, 1_000, 11_000), None); + } + + #[test] + fn zero_slide_steps_by_window_size() { + // Tumbling configs may leave slideIntervalMs at 0; dividing by it panicked. + assert_eq!(retained_windows(Some(1), 1_000, 0, 11_000), Some(12)); + } + + #[test] + fn zero_width_window_gets_no_slack() { + // Window and slide both 0 divided by zero, even with the default slack of 0. + assert_eq!(retained_windows(Some(3), 0, 0, 0), Some(3)); + assert_eq!(retained_windows(Some(3), 0, 0, 11_000), Some(3)); + } +} diff --git a/asap-query-engine/src/stores/simple_map_store/global.rs b/asap-query-engine/src/stores/simple_map_store/global.rs index 6958ed1f..ae841b64 100644 --- a/asap-query-engine/src/stores/simple_map_store/global.rs +++ b/asap-query-engine/src/stores/simple_map_store/global.rs @@ -2,9 +2,10 @@ use crate::data_model::{ AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, }; use crate::stores::simple_map_store::common::{ - resolve_exact_windows, sort_buckets_chronologically, EpochID, InternTable, MetricBucketMap, - MutableEpoch, SealedEpoch, TimestampRange, + resolve_exact_windows, retained_windows, sort_buckets_chronologically, EpochID, EvictedEntry, + InternTable, MetricBucketMap, MutableEpoch, SealedEpoch, TimestampRange, }; +use crate::stores::simple_map_store::spill::{resolve_evicted, SpillEntry, SpillWriter}; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::{Arc, Mutex, RwLock}; @@ -24,15 +25,9 @@ struct PerKeyState { /// Sealed (immutable) epochs stored as flat sorted Vecs (Optimization 2). sealed_epochs: BTreeMap, - /// Monotonically increasing ID of the current epoch. - current_epoch_id: EpochID, - - /// Max distinct time-windows per epoch before sealing. - /// None = unlimited (set on first insert from num_aggregates_to_retain). - epoch_capacity: Option, - - /// Max total epochs (1 current + sealed) to retain. - max_epochs: usize, + /// Distinct windows CircularBuffer keeps. None = unlimited (set on first + /// insert from the planned retention plus slack). + window_limit: Option, } impl PerKeyState { @@ -41,74 +36,34 @@ impl PerKeyState { intern: InternTable::new(), current_epoch: MutableEpoch::new(), sealed_epochs: BTreeMap::new(), - current_epoch_id: 0, - epoch_capacity: None, - max_epochs: 4, + window_limit: None, } } - /// Set epoch_capacity on first insert (no-op after first call). - fn configure_epochs(&mut self, num_aggregates_to_retain: Option) { - if self.epoch_capacity.is_none() { - if let Some(cap) = num_aggregates_to_retain { - self.epoch_capacity = Some(cap as usize); + /// Set window_limit on first insert (no-op after first call). + fn configure_window_limit(&mut self, retained_windows: Option) { + if self.window_limit.is_none() { + if let Some(limit) = retained_windows { + self.window_limit = Some(limit as usize); } } } - /// Seal the current epoch when full, then evict the minimum number of oldest windows - /// to keep total distinct windows ≤ `epoch_capacity * max_epochs`. - /// Returns the evicted windows so the caller can clean up `read_counts`. - fn maybe_rotate_epoch(&mut self) -> Vec { - let capacity = match self.epoch_capacity { - Some(c) if c > 0 => c, + /// Evict the oldest distinct windows beyond `window_limit`, returning the evicted + /// entries. CircularBuffer never seals, so every window stays in the current epoch + /// and its distinct-window count is exact even when a window's series arrive in + /// batches interleaved with other windows' (one batch per precompute worker). + fn evict_beyond_window_limit(&mut self) -> Vec { + let limit = match self.window_limit { + Some(l) if l > 0 => l, _ => return Vec::new(), // unlimited }; - let retention_limit = capacity * self.max_epochs; - - // Step 1: seal current epoch if it has hit the window capacity threshold. - if self.current_epoch.window_count() >= capacity { - let hint = self.current_epoch.len(); - let old = std::mem::replace(&mut self.current_epoch, MutableEpoch::with_capacity(hint)); - self.sealed_epochs.insert(self.current_epoch_id, old.seal()); - self.current_epoch_id += 1; - } - - // Step 2: evict oldest windows until total distinct windows ≤ retention_limit. - let total: usize = self.current_epoch.window_count() - + self - .sealed_epochs - .values() - .map(|e| e.distinct_window_count()) - .sum::(); - - if total <= retention_limit { + let over = self.current_epoch.window_count().saturating_sub(limit); + if over == 0 { return Vec::new(); } - let mut over = total - retention_limit; - let mut evicted = Vec::new(); - - while over > 0 { - let oldest_id = match self.sealed_epochs.keys().next().copied() { - Some(id) => id, - None => break, - }; - let oldest_windows = self.sealed_epochs[&oldest_id].unique_windows(); - let n_evict = over.min(oldest_windows.len()); - let to_remove = oldest_windows[..n_evict].to_vec(); - over -= n_evict; - evicted.extend_from_slice(&to_remove); - - if n_evict == oldest_windows.len() { - self.sealed_epochs.remove(&oldest_id); - } else { - self.sealed_epochs - .get_mut(&oldest_id) - .unwrap() - .remove_windows(&to_remove); - } - } - evicted + let to_remove = self.current_epoch.oldest_windows(over); + self.current_epoch.remove_windows(&to_remove) } } @@ -142,6 +97,12 @@ pub struct SimpleMapStoreGlobal { // Policy for cleaning up old aggregates cleanup_policy: CleanupPolicy, + + // Where evicted windows are archived; None drops them. + pub(super) spill: Option>, + + // Query delay CircularBuffer retention covers beyond the planned windows. + pub(super) retention_slack_ms: u64, } impl SimpleMapStoreGlobal { @@ -156,6 +117,8 @@ impl SimpleMapStoreGlobal { }), streaming_config: RwLock::new(streaming_config), cleanup_policy, + spill: None, + retention_slack_ms: 0, } } @@ -226,7 +189,7 @@ type GroupedBatch = HashMap< struct BatchConfig { metric: String, is_delta: bool, - num_aggregates_to_retain: Option, + retained_windows: Option, read_count_threshold: Option, } @@ -275,7 +238,12 @@ impl Store for SimpleMapStoreGlobal { metric: aggregation_config.metric.clone(), is_delta: aggregation_config.aggregation_type == AggregationType::DeltaSetAggregator, - num_aggregates_to_retain: aggregation_config.num_aggregates_to_retain, + retained_windows: retained_windows( + aggregation_config.num_aggregates_to_retain, + aggregation_config.window_size_ms, + aggregation_config.slide_interval_ms, + self.retention_slack_ms, + ), read_count_threshold: aggregation_config.read_count_threshold, }, u64::MAX, @@ -292,6 +260,8 @@ impl Store for SimpleMapStoreGlobal { let lock_wait_start = Instant::now(); let mut data = self.lock.lock().unwrap(); + // Evicted entries to spill once the lock is released, keys already resolved. + let mut to_spill: Vec<(StoreKey, Vec)> = Vec::new(); #[cfg(feature = "lock_profiling")] { @@ -319,8 +289,8 @@ impl Store for SimpleMapStoreGlobal { let batch_len = items.len() as u64; - // Ensure PerKeyState exists and configure epoch capacity once per group (Opt 4). - // configure_epochs is a no-op after the first call, so calling it once here + // Ensure PerKeyState exists and configure the window limit once per group (Opt 4). + // configure_window_limit is a no-op after the first call, so calling it once here // avoids the is_none() check on every inner iteration. { let per_key = data @@ -328,10 +298,11 @@ impl Store for SimpleMapStoreGlobal { .entry(store_key) .or_insert_with(PerKeyState::new); if !cfg.is_delta { - per_key.configure_epochs(cfg.num_aggregates_to_retain); + per_key.configure_window_limit(cfg.retained_windows); } } // per_key borrow ends here + let mut evicted = Vec::new(); for (output, precompute) in items { // Get per_key fresh each iteration so the borrow of data.stores ends before // the cleanup branches borrow data.read_counts (different field — NLL splits @@ -347,63 +318,70 @@ impl Store for SimpleMapStoreGlobal { .current_epoch .insert(metric_id, timestamp_range, Arc::from(precompute)); - // Apply retention policy if configured (but exclude DeltaSetAggregator). + // Apply ReadBased retention if configured (but exclude DeltaSetAggregator). // per_key is last used above; NLL ends its borrow so data.read_counts can - // be accessed in the cleanup branches below. - if !cfg.is_delta { - match self.cleanup_policy { - CleanupPolicy::CircularBuffer => { - let dropped_windows = data - .stores - .get_mut(&store_key) - .unwrap() - .maybe_rotate_epoch(); - if !dropped_windows.is_empty() { - if let Some(rc_map) = data.read_counts.get_mut(&store_key) { - for window in &dropped_windows { - rc_map.remove(window); - } - } - for window in &dropped_windows { - debug!( - "Removed old aggregate for {} aggregation_id {} window {}-{} (epoch rotation)", - cfg.metric, store_key, window.0, window.1 - ); - } - } - } - CleanupPolicy::ReadBased => { - if let Some(threshold) = cfg.read_count_threshold { - let rc_map = data.read_counts.entry(store_key).or_default(); - let windows_to_remove: Vec = rc_map - .iter() - .filter(|(_, &count)| count >= threshold) - .map(|(range, _)| *range) - .collect(); - - if !windows_to_remove.is_empty() { - for window in &windows_to_remove { - debug!( - "Removed aggregate for {} aggregation_id {} window {}-{} (read_count >= threshold: {})", - cfg.metric, store_key, window.0, window.1, threshold - ); - rc_map.remove(window); - } - - let per_key = data.stores.get_mut(&store_key).unwrap(); - per_key.current_epoch.remove_windows(&windows_to_remove); - per_key.sealed_epochs.retain(|_, epoch| { - epoch.remove_windows(&windows_to_remove); - !epoch.is_empty() - }); - } + // be accessed below. + if !cfg.is_delta && matches!(self.cleanup_policy, CleanupPolicy::ReadBased) { + if let Some(threshold) = cfg.read_count_threshold { + let rc_map = data.read_counts.entry(store_key).or_default(); + let windows_to_remove: Vec = rc_map + .iter() + .filter(|(_, &count)| count >= threshold) + .map(|(range, _)| *range) + .collect(); + + if !windows_to_remove.is_empty() { + for window in &windows_to_remove { + debug!( + "Removed aggregate for {} aggregation_id {} window {}-{} (read_count >= threshold: {})", + cfg.metric, store_key, window.0, window.1, threshold + ); + rc_map.remove(window); } + + let per_key = data.stores.get_mut(&store_key).unwrap(); + evicted + .extend(per_key.current_epoch.remove_windows(&windows_to_remove)); + per_key.sealed_epochs.retain(|_, epoch| { + evicted.extend(epoch.remove_windows(&windows_to_remove)); + !epoch.is_empty() + }); } - CleanupPolicy::NoCleanup => {} } } } + if !cfg.is_delta && matches!(self.cleanup_policy, CleanupPolicy::CircularBuffer) { + let dropped = data + .stores + .get_mut(&store_key) + .unwrap() + .evict_beyond_window_limit(); + let mut dropped_windows: Vec = + dropped.iter().map(|(window, _, _)| *window).collect(); + // Entries come in insertion order, so a window's entries aren't adjacent. + dropped_windows.sort_unstable(); + dropped_windows.dedup(); + if let Some(rc_map) = data.read_counts.get_mut(&store_key) { + for window in &dropped_windows { + rc_map.remove(window); + } + } + for window in &dropped_windows { + debug!( + "Removed old aggregate for {} aggregation_id {} window {}-{} (circular buffer)", + cfg.metric, store_key, window.0, window.1 + ); + } + evicted.extend(dropped); + } + if self.spill.is_some() && !evicted.is_empty() { + to_spill.push(( + store_key, + resolve_evicted(&data.stores[&store_key].intern, evicted), + )); + } + // Opt 4: one count update per group (was one per item) let current_count = data.items_inserted.entry(cfg.metric.clone()).or_insert(0); let old_count = *current_count; @@ -422,6 +400,12 @@ impl Store for SimpleMapStoreGlobal { batch_size ); } + drop(data); + if let Some(spill) = &self.spill { + for (aggregation_id, entries) in &to_spill { + spill.spill(*aggregation_id, entries); + } + } let batch_insert_duration = batch_insert_start_time.elapsed(); debug!( diff --git a/asap-query-engine/src/stores/simple_map_store/mod.rs b/asap-query-engine/src/stores/simple_map_store/mod.rs index 8f0ad012..2e9d06f7 100644 --- a/asap-query-engine/src/stores/simple_map_store/mod.rs +++ b/asap-query-engine/src/stores/simple_map_store/mod.rs @@ -1,6 +1,7 @@ mod common; pub mod global; pub mod per_key; +pub mod spill; use crate::data_model::{ AggregateCore, CleanupPolicy, LockStrategy, PrecomputedOutput, StreamingConfig, @@ -8,6 +9,7 @@ use crate::data_model::{ use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use global::SimpleMapStoreGlobal; use per_key::SimpleMapStorePerKey; +use spill::SpillWriter; use std::collections::HashMap; use std::sync::Arc; @@ -64,6 +66,25 @@ impl SimpleMapStore { } } + /// Write every window cleanup evicts to `spill` before it leaves memory. + pub fn with_spill(mut self, spill: Arc) -> Self { + match &mut self { + SimpleMapStore::Global(store) => store.spill = Some(spill), + SimpleMapStore::PerKey(store) => store.spill = Some(spill), + } + self + } + + /// Keep `retention_slack_ms` of extra windows under CircularBuffer, so queries + /// evaluated that far in the past still find their windows. + pub fn with_retention_slack_ms(mut self, retention_slack_ms: u64) -> Self { + match &mut self { + SimpleMapStore::Global(store) => store.retention_slack_ms = retention_slack_ms, + SimpleMapStore::PerKey(store) => store.retention_slack_ms = retention_slack_ms, + } + self + } + /// Replace the streaming config at runtime. Delegates to the active variant. pub fn update_streaming_config(&self, new_config: StreamingConfig) { match self { diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index 928f1f2c..450e4ca9 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -2,9 +2,10 @@ use crate::data_model::{ AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, }; use crate::stores::simple_map_store::common::{ - resolve_exact_windows, sort_buckets_chronologically, EpochID, InternTable, MetricBucketMap, - MetricID, MutableEpoch, SealedEpoch, TimestampRange, + resolve_exact_windows, retained_windows, sort_buckets_chronologically, EpochID, EvictedEntry, + InternTable, MetricBucketMap, MetricID, MutableEpoch, SealedEpoch, TimestampRange, }; +use crate::stores::simple_map_store::spill::{resolve_evicted, SpillWriter}; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use dashmap::DashMap; use std::collections::{BTreeMap, HashMap}; @@ -26,15 +27,9 @@ struct StoreKeyData { /// Sealed (immutable) epochs stored as flat sorted Vecs (Optimization 2). sealed_epochs: BTreeMap, - /// Monotonically increasing ID of the current epoch. - current_epoch_id: EpochID, - - /// Max distinct time-windows per epoch before sealing. - /// None = unlimited (set on first insert from num_aggregates_to_retain). - epoch_capacity: Option, - - /// Max total epochs (1 current + sealed) to retain before dropping the oldest. - max_epochs: usize, + /// Distinct windows CircularBuffer keeps. None = unlimited (set on first + /// insert from the planned retention plus slack). + window_limit: Option, /// Track how many times each timestamp range has been read. /// Behind Mutex so range queries can use a read lock on the outer RwLock. @@ -47,85 +42,49 @@ impl StoreKeyData { intern: InternTable::new(), current_epoch: MutableEpoch::new(), sealed_epochs: BTreeMap::new(), - current_epoch_id: 0, - epoch_capacity: None, - max_epochs: 4, + window_limit: None, read_counts: Mutex::new(HashMap::new()), } } - /// Set epoch_capacity on first insert (no-op after first call). - fn configure_epochs(&mut self, num_aggregates_to_retain: Option) { - if self.epoch_capacity.is_none() { - if let Some(cap) = num_aggregates_to_retain { - self.epoch_capacity = Some(cap as usize); + /// Set window_limit on first insert (no-op after first call). + fn configure_window_limit(&mut self, retained_windows: Option) { + if self.window_limit.is_none() { + if let Some(limit) = retained_windows { + self.window_limit = Some(limit as usize); } } } - /// Seal the current epoch when full, then evict the minimum number of oldest windows - /// to keep total distinct windows ≤ `epoch_capacity * max_epochs`. - /// - /// Matches legacy per-window eviction semantics: only the exact number of windows - /// needed to reach the retention limit are removed, which may be fewer than a full epoch. - fn maybe_rotate_epoch(&mut self) { - let capacity = match self.epoch_capacity { - Some(c) if c > 0 => c, - _ => return, // unlimited + /// Evict the oldest distinct windows beyond `window_limit`, returning the evicted + /// entries. CircularBuffer never seals, so every window stays in the current epoch + /// and its distinct-window count is exact even when a window's series arrive in + /// batches interleaved with other windows' (one batch per precompute worker). + fn evict_beyond_window_limit(&mut self) -> Vec { + let limit = match self.window_limit { + Some(l) if l > 0 => l, + _ => return Vec::new(), // unlimited }; - let retention_limit = capacity * self.max_epochs; - - // Step 1: seal current epoch if it has hit the window capacity threshold. - if self.current_epoch.window_count() >= capacity { - let hint = self.current_epoch.len(); - let old = std::mem::replace(&mut self.current_epoch, MutableEpoch::with_capacity(hint)); - self.sealed_epochs.insert(self.current_epoch_id, old.seal()); - self.current_epoch_id += 1; - } - - // Step 2: evict oldest windows until total distinct windows ≤ retention_limit. - // Uses O(E) distinct_window_count() calls (E ≤ max_epochs, a small constant). - let total: usize = self.current_epoch.window_count() - + self - .sealed_epochs - .values() - .map(|e| e.distinct_window_count()) - .sum::(); - - if total <= retention_limit { - return; - } - let mut over = total - retention_limit; - - while over > 0 { - let oldest_id = match self.sealed_epochs.keys().next().copied() { - Some(id) => id, - None => break, - }; - let oldest_windows = self.sealed_epochs[&oldest_id].unique_windows(); - let n_evict = over.min(oldest_windows.len()); - let to_remove = oldest_windows[..n_evict].to_vec(); - over -= n_evict; - - { - let read_counts = self.read_counts.get_mut().unwrap(); - for w in &to_remove { - read_counts.remove(w); - } - } - if n_evict == oldest_windows.len() { - self.sealed_epochs.remove(&oldest_id); - } else { - self.sealed_epochs - .get_mut(&oldest_id) - .unwrap() - .remove_windows(&to_remove); - } + let over = self.current_epoch.window_count().saturating_sub(limit); + if over == 0 { + return Vec::new(); } + let to_remove = self.current_epoch.oldest_windows(over); + let read_counts = self.read_counts.get_mut().unwrap(); + for w in &to_remove { + read_counts.remove(w); + } + self.current_epoch.remove_windows(&to_remove) } - /// Apply ReadBased cleanup across current and sealed epochs. - fn cleanup_read_based(&mut self, metric: &str, aggregation_id: u64, threshold: u64) { + /// Apply ReadBased cleanup across current and sealed epochs, returning the evicted + /// entries. + fn cleanup_read_based( + &mut self, + metric: &str, + aggregation_id: u64, + threshold: u64, + ) -> Vec { let read_counts = self.read_counts.get_mut().unwrap(); let windows_to_remove: Vec = read_counts @@ -135,7 +94,7 @@ impl StoreKeyData { .collect(); if windows_to_remove.is_empty() { - return; + return Vec::new(); } for window in &windows_to_remove { @@ -147,13 +106,14 @@ impl StoreKeyData { } // Remove from current epoch. - self.current_epoch.remove_windows(&windows_to_remove); + let mut evicted = self.current_epoch.remove_windows(&windows_to_remove); // Remove from sealed epochs; drop any that become empty. self.sealed_epochs.retain(|_, epoch| { - epoch.remove_windows(&windows_to_remove); + evicted.extend(epoch.remove_windows(&windows_to_remove)); !epoch.is_empty() }); + evicted } } @@ -176,6 +136,12 @@ pub struct SimpleMapStorePerKey { // Policy for cleaning up old aggregates cleanup_policy: CleanupPolicy, + + // Where evicted windows are archived; None drops them. + pub(super) spill: Option>, + + // Query delay CircularBuffer retention covers beyond the planned windows. + pub(super) retention_slack_ms: u64, } impl SimpleMapStorePerKey { @@ -187,6 +153,8 @@ impl SimpleMapStorePerKey { items_inserted: DashMap::new(), streaming_config: RwLock::new(streaming_config), cleanup_policy, + spill: None, + retention_slack_ms: 0, } } @@ -251,21 +219,22 @@ impl SimpleMapStorePerKey { aggregation_id: u64, num_aggregates_to_retain: Option, read_count_threshold: Option, - ) { + ) -> Vec { match self.cleanup_policy { CleanupPolicy::CircularBuffer => { - // configure_epochs was already called before insert; - // rotation is handled by maybe_rotate_epoch after each insert batch. + // configure_window_limit was already called before insert; + // eviction is handled by evict_beyond_window_limit after each insert. // Nothing additional needed here. let _ = (num_aggregates_to_retain, metric, aggregation_id); + Vec::new() } - CleanupPolicy::ReadBased => { - if let Some(threshold) = read_count_threshold { - data.cleanup_read_based(metric, aggregation_id, threshold); - } - } + CleanupPolicy::ReadBased => match read_count_threshold { + Some(threshold) => data.cleanup_read_based(metric, aggregation_id, threshold), + None => Vec::new(), + }, CleanupPolicy::NoCleanup => { // Do nothing - no cleanup + Vec::new() } } } @@ -365,9 +334,14 @@ impl SimpleMapStorePerKey { .get_aggregation_config(aggregation_id) .ok_or_else(|| format!("Aggregation config not found for {}", aggregation_id))?; - // Configure epoch capacity on first insert (Optimization 2) + // Configure the window limit on first insert if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { - data.configure_epochs(aggregation_config.num_aggregates_to_retain); + data.configure_window_limit(retained_windows( + aggregation_config.num_aggregates_to_retain, + aggregation_config.window_size_ms, + aggregation_config.slide_interval_ms, + self.retention_slack_ms, + )); } for (output, precompute) in items { @@ -378,24 +352,33 @@ impl SimpleMapStorePerKey { // Insert into current (mutable) epoch. data.current_epoch .insert(metric_id, timestamp_range, Arc::from(precompute)); + } - // After each item, check if we should rotate (CircularBuffer, Optimization 2) - if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator - && matches!(self.cleanup_policy, CleanupPolicy::CircularBuffer) - { - data.maybe_rotate_epoch(); - } + let mut evicted = Vec::new(); + if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator + && matches!(self.cleanup_policy, CleanupPolicy::CircularBuffer) + { + evicted.extend(data.evict_beyond_window_limit()); } // Apply retention policy if configured (but exclude DeltaSetAggregator) if aggregation_config.aggregation_type != AggregationType::DeltaSetAggregator { - self.cleanup_old_aggregates( + evicted.extend(self.cleanup_old_aggregates( &mut data, metric, aggregation_id, aggregation_config.num_aggregates_to_retain, aggregation_config.read_count_threshold, - ); + )); + } + // Keys are resolved under the lock; serializing and writing happen after it. + let to_spill = self + .spill + .as_ref() + .map(|spill| (spill, resolve_evicted(&data.intern, evicted))); + drop(data); + if let Some((spill, entries)) = to_spill { + spill.spill(aggregation_id, &entries); } #[cfg(feature = "lock_profiling")] diff --git a/asap-query-engine/src/stores/simple_map_store/spill.rs b/asap-query-engine/src/stores/simple_map_store/spill.rs new file mode 100644 index 00000000..df70903a --- /dev/null +++ b/asap-query-engine/src/stores/simple_map_store/spill.rs @@ -0,0 +1,313 @@ +//! Write-only archive of windows evicted from the in-memory store. +//! +//! One append-only file per aggregation, `agg_.bin`. Each record is, in +//! little-endian u64s: window start ms, window end ms, key length, key bytes +//! (bincode label values; length `NO_KEY` and no bytes for an unkeyed +//! window), accumulator length, accumulator bytes (`serialize_to_bytes`). + +use super::common::{EvictedEntry, InternTable, TimestampRange}; +use crate::data_model::{AggregateCore, KeyByLabelValues}; +use dashmap::DashMap; +use lazy_static::lazy_static; +use prometheus::{register_counter, Counter}; +use std::fs::{self, File, OpenOptions}; +use std::io::Write; +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use tracing::error; + +/// Key length recorded for a window without a label key. +pub const NO_KEY: u64 = u64::MAX; + +lazy_static! { + pub static ref SPILL_WRITE_ERRORS_TOTAL: Counter = register_counter!( + "asap_store_spill_write_errors_total", + "Evicted-window batches that could not be written to the spill directory" + ) + .unwrap(); +} + +/// An evicted entry with its label key resolved, so it can be spilled after +/// the store lock (which guards the intern table) is released. +pub type SpillEntry = ( + TimestampRange, + Option, + Arc, +); + +/// Resolve each evicted entry's interned key. +pub fn resolve_evicted(intern: &InternTable, evicted: Vec) -> Vec { + evicted + .into_iter() + .map(|(window, metric_id, aggregate)| { + (window, intern.resolve(metric_id).clone(), aggregate) + }) + .collect() +} + +pub struct SpillWriter { + dir: PathBuf, + // One lock per file, so writes for different aggregations don't wait on each other. + files: DashMap>>, +} + +impl SpillWriter { + pub fn new(dir: impl Into) -> std::io::Result { + let dir = dir.into(); + fs::create_dir_all(&dir)?; + // Register now so /metrics reports 0 rather than omitting the counter. + lazy_static::initialize(&SPILL_WRITE_ERRORS_TOTAL); + Ok(Self { + dir, + files: DashMap::new(), + }) + } + + /// Appends `entries` to the aggregation's file. A failure is logged and + /// counted, not returned: the windows leave memory either way. + pub fn spill(&self, aggregation_id: u64, entries: &[SpillEntry]) { + if entries.is_empty() { + return; + } + if let Err(e) = self.append(aggregation_id, entries) { + SPILL_WRITE_ERRORS_TOTAL.inc(); + error!( + "spilling {} evicted entries of aggregation {aggregation_id} failed: {e}", + entries.len() + ); + } + } + + fn append( + &self, + aggregation_id: u64, + entries: &[SpillEntry], + ) -> Result<(), Box> { + let mut buf = Vec::new(); + for (window, key, aggregate) in entries { + buf.extend_from_slice(&window.0.to_le_bytes()); + buf.extend_from_slice(&window.1.to_le_bytes()); + match key { + Some(key) => put_len_prefixed(&mut buf, &key.serialize_to_bytes()?), + None => buf.extend_from_slice(&NO_KEY.to_le_bytes()), + } + put_len_prefixed(&mut buf, &aggregate.serialize_to_bytes()?); + } + let file = self + .files + .entry(aggregation_id) + .or_try_insert_with(|| { + OpenOptions::new() + .create(true) + .append(true) + .open(self.dir.join(format!("agg_{aggregation_id}.bin"))) + .map(|file| Arc::new(Mutex::new(file))) + })? + .clone(); + // One write per batch, so a killed process loses at most the batch in flight. + file.lock().unwrap().write_all(&buf)?; + Ok(()) + } +} + +fn put_len_prefixed(buf: &mut Vec, bytes: &[u8]) { + buf.extend_from_slice(&(bytes.len() as u64).to_le_bytes()); + buf.extend_from_slice(bytes); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data_model::{ + AggregationType, CleanupPolicy, KeyByLabelValues, LockStrategy, PrecomputedOutput, + SerializableToSink, StreamingConfig, WindowType, + }; + use crate::precompute_operators::SumAccumulator; + use crate::stores::Store; + use crate::{AggregationConfig, SimpleMapStore}; + use promql_utilities::data_model::KeyByLabelNames; + use std::collections::HashMap; + use std::sync::Arc; + + /// (start, end, key, accumulator bytes) per record. + type Record = (u64, u64, Option, Vec); + + fn read_records(path: &std::path::Path) -> Vec { + let bytes = fs::read(path).unwrap(); + let mut pos = 0; + let next_u64 = |pos: &mut usize| { + let value = u64::from_le_bytes(bytes[*pos..*pos + 8].try_into().unwrap()); + *pos += 8; + value + }; + let mut records = Vec::new(); + while pos < bytes.len() { + let start = next_u64(&mut pos); + let end = next_u64(&mut pos); + let key = match next_u64(&mut pos) { + NO_KEY => None, + len => { + let key = &bytes[pos..pos + len as usize]; + pos += len as usize; + Some(KeyByLabelValues::deserialize_from_bytes(key).unwrap()) + } + }; + let len = next_u64(&mut pos) as usize; + records.push((start, end, key, bytes[pos..pos + len].to_vec())); + pos += len; + } + records + } + + fn store( + strategy: LockStrategy, + policy: CleanupPolicy, + retain: Option, + read_count_threshold: Option, + spill_dir: &std::path::Path, + ) -> SimpleMapStore { + let config = AggregationConfig::new( + 1, + AggregationType::Sum, + "".to_string(), + HashMap::new(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + "".to_string(), + 60_000, + 60_000, + WindowType::Tumbling, + "".to_string(), + "cpu_usage".to_string(), + retain, + read_count_threshold, + None, + None, + ); + let streaming = Arc::new(StreamingConfig::new(HashMap::from([(1, config)]))); + SimpleMapStore::new_with_strategy(streaming, policy, strategy) + .with_spill(Arc::new(SpillWriter::new(spill_dir).unwrap())) + } + + fn insert(store: &SimpleMapStore, window: u64, key: Option<&str>) { + let output = PrecomputedOutput::new( + window * 60_000, + (window + 1) * 60_000, + key.map(|k| KeyByLabelValues::new_with_labels(vec![k.to_string()])), + 1, + ); + store + .insert_precomputed_output(output, Box::new(SumAccumulator::with_sum(window as f64))) + .unwrap(); + } + + fn sum_bytes(value: f64) -> Vec { + SumAccumulator::with_sum(value) + .serialize_to_bytes() + .unwrap() + } + + #[test] + fn circular_buffer_spills_each_evicted_window_once() { + for strategy in [LockStrategy::PerKey, LockStrategy::Global] { + let dir = tempfile::tempdir().unwrap(); + let store = store( + strategy, + CleanupPolicy::CircularBuffer, + Some(2), + None, + dir.path(), + ); + for window in 0..5 { + insert(&store, window, Some("host-a")); + } + + let key = Some(KeyByLabelValues::new_with_labels(vec!["host-a".into()])); + let expected: Vec = (0..3) + .map(|w| { + ( + w * 60_000, + (w + 1) * 60_000, + key.clone(), + sum_bytes(w as f64), + ) + }) + .collect(); + assert_eq!( + read_records(&dir.path().join("agg_1.bin")), + expected, + "{strategy:?}" + ); + } + } + + #[test] + fn read_based_spills_windows_read_past_the_threshold() { + for strategy in [LockStrategy::PerKey, LockStrategy::Global] { + let dir = tempfile::tempdir().unwrap(); + let store = store( + strategy, + CleanupPolicy::ReadBased, + None, + Some(1), + dir.path(), + ); + insert(&store, 0, None); + store + .query_precomputed_output("cpu_usage", 1, 0, 60_000) + .unwrap(); + // Cleanup runs on insert. + insert(&store, 1, None); + + assert_eq!( + read_records(&dir.path().join("agg_1.bin")), + vec![(0, 60_000, None, sum_bytes(0.0))], + "{strategy:?}" + ); + } + } + + /// The harness fails a run when the counter is missing, so it must be + /// exported before any write fails. + #[test] + fn error_counter_is_exported_once_spill_is_on() { + let dir = tempfile::tempdir().unwrap(); + SpillWriter::new(dir.path()).unwrap(); + assert!(prometheus::gather() + .iter() + .any(|family| family.get_name() == "asap_store_spill_write_errors_total")); + } + + #[test] + fn failed_spill_is_counted_and_still_evicts() { + for strategy in [LockStrategy::PerKey, LockStrategy::Global] { + let dir = tempfile::tempdir().unwrap(); + let spill_dir = dir.path().join("spill"); + let store = store( + strategy, + CleanupPolicy::CircularBuffer, + Some(1), + None, + &spill_dir, + ); + fs::remove_dir(&spill_dir).unwrap(); + let errors_before = SPILL_WRITE_ERRORS_TOTAL.get(); + + insert(&store, 0, None); + insert(&store, 1, None); + + assert!( + SPILL_WRITE_ERRORS_TOTAL.get() > errors_before, + "{strategy:?}" + ); + let evicted = store + .query_precomputed_output_exact("cpu_usage", 1, 0, 60_000) + .unwrap(); + assert!( + evicted.is_empty(), + "{strategy:?}: window 0 must leave memory" + ); + } + } +} diff --git a/asap-query-engine/src/tests/exact_window_grid_adversarial_tests.rs b/asap-query-engine/src/tests/exact_window_grid_adversarial_tests.rs index b42e3c26..6b4ef30f 100644 --- a/asap-query-engine/src/tests/exact_window_grid_adversarial_tests.rs +++ b/asap-query-engine/src/tests/exact_window_grid_adversarial_tests.rs @@ -20,7 +20,7 @@ //! //! - **Store-level** (`SimpleMapStore` directly): pin the `Store` trait //! contract itself -- empty-range edge cases, boundary containment, -//! overlapping sliding windows, and sealed/mutable epoch spans. These +//! overlapping sliding windows, and spans after CircularBuffer eviction. These //! exercise the store's own (unchanged) `query_precomputed_output`, so they //! document the ground truth the rest of the pipeline must reproduce. //! - **Engine-level** (`SimpleEngine` + PromQL): exercise the actual query @@ -345,19 +345,17 @@ mod tests { ); } - /// Range query spanning both sealed epochs and the still-open (mutable) - /// current epoch, in one call. Mirrors - /// `store_correctness_tests::test_exact_query_correct_after_epoch_rotation` + /// Range query over everything CircularBuffer retained, in one call. Mirrors + /// `store_correctness_tests::test_exact_query_correct_after_eviction` /// but drives the tolerant range scan (`query_precomputed_output`) /// instead of the exact lookup, and checks correctness across the whole /// retained span in a single query rather than probing one window at a /// time. /// - /// capacity=2, 10 inserts => retention_limit=8, evicting windows 0 and 1 - /// (oldest sealed epoch). Windows 2..10 remain, spread across multiple - /// sealed epochs plus the current (newest, still "mutable") epoch. + /// retain=2, 11 inserts => windows 9 and 10 remain, windows 0..9 are + /// evicted. #[test] - fn store_query_range_spans_sealed_and_mutable_epochs() { + fn store_query_range_spans_retained_windows_after_eviction() { let store = make_store( vec![( 1, @@ -372,7 +370,7 @@ mod tests { )], CleanupPolicy::CircularBuffer, ); - let n = 10u64; + let n = 11u64; for i in 0..n { let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); store.insert_precomputed_output(out, acc).unwrap(); @@ -382,15 +380,16 @@ mod tests { .query_precomputed_output("cpu_usage", 1, 0, n * 60_000) .unwrap(); let ts = timestamps_for_none_key(&result); - let expected: Vec<(u64, u64)> = (2..n).map(|i| (i * 60_000, (i + 1) * 60_000)).collect(); + let expected: Vec<(u64, u64)> = + (n - 2..n).map(|i| (i * 60_000, (i + 1) * 60_000)).collect(); assert_eq!( ts, expected, - "range query spanning sealed + mutable epochs must return exactly windows 2..10, \ - with windows 0,1 evicted" + "range query over the retained windows must return exactly windows 9 and 10, \ + with windows 0..9 evicted" ); let buckets = result.get(&None).unwrap(); - for i in 2..n { + for i in n - 2..n { let range = (i * 60_000, (i + 1) * 60_000); let acc = &buckets.iter().find(|(r, _)| *r == range).unwrap().1; let expected_json = SumAccumulator::with_sum(i as f64) @@ -399,7 +398,7 @@ mod tests { assert_eq!( acc.serialize_to_json().unwrap(), expected_json, - "window {i} must return its own value across the sealed/mutable epoch boundary" + "window {i} must return its own value after eviction" ); } } diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index 78f0d4d2..4a0c6545 100644 --- a/asap-query-engine/src/tests/store_correctness_tests.rs +++ b/asap-query-engine/src/tests/store_correctness_tests.rs @@ -192,7 +192,7 @@ pub fn run_contract_suite(strategy: LockStrategy) { test_exact_query_sees_window_inserted_after_a_prior_exact_query(strategy); test_exact_query_miss_then_hit_after_insert(strategy); test_exact_query_correct_across_interleaved_inserts_and_queries(strategy); - test_exact_query_correct_after_epoch_rotation(strategy); + test_exact_query_correct_after_eviction(strategy); // Batched exact-query correctness (#609) test_batch_exact_query_empty_windows_returns_empty(strategy); @@ -211,10 +211,12 @@ pub fn run_contract_suite(strategy: LockStrategy) { // Cleanup policies test_cleanup_circular_buffer_evicts_oldest_window(strategy); test_cleanup_circular_buffer_retains_newest_windows(strategy); + test_retention_slack_keeps_windows_for_delayed_queries(strategy); + test_circular_buffer_keeps_every_series_of_retained_windows(strategy); test_cleanup_read_based_evicts_after_threshold_reads(strategy); test_cleanup_read_based_unread_window_is_retained(strategy); test_delta_set_aggregator_bypasses_cleanup(strategy); - test_buckets_returned_in_chronological_order_after_epoch_rotation(strategy); + test_buckets_returned_in_chronological_order_after_eviction(strategy); // Keyed (label-grouped) entries test_keyed_entries_grouped_by_key(strategy); @@ -508,23 +510,21 @@ fn test_exact_query_correct_across_interleaved_inserts_and_queries(strategy: Loc } } -fn test_exact_query_correct_after_epoch_rotation(strategy: LockStrategy) { - // capacity=2, max_epochs=4 (hardcoded in StoreKeyData/PerKeyState) => - // retention_limit = 8. Inserting 10 windows evicts exactly the oldest 2 - // (windows 0 and 1, both in the oldest sealed epoch), leaving windows - // 2..10 spread across sealed epochs and the current (still-open) epoch. +fn test_exact_query_correct_after_eviction(strategy: LockStrategy) { + // retain=2 => the store keeps the 2 newest windows. After 11 inserts, + // windows 9 and 10 remain and windows 0..9 are evicted. let store = make_store( strategy, CleanupPolicy::CircularBuffer, &[(1, AggregationType::Sum, Some(2), None)], ); - let n = 10u64; + let n = 11u64; for i in 0..n { let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); store.insert_precomputed_output(out, acc).unwrap(); } - for i in 0u64..2 { + for i in 0u64..n - 2 { let evicted = store .query_precomputed_output_exact("cpu_usage", 1, i * 60_000, (i + 1) * 60_000) .unwrap(); @@ -535,14 +535,14 @@ fn test_exact_query_correct_after_epoch_rotation(strategy: LockStrategy) { ); } - for i in 2u64..n { + for i in n - 2..n { let result = store .query_precomputed_output_exact("cpu_usage", 1, i * 60_000, (i + 1) * 60_000) .unwrap(); assert_eq!( total_bucket_count(&result), 1, - "[{}] window {i} must be retrievable via exact query after epoch rotation", + "[{}] window {i} must be retrievable via exact query after eviction", label(strategy) ); let expected = SumAccumulator::with_sum(i as f64) @@ -552,8 +552,8 @@ fn test_exact_query_correct_after_epoch_rotation(strategy: LockStrategy) { assert_eq!( actual, expected, - "[{}] window {i} must return its own value after epoch rotation, \ - not a value from a sealed epoch's stale offset", + "[{}] window {i} must return its own value after eviction, \ + not a value from a stale offset", label(strategy) ); } @@ -886,14 +886,14 @@ fn test_earliest_timestamp_tracked_per_agg_id(strategy: LockStrategy) { // ── cleanup: circular buffer ────────────────────────────────────────────────── fn test_cleanup_circular_buffer_evicts_oldest_window(strategy: LockStrategy) { - // retention_limit = num_aggregates_to_retain * 4 = 2 * 4 = 8. - // Inserting a 9th window triggers eviction of the oldest 1. + // retention_limit = num_aggregates_to_retain = 2. + // Inserting a 3rd window triggers eviction of the oldest 1. let store = make_store( strategy, CleanupPolicy::CircularBuffer, &[(1, AggregationType::Sum, Some(2), None)], ); - for i in 0u64..9 { + for i in 0u64..3 { let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); store.insert_precomputed_output(out, acc).unwrap(); } @@ -917,17 +917,127 @@ fn test_cleanup_circular_buffer_retains_newest_windows(strategy: LockStrategy) { let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); store.insert_precomputed_output(out, acc).unwrap(); } + // Retains exactly num_aggregates_to_retain windows, the count the planner + // sizes memory for. let result = store - .query_precomputed_output("cpu_usage", 1, 60_000, 9 * 60_000) + .query_precomputed_output("cpu_usage", 1, 0, 9 * 60_000) .unwrap(); assert_eq!( - total_bucket_count(&result), - 8, - "[{}] circular buffer must retain the 8 newest windows after eviction", + timestamps_for_none_key(&result), + vec![(7 * 60_000, 8 * 60_000), (8 * 60_000, 9 * 60_000)], + "[{}] circular buffer must retain exactly the 2 newest windows after eviction", label(strategy) ); } +/// An instant query planned with retain=1 but evaluated 10s in the past (the +/// harness's query_time_offset) found its window already evicted and fell back +/// to Prometheus on every request. Slack keeps that window in memory. +fn test_retention_slack_keeps_windows_for_delayed_queries(strategy: LockStrategy) { + let window_ms = 1_000; + let slack_ms = 10_000; + let config = AggregationConfig::new( + 1, + AggregationType::Sum, + "".to_string(), + HashMap::new(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + "".to_string(), + window_ms, + window_ms, + WindowType::Tumbling, + "".to_string(), + "cpu_usage".to_string(), + Some(1), + None, + None, + None, + ); + let streaming = Arc::new(StreamingConfig::new(HashMap::from([(1, config)]))); + let store = + SimpleMapStore::new_with_strategy(streaming, CleanupPolicy::CircularBuffer, strategy) + .with_retention_slack_ms(slack_ms); + let n = 30u64; + for i in 0..n { + let (out, acc) = sum_entry(1, i * window_ms, (i + 1) * window_ms, i as f64); + store.insert_precomputed_output(out, acc).unwrap(); + } + + // 1 planned + 10 slack windows: the window 10s before the newest is kept, + // the one before that is evicted. + let retained = (n - 11..n).map(|i| (i * window_ms, (i + 1) * window_ms)); + let result = store + .query_precomputed_output("cpu_usage", 1, 0, n * window_ms) + .unwrap(); + assert_eq!( + timestamps_for_none_key(&result), + retained.collect::>(), + "[{}] retain=1 plus 10s of 1s windows must keep exactly the newest 11", + label(strategy) + ); +} + +/// Sealing an epoch once it held N windows split a window whose series arrived +/// after the seal across two epochs; the window was counted twice and its early +/// series evicted while it was still retained. Precompute workers each send +/// their own batch per window, so a window's series arrive interleaved with +/// other windows' — modelled here by inserting window w's series b, then +/// window w+1's series a, then window w's series c. +fn test_circular_buffer_keeps_every_series_of_retained_windows(strategy: LockStrategy) { + let series = |host: &str, w: u64| { + keyed_entry( + 1, + w * 60_000, + (w + 1) * 60_000, + key(&[host]), + Box::new(SumAccumulator::with_sum(1.0)), + ) + }; + for retain in [1u64, 3] { + let store = make_store( + strategy, + CleanupPolicy::CircularBuffer, + &[(1, AggregationType::Sum, Some(retain), None)], + ); + let n = 6u64; + let (out, acc) = series("a", 0); + store.insert_precomputed_output(out, acc).unwrap(); + for w in 0..n { + let (out, acc) = series("b", w); + store.insert_precomputed_output(out, acc).unwrap(); + if w + 1 < n { + let (out, acc) = series("a", w + 1); + store.insert_precomputed_output(out, acc).unwrap(); + } + let (out, acc) = series("c", w); + store.insert_precomputed_output(out, acc).unwrap(); + } + + for w in n - retain..n { + let window = store + .query_precomputed_output_exact("cpu_usage", 1, w * 60_000, (w + 1) * 60_000) + .unwrap(); + assert_eq!( + total_bucket_count(&window), + 3, + "[{}] retain={retain}: window {w} must keep all 3 series", + label(strategy) + ); + } + let all = store + .query_precomputed_output("cpu_usage", 1, 0, n * 60_000) + .unwrap(); + assert_eq!( + total_bucket_count(&all), + 3 * retain as usize, + "[{}] retain={retain}: exactly {retain} windows of 3 series remain", + label(strategy) + ); + } +} + // ── cleanup: read-based ─────────────────────────────────────────────────────── fn test_cleanup_read_based_evicts_after_threshold_reads(strategy: LockStrategy) { @@ -1005,7 +1115,7 @@ fn test_cleanup_read_based_unread_window_is_retained(strategy: LockStrategy) { fn test_delta_set_aggregator_bypasses_cleanup(strategy: LockStrategy) { // The store skips cleanup entirely when aggregation_type == "DeltaSetAggregator". - // retention_limit = 2 * 4 = 8. Inserting 10 windows must not evict any. + // retention_limit = 2. Inserting 10 windows must not evict any. let store = make_store( strategy, CleanupPolicy::CircularBuffer, @@ -1045,15 +1155,13 @@ fn test_delta_set_aggregator_bypasses_cleanup(strategy: LockStrategy) { /// moment an order-sensitive accumulator (`DeltaSetAggregator` included, if /// its rotation exemption is ever relaxed) hits this path. /// -/// capacity=2 with 7 inserts forces 3 epoch seals, leaving exactly 1 window -/// in the current epoch (the newest) alongside 3 sealed epochs (the 6 -/// oldest) — the exact shape under which "current checked first" prepends a -/// newer window ahead of older ones. +/// retain=2 with 7 inserts leaves windows 5 and 6; the returned order must +/// still be chronological. /// /// Deliberately does NOT use `timestamps_for_none_key` — that helper sorts /// before returning, which would mask exactly the bug this test exists to /// catch. -fn test_buckets_returned_in_chronological_order_after_epoch_rotation(strategy: LockStrategy) { +fn test_buckets_returned_in_chronological_order_after_eviction(strategy: LockStrategy) { let store = make_store( strategy, CleanupPolicy::CircularBuffer, @@ -1076,8 +1184,8 @@ fn test_buckets_returned_in_chronological_order_after_epoch_rotation(strategy: L .collect(); assert_eq!( returned_order.len(), - n as usize, - "[{}] no windows should have been evicted yet (7 <= retention_limit 8)", + 2, + "[{}] the 2 newest windows should remain (retention_limit 2)", label(strategy) ); @@ -1087,7 +1195,7 @@ fn test_buckets_returned_in_chronological_order_after_epoch_rotation(strategy: L returned_order, chronological, "[{}] buckets must be returned in chronological (ascending start) order \ - even after epoch rotation has occurred", + even after eviction has occurred", label(strategy) ); } diff --git a/asap-tools/experiments/config/config.yaml b/asap-tools/experiments/config/config.yaml index 8539b4ff..97cc3c36 100644 --- a/asap-tools/experiments/config/config.yaml +++ b/asap-tools/experiments/config/config.yaml @@ -85,6 +85,10 @@ backend: # Query engine options query_engine: lock_strategy: "per-key" # Lock strategy for SimpleMapStore: "global" or "per-key" (Rust query engine only) + # Append windows the store evicts to query_engine_output/spill instead of + # dropping them; the run records their size and fails if any write failed. + # Spill files stay on the node; results syncs skip them. + spill: false # Controller configuration controller: @@ -104,11 +108,11 @@ controller: # Aggregate cleanup configuration # Policy options: -# - "circular_buffer": Keep N most recent aggregates (fixed-count cleanup) -# - "read_based": Remove aggregates after N reads (default) +# - "circular_buffer": Keep N most recent aggregates (fixed-count cleanup) (default) +# - "read_based": Remove aggregates after N reads; windows no query reads are never removed # - "no_cleanup": Never clean up aggregates (retain all indefinitely) aggregate_cleanup: - policy: "read_based" + policy: "circular_buffer" # Sketch parameters configuration # These control the accuracy and memory usage of sketch data structures diff --git a/asap-tools/experiments/constants.py b/asap-tools/experiments/constants.py index 487e7e2c..70bd484c 100644 --- a/asap-tools/experiments/constants.py +++ b/asap-tools/experiments/constants.py @@ -44,6 +44,9 @@ PROMETHEUS_REMOTE_WRITE_ENDPOINT = "/api/v1/write" HTTP_READINESS_CURL_TIMEOUT_SECONDS = 5 QUERY_ENGINE_RUNTIME_INFO_ENDPOINT = "/api/v1/status/runtimeinfo" +# Query engine spill files live in this subdirectory of its output dir, which +# the container mounts, so host and container paths agree. Results syncs skip it. +QUERY_ENGINE_SPILL_SUBDIR = "spill" # VictoriaMetrics configuration files VMAGENT_SCRAPE_CONFIG_FILE = "vmagent_scrape.yml" diff --git a/asap-tools/experiments/experiment_run_clickhouse.py b/asap-tools/experiments/experiment_run_clickhouse.py index 7280945d..4aa46b62 100644 --- a/asap-tools/experiments/experiment_run_clickhouse.py +++ b/asap-tools/experiments/experiment_run_clickhouse.py @@ -174,6 +174,7 @@ def _generate_clickhouse_controller_output( dataset_cfg, cfg.get("sketch_parameters", None), cfg.get("windowing", None), + cleanup_policy=cfg.aggregate_cleanup.policy, ) local_planner_input = os.path.join(local_controller_dir, "planner_input.yaml") with open(local_planner_input, "w") as planner_input_file: @@ -667,6 +668,10 @@ def generate_controller_output() -> None: "forward_unsupported_queries": True, }, http_port=asap_http_port, + spill=cfg.query_engine.spill, + retention_slack_ms=config.get_retention_slack_ms( + ep.query_groups, data_ingestion_interval_ms + ), ingest_json_config={ "path": remote_data_file, "metric_name": str(dataset_cfg.metric_name), @@ -712,6 +717,7 @@ def generate_controller_output() -> None: pre_query_wait_seconds=steady_state_wait, monitor_interval_seconds=monitor_interval_seconds, ) + spill_error = query_engine_service.record_spill() sync.rsync_experiment_data( provider, @@ -722,6 +728,10 @@ def generate_controller_output() -> None: if not no_teardown: query_engine_service.stop() + # Raised only after sync and teardown, so results come back and the + # node is clean for the next run. + if spill_error: + raise RuntimeError(spill_error) else: # --- baseline mode --- diff --git a/asap-tools/experiments/experiment_run_e2e.py b/asap-tools/experiments/experiment_run_e2e.py index 13273a12..2f13cd41 100644 --- a/asap-tools/experiments/experiment_run_e2e.py +++ b/asap-tools/experiments/experiment_run_e2e.py @@ -428,6 +428,11 @@ def main(cfg: DictConfig): lock_strategy=args.lock_strategy, backend_config=backend_config, http_port=http_port, + spill=args.spill, + retention_slack_ms=config.get_retention_slack_ms( + cfg.experiment_params.get("query_groups"), + data_ingestion_interval_ms, + ), remote_write_port=args.remote_write_base_port, ) # For precompute mode the query engine IS the Prometheus remote-write @@ -551,6 +556,10 @@ def main(cfg: DictConfig): ): sync.copy_prometheus_data(provider, local_experiment_dir, args.node_offset) + spill_error = None + if experiment_mode == constants.SKETCHDB_EXPERIMENT_NAME: + spill_error = query_engine_service.record_spill() + # Skip teardown if the no_teardown flag is set if not args.no_teardown: if experiment_mode == constants.SKETCHDB_EXPERIMENT_NAME: @@ -597,6 +606,10 @@ def main(cfg: DictConfig): local_experiment_dir, node_offset=args.node_offset, ) + # Raised only after teardown and sync, so results come back and the + # node is clean for the next run. + if spill_error: + raise RuntimeError(spill_error) if __name__ == "__main__": diff --git a/asap-tools/experiments/experiment_run_grafana_demo.py b/asap-tools/experiments/experiment_run_grafana_demo.py index 4f639697..9af9ca82 100644 --- a/asap-tools/experiments/experiment_run_grafana_demo.py +++ b/asap-tools/experiments/experiment_run_grafana_demo.py @@ -360,6 +360,12 @@ def main(cfg: DictConfig): lock_strategy=args.lock_strategy, backend_config=backend_config, http_port=http_port, + # The demo has no end-of-run step to check spill writes, so it never spills. + spill=False, + retention_slack_ms=config.get_retention_slack_ms( + cfg.experiment_params.get("query_groups"), + data_ingestion_interval_ms, + ), remote_write_port=args.remote_write_base_port, ) diff --git a/asap-tools/experiments/experiment_utils/config.py b/asap-tools/experiments/experiment_utils/config.py index d0efd7ce..5c64beb8 100644 --- a/asap-tools/experiments/experiment_utils/config.py +++ b/asap-tools/experiments/experiment_utils/config.py @@ -538,6 +538,21 @@ def get_prometheus_data_ingestion_interval_ms(prometheus_config): raise ValueError(f"Invalid scrape interval string: {s}") +def get_retention_slack_ms(query_groups, data_ingestion_interval_ms: int) -> int: + """How far in the past the query client evaluates (the largest + client_options.query_time_offset, in seconds), plus one ingestion interval + for windows still closing. The query engine keeps this much extra + CircularBuffer history so delayed queries find their windows.""" + max_offset_s = max( + ( + (group.get("client_options") or {}).get("query_time_offset", 0) + for group in query_groups or [] + ), + default=0, + ) + return int(max_offset_s * 1000) + data_ingestion_interval_ms + + PLANNER_OUTPUT_FILES = ("streaming_config.yaml", "inference_config.yaml") QUERY_LANGUAGE_PROMQL = "promql" QUERY_LANGUAGE_SQL = "sql" @@ -702,6 +717,7 @@ def __init__(self, cfg: DictConfig): # Query engine options self.lock_strategy = cfg.query_engine.lock_strategy + self.spill = cfg.query_engine.spill # Controller configuration self.controller_punting = cfg.controller.punting @@ -968,6 +984,8 @@ def generate_sql_planner_input( dataset_cfg: Any, sketch_parameters: Any = None, windowing: Any = None, + *, + cleanup_policy: str, ) -> str: """Generate the YAML input file for asap-planner in SQL mode. @@ -998,6 +1016,7 @@ def generate_sql_planner_input( windowing: Optional DictConfig/dict mirroring ``config.yaml``'s top-level ``windowing`` section. When ``None``, the planner uses its default tumbling-window behavior. + cleanup_policy: ``aggregate_cleanup.policy`` from ``config.yaml``. Returns: YAML string ready to write to disk and pass to asap-planner. @@ -1065,7 +1084,7 @@ def generate_sql_planner_input( planner_input = { "tables": tables, "query_groups": planner_query_groups, - "aggregate_cleanup": {"policy": "read_based"}, + "aggregate_cleanup": {"policy": cleanup_policy}, } if sketch_parameters is not None: if isinstance(sketch_parameters, (DictConfig, ListConfig)): diff --git a/asap-tools/experiments/experiment_utils/services/query_engine.py b/asap-tools/experiments/experiment_utils/services/query_engine.py index a98d7c84..29db40a8 100644 --- a/asap-tools/experiments/experiment_utils/services/query_engine.py +++ b/asap-tools/experiments/experiment_utils/services/query_engine.py @@ -2,6 +2,7 @@ Query Engine service management for experiments. """ +import json import os import subprocess from typing import Optional @@ -13,6 +14,19 @@ from .base import BaseService from experiment_utils.providers.base import InfrastructureProvider +SPILL_SUBDIR = constants.QUERY_ENGINE_SPILL_SUBDIR +SPILL_ERRORS_METRIC = "asap_store_spill_write_errors_total" +SPILL_REPORT_FILE = "spill_report.json" + + +def parse_spill_write_errors(metrics_text: str) -> int: + """The spill write-error count from the engine's Prometheus /metrics text.""" + for line in metrics_text.splitlines(): + fields = line.split() + if len(fields) == 2 and fields[0] == SPILL_ERRORS_METRIC: + return int(float(fields[1])) + raise ValueError(f"{SPILL_ERRORS_METRIC} missing from the engine's /metrics") + def resolve_backend_config( backend: dict, @@ -82,6 +96,10 @@ def __init__( """ super().__init__(provider, use_container, node_offset) self.container_name = constants.QUERY_ENGINE_RS_CONTAINER_NAME + # Set by start() when spill is on; read by record_spill(). + self.remote_spill_dir: Optional[str] = None + self.local_spill_report: Optional[str] = None + self.http_port: Optional[int] = None # ------------------------------------------------------------------ # Config-file helpers (used by both bare-metal and containerized paths) @@ -100,6 +118,8 @@ def _build_engine_config( lock_strategy: str, profile_query_engine: bool, ingest_json_config: Optional[dict], + spill: bool, + retention_slack_ms: int, ) -> dict: """ Build an EngineConfig dict matching asap-query-engine's engine_config.rs schema. @@ -119,6 +139,8 @@ def _build_engine_config( remote_write_port: Port to listen on for Prometheus remote write; should match streaming.remote_write.base_port in the Hydra config lock_strategy: Lock strategy for SimpleMapStore ('global' or 'per-key') + spill: Append evicted windows under /spill instead of dropping them + retention_slack_ms: Extra CircularBuffer history for queries evaluated in the past Returns: Dict matching the EngineConfig YAML schema @@ -137,6 +159,13 @@ def _build_engine_config( "port": remote_write_port, } + store = { + "lock_strategy": lock_strategy, + "retention_slack_ms": retention_slack_ms, + } + if spill: + store["spill_dir"] = os.path.join(output_dir, SPILL_SUBDIR) + return { "output_dir": output_dir, "log_level": log_level, @@ -144,7 +173,7 @@ def _build_engine_config( "streaming_engine": streaming_engine, "http_server": {"port": http_port}, "backend": backend, # already fully resolved by caller - "store": {"lock_strategy": lock_strategy}, + "store": store, "ingest": ingest, "inference_config": os.path.join( controller_config_dir, "inference_config.yaml" @@ -212,6 +241,8 @@ def start( lock_strategy: str, backend_config: dict, http_port: int, + spill: bool, + retention_slack_ms: int, remote_write_port: int = 8080, ingest_json_config: Optional[dict] = None, ) -> None: @@ -235,9 +266,20 @@ def start( Matches the BackendConfig tagged union in asap-query-engine/src/engine_config.rs. http_port: Port for the query engine's HTTP API server + spill: Append evicted windows to disk; check them with record_spill() + retention_slack_ms: Extra CircularBuffer history for queries evaluated in + the past; see config.get_retention_slack_ms remote_write_port: Port the precompute engine listens on for Prometheus remote write; should match streaming.remote_write.base_port (default 8080) """ + if spill: + self.remote_spill_dir = os.path.join( + experiment_output_dir, "query_engine_output", SPILL_SUBDIR + ) + self.local_spill_report = os.path.join( + local_experiment_dir, "query_engine_output", SPILL_REPORT_FILE + ) + self.http_port = http_port if self.use_container: self._start_containerized( experiment_output_dir, @@ -253,6 +295,8 @@ def start( remote_write_port, lock_strategy, ingest_json_config, + spill, + retention_slack_ms, ) else: self._start_bare_metal( @@ -269,6 +313,8 @@ def start( remote_write_port, lock_strategy, ingest_json_config, + spill, + retention_slack_ms, ) def _start_bare_metal( @@ -286,6 +332,8 @@ def _start_bare_metal( remote_write_port: int, lock_strategy: str, ingest_json_config: Optional[dict], + spill: bool, + retention_slack_ms: int, ) -> None: """Start Rust QueryEngine using bare metal deployment.""" output_dir = os.path.join(experiment_output_dir, "query_engine_output") @@ -324,6 +372,8 @@ def _start_bare_metal( lock_strategy=lock_strategy, profile_query_engine=profile_query_engine, ingest_json_config=ingest_json_config, + spill=spill, + retention_slack_ms=retention_slack_ms, ) self._write_engine_config_to_remote( config_dict=config, @@ -370,6 +420,8 @@ def _start_containerized( remote_write_port: int, lock_strategy: str, ingest_json_config: Optional[dict], + spill: bool, + retention_slack_ms: int, ) -> None: """Start Rust QueryEngine using containerized deployment with Jinja template.""" output_dir = os.path.join(experiment_output_dir, "query_engine_output") @@ -393,6 +445,8 @@ def _start_containerized( lock_strategy=lock_strategy, profile_query_engine=profile_query_engine, ingest_json_config=ingest_json_config, + spill=spill, + retention_slack_ms=retention_slack_ms, ) # Write the config to the host path that is volume-mounted as /app/outputs, # so the container finds it at /app/outputs/engine_config.yaml. @@ -452,6 +506,46 @@ def _start_containerized( print(f"Failed to start Rust QueryEngine container: {e}") raise + def record_spill(self) -> Optional[str]: + """ + Write the spill directory's size and the engine's spill write-error count + to query_engine_output/spill_report.json. Call while the engine is still + running. Returns why the run should fail (a failed spill write, or the + check itself failing), or None. Never raises, so callers can sync results + and tear down before failing the run. No-op when spill is off. + """ + if self.remote_spill_dir is None: + return None + try: + spill_bytes = int( + self._run_on_node(f"du -sb {self.remote_spill_dir} | cut -f1") + ) + metrics = self._run_on_node( + f"curl -fsS --max-time {constants.HTTP_READINESS_CURL_TIMEOUT_SECONDS} " + f"http://127.0.0.1:{self.http_port}/metrics" + ) + write_errors = parse_spill_write_errors(metrics) + with open(self.local_spill_report, "w") as f: + json.dump({"spill_bytes": spill_bytes, "write_errors": write_errors}, f) + except Exception as e: + return f"spill check failed: {e}" + print(f"Spill: {spill_bytes} bytes on disk, {write_errors} failed writes") + if write_errors > 0: + return f"{write_errors} spill writes failed; see the query engine log" + return None + + def _run_on_node(self, cmd: str) -> str: + result = self.provider.execute_command( + node_idx=self.node_offset, + cmd=cmd, + cmd_dir=None, + nohup=False, + popen=False, + ignore_errors=False, + ) + assert isinstance(result, subprocess.CompletedProcess) + return result.stdout.strip() + def stop(self, **kwargs) -> None: """ Stop the Rust query engine process. diff --git a/asap-tools/experiments/experiment_utils/sync.py b/asap-tools/experiments/experiment_utils/sync.py index d0dc4983..bab49b67 100644 --- a/asap-tools/experiments/experiment_utils/sync.py +++ b/asap-tools/experiments/experiment_utils/sync.py @@ -38,10 +38,12 @@ def rsync_experiment_data( local_experiment_dir: str, node_offset: int, ): - """Sync experiment data from remote to local machine.""" + """Sync experiment data from remote to local machine. The query engine's + spill files (GBs per run) stay on the node.""" if not provider.is_remote(): return - cmd = 'rsync -azh -e "ssh {}" {}@node{}.{}:{}/ {}/'.format( + cmd = 'rsync -azh --exclude "query_engine_output/{}/" -e "ssh {}" {}@node{}.{}:{}/ {}/'.format( + constants.QUERY_ENGINE_SPILL_SUBDIR, constants.SSH_OPTIONS, provider.username, node_offset, diff --git a/asap-tools/experiments/tests/test_query_engine_spill.py b/asap-tools/experiments/tests/test_query_engine_spill.py new file mode 100644 index 00000000..54420129 --- /dev/null +++ b/asap-tools/experiments/tests/test_query_engine_spill.py @@ -0,0 +1,108 @@ +"""Tests for the query engine's spill config and end-of-run spill report.""" + +import json +import os +import subprocess +import tempfile +import unittest +from unittest.mock import Mock + +from experiment_utils.services.query_engine import QueryEngineRustService + +METRICS = """# HELP asap_store_spill_write_errors_total Evicted-window batches ... +# TYPE asap_store_spill_write_errors_total counter +asap_store_spill_write_errors_total {errors} +""" + + +def completed(stdout): + return subprocess.CompletedProcess(args="", returncode=0, stdout=stdout, stderr="") + + +def engine_store(spill): + service = QueryEngineRustService(Mock(), use_container=False, node_offset=0) + config = service._build_engine_config( + output_dir="/out/query_engine_output", + data_ingestion_interval_ms=1000, + log_level="INFO", + streaming_engine="precompute", + controller_config_dir="/ctl", + backend={"type": "prometheus", "forward_unsupported_queries": True}, + http_port=8088, + remote_write_port=8080, + lock_strategy="per-key", + profile_query_engine=False, + ingest_json_config=None, + spill=spill, + retention_slack_ms=11_000, + ) + return config["store"] + + +class SpillConfigTest(unittest.TestCase): + def test_spill_off_leaves_spill_dir_unset(self): + self.assertEqual( + engine_store(False), + {"lock_strategy": "per-key", "retention_slack_ms": 11_000}, + ) + + def test_spill_on_writes_under_the_engine_output_dir(self): + self.assertEqual( + engine_store(True)["spill_dir"], "/out/query_engine_output/spill" + ) + + +class RecordSpillTest(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.provider = Mock() + self.service = QueryEngineRustService( + self.provider, use_container=False, node_offset=0 + ) + self.service.remote_spill_dir = "/out/query_engine_output/spill" + self.service.local_spill_report = os.path.join(self.tmp.name, "report.json") + self.service.http_port = 8088 + + def tearDown(self): + self.tmp.cleanup() + + def test_records_size_and_error_count(self): + self.provider.execute_command.side_effect = [ + completed("4096\n"), + completed(METRICS.format(errors="0")), + ] + self.assertIsNone(self.service.record_spill()) + with open(self.service.local_spill_report) as f: + self.assertEqual(json.load(f), {"spill_bytes": 4096, "write_errors": 0}) + + def test_failed_writes_fail_the_run(self): + self.provider.execute_command.side_effect = [ + completed("4096\n"), + completed(METRICS.format(errors="2")), + ] + self.assertIn("2 spill writes failed", self.service.record_spill()) + + def test_missing_error_counter_fails_the_run(self): + # An engine without spill wiring must not pass as "no errors". + self.provider.execute_command.side_effect = [ + completed("4096\n"), + completed("other_metric 1\n"), + ] + self.assertIn("missing", self.service.record_spill()) + + def test_failed_check_is_returned_not_raised(self): + # Raising here skipped the rsync and teardown that follow, leaving the + # stack running on the node. + self.provider.execute_command.side_effect = [completed("\n")] + self.assertIn("spill check failed", self.service.record_spill()) + + def test_spill_off_is_a_no_op(self): + service = QueryEngineRustService( + self.provider, use_container=False, node_offset=0 + ) + self.assertIsNone(service.record_spill()) + self.provider.execute_command.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/asap-tools/experiments/tests/test_retention_slack.py b/asap-tools/experiments/tests/test_retention_slack.py new file mode 100644 index 00000000..caba93fb --- /dev/null +++ b/asap-tools/experiments/tests/test_retention_slack.py @@ -0,0 +1,31 @@ +"""Tests for the query engine's retention slack derived from the workload.""" + +import unittest + +from omegaconf import OmegaConf + +from experiment_utils.config import get_retention_slack_ms + + +class RetentionSlackTest(unittest.TestCase): + def test_largest_query_time_offset_plus_one_ingestion_interval(self): + # A 1s-window instant query evaluated 10s in the past fell back to + # Prometheus every time until retention covered the offset. + groups = OmegaConf.create( + [ + {"client_options": {"query_time_offset": 10}}, + {"client_options": {"query_time_offset": 4}}, + ] + ) + self.assertEqual(get_retention_slack_ms(groups, 1000), 11_000) + + def test_missing_offset_counts_as_zero(self): + groups = OmegaConf.create([{"client_options": {}}, {}]) + self.assertEqual(get_retention_slack_ms(groups, 15_000), 15_000) + + def test_no_query_groups(self): + self.assertEqual(get_retention_slack_ms(None, 1000), 1000) + + +if __name__ == "__main__": + unittest.main() diff --git a/asap-tools/experiments/tests/test_sql_planner_input.py b/asap-tools/experiments/tests/test_sql_planner_input.py index 007c67b9..a8931931 100644 --- a/asap-tools/experiments/tests/test_sql_planner_input.py +++ b/asap-tools/experiments/tests/test_sql_planner_input.py @@ -32,11 +32,13 @@ def test_unknown_controller_option_is_rejected(self): } ) with self.assertRaisesRegex(ValueError, "latency_sla"): - generate_sql_planner_input(groups, dataset_cfg) + generate_sql_planner_input( + groups, dataset_cfg, cleanup_policy="circular_buffer" + ) class LatencySlaMsTest(unittest.TestCase): - def _planner_latency(self, controller_options): + def _planner_input(self, controller_options, cleanup_policy="circular_buffer"): with tempfile.TemporaryDirectory() as tmp: sql_path = os.path.join(tmp, "q.sql") with open(sql_path, "w") as f: @@ -58,13 +60,23 @@ def _planner_latency(self, controller_options): }, } ) - planner_input = yaml.safe_load( - generate_sql_planner_input(groups, dataset_cfg) + return yaml.safe_load( + generate_sql_planner_input( + groups, dataset_cfg, cleanup_policy=cleanup_policy + ) ) + + def _planner_latency(self, controller_options): + planner_input = self._planner_input(controller_options) return planner_input["query_groups"][0]["controller_options"].get( "latency_sla_ms" ) + def test_cleanup_policy_comes_from_the_experiment_config(self): + for policy in ["circular_buffer", "read_based", "no_cleanup"]: + planner_input = self._planner_input({"accuracy_sla": 0.99}, policy) + self.assertEqual(planner_input["aggregate_cleanup"], {"policy": policy}) + def test_omitted_latency_sla_ms_means_no_limit(self): # Same meaning as the PromQL path and the planner. self.assertIsNone(self._planner_latency({"accuracy_sla": 0.99}))