From b25a6553a9425657ea5adacdde3d7069079b1015 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 1 Sep 2026 20:50:33 -0400 Subject: [PATCH 1/6] fix(query-engine): honor CountMinSketch sum/count subtype --- .../precompute_engine/accumulator_factory.rs | 144 +++++++++++++++++- .../src/precompute_engine/engine.rs | 97 +++++++++++- .../src/precompute_engine/worker.rs | 61 ++++++++ 3 files changed, 293 insertions(+), 9 deletions(-) diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index 7044cd2..c53f74e 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -524,14 +524,16 @@ pub struct CmsAccumulatorUpdater { acc: CountMinSketchAccumulator, row_num: usize, col_num: usize, + count_events: bool, } impl CmsAccumulatorUpdater { - pub fn new(row_num: usize, col_num: usize) -> Self { + pub fn new(row_num: usize, col_num: usize, count_events: bool) -> Self { Self { acc: CountMinSketchAccumulator::new(row_num, col_num), row_num, col_num, + count_events, } } } @@ -545,7 +547,8 @@ impl AccumulatorUpdater for CmsAccumulatorUpdater { } fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { - self.acc.inner.update(&key.to_semicolon_str(), value); + let weight = if self.count_events { 1.0 } else { value }; + self.acc.inner.update(&key.to_semicolon_str(), weight); } impl_accumulator_methods!(acc); @@ -732,6 +735,24 @@ fn cms_params(config: &AggregationConfig) -> Result<(usize, usize), String> { Ok((row_num, col_num)) } +/// Resolve the weighting semantics for a plain Count-Min Sketch. +/// +/// Unlike the heap variant, plain CMS uses `aggregation_sub_type` to +/// distinguish approximate SUM from approximate COUNT. Do not silently +/// default malformed configs: the wrong weighting produces plausible but +/// incorrect results. +fn cms_count_events_for_sub_type(sub_type: &str) -> Result { + if sub_type.eq_ignore_ascii_case("count") { + Ok(true) + } else if sub_type.eq_ignore_ascii_case("sum") { + Ok(false) + } else { + Err(format!( + "CountMinSketch requires aggregation_sub_type 'sum' or 'count', got '{sub_type}'" + )) + } +} + /// Extract `(row_num, col_num, k)` for HydraKLL configs. fn hydra_kll_params(config: &AggregationConfig) -> Result<(usize, usize, u16), String> { let (row_num, col_num) = cms_params(config)?; @@ -834,7 +855,9 @@ pub fn create_accumulator_updater( "Increase" | "increase" => Ok(Box::new(MultipleIncreaseAccumulatorUpdater::new())), "CountMinSketch" | "count_min_sketch" | "CMS" | "cms" => { let (row_num, col_num) = cms_params(config)?; - Ok(Box::new(CmsAccumulatorUpdater::new(row_num, col_num))) + Ok(Box::new(CmsAccumulatorUpdater::new( + row_num, col_num, false, + ))) } "HydraKLL" | "hydra_kll" => { let (row_num, col_num, k) = hydra_kll_params(config)?; @@ -867,7 +890,12 @@ pub fn create_accumulator_updater( AggregationType::Increase => Ok(Box::new(IncreaseAccumulatorUpdater::new())), AggregationType::CountMinSketch => { let (row_num, col_num) = cms_params(config)?; - Ok(Box::new(CmsAccumulatorUpdater::new(row_num, col_num))) + let count_events = cms_count_events_for_sub_type(sub_type)?; + Ok(Box::new(CmsAccumulatorUpdater::new( + row_num, + col_num, + count_events, + ))) } AggregationType::CountMinSketchWithHeap => { let (row_num, col_num, heap_size) = cms_heap_params(config)?; @@ -1032,7 +1060,7 @@ mod tests { ))); assert!(config_is_keyed(&make_config( AggregationType::CountMinSketch, - "" + "sum" ))); assert!(config_is_keyed(&make_config(AggregationType::HydraKLL, ""))); @@ -1079,7 +1107,11 @@ mod tests { }; for (agg_type, sub_type, params) in [ (AggregationType::DatasketchesKLL, "", kll_params_required()), - (AggregationType::CountMinSketch, "", cms_params_required()), + ( + AggregationType::CountMinSketch, + "sum", + cms_params_required(), + ), ] { let config = make_config_with_params(agg_type, sub_type, params); let updater = create_accumulator_updater(&config).unwrap(); @@ -1323,7 +1355,7 @@ mod tests { let config = AggregationConfig::new( 1, AggregationType::CountMinSketch, - String::new(), + "sum".to_string(), HashMap::new(), promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), @@ -1359,7 +1391,7 @@ mod tests { let config = AggregationConfig::new( 21, AggregationType::CountMinSketch, - String::new(), + "sum".to_string(), params, promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), @@ -1397,6 +1429,102 @@ mod tests { p } + fn cms_config(sub_type: &str) -> AggregationConfig { + AggregationConfig::new( + 100, + AggregationType::CountMinSketch, + sub_type.to_string(), + cms_params_required(), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + 1_000, + 1_000, + WindowType::Tumbling, + "test_metric".to_string(), + "test_metric".to_string(), + None, + None, + None, + None, + ) + } + + #[test] + fn test_cms_count_subtype_uses_unit_weight() { + let config = cms_config("count"); + let mut updater = create_accumulator_updater(&config).unwrap(); + let key = KeyByLabelValues::new_with_labels(vec!["host-a".to_string()]); + + for _ in 0..5 { + updater.update_keyed(&key, 1_000.0, 0); + } + + let acc = updater.take_accumulator(); + let cms = acc + .as_any() + .downcast_ref::() + .expect("CountMinSketch accumulator"); + assert_eq!(cms.query_key(&key), 5.0); + } + + #[test] + fn test_cms_sum_subtype_uses_sample_weight() { + let config = cms_config("sum"); + let mut updater = create_accumulator_updater(&config).unwrap(); + let key = KeyByLabelValues::new_with_labels(vec!["host-a".to_string()]); + + for _ in 0..5 { + updater.update_keyed(&key, 10.0, 0); + } + + let acc = updater.take_accumulator(); + let cms = acc + .as_any() + .downcast_ref::() + .expect("CountMinSketch accumulator"); + assert_eq!(cms.query_key(&key), 50.0); + } + + #[test] + fn test_cms_rejects_empty_subtype() { + let config = cms_config(""); + let err = match create_accumulator_updater(&config) { + Ok(_) => panic!("empty CountMinSketch subtype must fail"), + Err(err) => err, + }; + assert!(err.contains("sum") && err.contains("count")); + } + + #[test] + fn test_cms_rejects_unknown_subtype() { + let config = cms_config("frequency"); + let err = match create_accumulator_updater(&config) { + Ok(_) => panic!("unknown CountMinSketch subtype must fail"), + Err(err) => err, + }; + assert!(err.contains("frequency")); + } + + #[test] + fn test_cms_accepts_case_insensitive_subtype() { + for sub_type in ["COUNT", "SuM"] { + create_accumulator_updater(&cms_config(sub_type)) + .unwrap_or_else(|err| panic!("subtype '{sub_type}' should be accepted: {err}")); + } + } + + #[test] + fn test_cms_rejects_whitespace_padded_subtype() { + let config = cms_config(" count "); + let err = match create_accumulator_updater(&config) { + Ok(_) => panic!("whitespace-padded CountMinSketch subtype must fail"), + Err(err) => err, + }; + assert!(err.contains(" count ")); + } + fn cms_heap_params_required() -> std::collections::HashMap { let mut p = std::collections::HashMap::new(); p.insert("depth".to_string(), serde_json::json!(3_u64)); diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index 7fbc9ce..abe882b 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -1,4 +1,5 @@ -use crate::data_model::StreamingConfig; +use crate::data_model::{AggregationType, StreamingConfig}; +use crate::precompute_engine::accumulator_factory::create_accumulator_updater; use crate::precompute_engine::config::PrecomputeEngineConfig; use crate::precompute_engine::ingest_source::{IngestContext, IngestSource}; use crate::precompute_engine::output_sink::OutputSink; @@ -154,6 +155,8 @@ impl PrecomputeEngine { /// Start the precompute engine. This spawns worker tasks and all registered /// ingest sources, then blocks until shutdown. pub async fn run(mut self) -> Result<(), Box> { + validate_startup_aggregation_configs(&self.streaming_config)?; + let num_workers = self.config.num_workers; let receivers = self @@ -252,3 +255,95 @@ impl PrecomputeEngine { Ok(()) } } + +/// Validate aggregation configs before starting any worker or ingest task. +/// +/// Plain Count-Min Sketch configs carry their SUM-versus-COUNT contract in +/// `aggregation_sub_type`; allowing an invalid value to reach the lazy worker +/// path would leave the engine running while silently losing that contract. +fn validate_startup_aggregation_configs( + streaming_config: &StreamingConfig, +) -> Result<(), Box> { + for (&aggregation_id, config) in streaming_config.get_all_aggregation_configs() { + if config.aggregation_type != AggregationType::CountMinSketch { + continue; + } + + create_accumulator_updater(config).map(|_| ()).map_err( + |err| -> Box { + format!("invalid aggregation config for aggregation_id {aggregation_id}: {err}") + .into() + }, + )?; + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::data_model::{AggregationType, StreamingConfig, WindowType}; + use crate::precompute_engine::config::LateDataPolicy; + use crate::precompute_engine::ingest_source::{IngestContext, IngestSource}; + use crate::precompute_engine::output_sink::NoopOutputSink; + use async_trait::async_trait; + use serde_json::json; + + struct ShutdownSource; + + #[async_trait] + impl IngestSource for ShutdownSource { + async fn run( + self: Box, + ctx: IngestContext, + ) -> Result<(), Box> { + ctx.router.broadcast_shutdown().await + } + } + + #[tokio::test] + async fn run_rejects_invalid_cms_subtype_before_starting_workers() { + let mut parameters = HashMap::new(); + parameters.insert("depth".to_string(), json!(3_u64)); + parameters.insert("width".to_string(), json!(128_u64)); + let cms = AggregationConfig::new( + 1, + AggregationType::CountMinSketch, + String::new(), + parameters, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![ + "host".to_string() + ]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + 1_000, + 1_000, + WindowType::Tumbling, + "requests_total".to_string(), + "requests_total".to_string(), + None, + None, + None, + None, + ); + let engine = PrecomputeEngine::new( + PrecomputeEngineConfig { + num_workers: 1, + late_data_policy: LateDataPolicy::Drop, + ..PrecomputeEngineConfig::default() + }, + Arc::new(StreamingConfig::new(HashMap::from([(1, cms)]))), + Arc::new(NoopOutputSink::new()), + vec![Box::new(ShutdownSource)], + ); + + let result = engine.run().await; + let err = match result { + Ok(()) => panic!("invalid CMS subtype must fail before startup"), + Err(err) => err, + }; + assert!(err.to_string().contains("aggregation_id 1")); + assert!(err.to_string().contains("sum") && err.to_string().contains("count")); + } +} diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index f2e151b..b59a14e 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -1242,6 +1242,7 @@ mod tests { use crate::precompute_operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; use crate::precompute_operators::multiple_sum_accumulator::MultipleSumAccumulator; use crate::precompute_operators::sum_accumulator::SumAccumulator; + use crate::precompute_operators::CountMinSketchAccumulator; use asap_sketchlib::KllSketch; use asap_types::enums::{AggregationType, WindowType}; @@ -1349,6 +1350,66 @@ mod tests { .collect() } + #[test] + fn test_count_min_sketch_count_subtype_counts_events_through_worker() { + let mut config = make_agg_config_full( + 6, + "requests_total", + AggregationType::CountMinSketch, + "count", + 1_000, + 1_000, + vec![], + vec!["host"], + ); + config + .parameters + .insert("depth".to_string(), serde_json::json!(3_u64)); + config + .parameters + .insert("width".to_string(), serde_json::json!(128_u64)); + + let sink = Arc::new(CapturingOutputSink::new()); + let mut worker = make_worker( + arc_configs(HashMap::from([(6, config)])), + sink.clone(), + false, + 0, + LateDataPolicy::Drop, + ); + + worker + .process_group_samples( + 6, + "", + vec![ + ("requests_total{host=\"A\"}".to_string(), 100, 100.0), + ("requests_total{host=\"A\"}".to_string(), 200, 200.0), + ], + ) + .unwrap(); + + worker + .process_group_samples( + 6, + "", + group_samples("requests_total{host=\"A\"}", vec![(5_000, 1.0)]), + ) + .unwrap(); + + let captured = sink.drain(); + let (_output, acc) = captured + .iter() + .find(|(output, _)| output.start_timestamp == 0) + .expect("worker should emit the closed [0, 1000) window"); + let cms = acc + .as_any() + .downcast_ref::() + .expect("worker should emit a CountMinSketch accumulator"); + let key = KeyByLabelValues::new_with_labels(vec!["A".to_string()]); + assert_eq!(cms.query_key(&key), 2.0); + } + // ----------------------------------------------------------------------- // Test: raw mode — each sample forwarded as SumAccumulator with sum==value // ----------------------------------------------------------------------- From 2f0c9139a4970bfe7362ed0a8f4d85c5d16d05fd Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 1 Sep 2026 21:21:50 -0400 Subject: [PATCH 2/6] fix(query-engine): invalidate changed aggregation state --- .../rs/asap_types/src/aggregation_config.rs | 2 +- .../rs/asap_types/src/capability_matching.rs | 167 ++++++++++++++++-- .../src/precompute_engine/engine.rs | 57 +++++- .../src/precompute_engine/worker.rs | 124 +++++++++++-- 4 files changed, 322 insertions(+), 28 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs index 9351610..02fb3a2 100644 --- a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs @@ -9,7 +9,7 @@ use crate::utils::normalize_spatial_filter; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::AggregationType; -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct AggregationConfig { pub aggregation_id: u64, pub aggregation_type: AggregationType, diff --git a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs index 2a6f3de..30218da 100644 --- a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs +++ b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs @@ -18,7 +18,11 @@ use promql_utilities::query_logics::enums::AggregationType; /// Returns the aggregation types that can serve this statistic. pub fn compatible_agg_types(stat: Statistic) -> &'static [AggregationType] { match stat { - Statistic::Sum => &[AggregationType::Sum, AggregationType::MultipleSum], + Statistic::Sum => &[ + AggregationType::Sum, + AggregationType::MultipleSum, + AggregationType::CountMinSketch, + ], Statistic::Count => &[ AggregationType::CountMinSketch, AggregationType::CountMinSketchWithHeap, @@ -187,6 +191,20 @@ pub fn topk_weighting_compatible( } } +/// Plain Count-Min Sketches are value-weighted or event-weighted according to +/// their subtype. A sketch with the other subtype cannot serve this statistic. +fn plain_cms_sub_type_compatible(stat: Statistic, config: &AggregationConfig) -> bool { + if config.aggregation_type != AggregationType::CountMinSketch { + return true; + } + + match stat { + Statistic::Sum => config.aggregation_sub_type.eq_ignore_ascii_case("sum"), + Statistic::Count => config.aggregation_sub_type.eq_ignore_ascii_case("count"), + _ => false, + } +} + /// Aggregation priority comparator: prefer larger `window_size_ms` (descending). /// This is a separate function so callers can swap the policy without touching matching logic. pub fn aggregation_priority(a: &AggregationConfig, b: &AggregationConfig) -> Ordering { @@ -244,6 +262,7 @@ pub fn find_compatible_aggregation( &c.spatial_filter_normalized, &requirements.spatial_filter_normalized, ) + && plain_cms_sub_type_compatible(stat, c) && topk_weighting_compatible(stat, c, requirements.topk_count_events); if !ok { debug!( @@ -466,6 +485,62 @@ mod tests { assert_eq!(result.unwrap().aggregation_id_for_value, 1); } + #[test] + fn plain_cms_matching_respects_sum_and_count_subtypes() { + let mut configs = HashMap::new(); + configs.insert( + 1, + make_config( + 1, + "cpu", + "CountMinSketch", + "sum", + 300_000, + "tumbling", + &[], + "", + ), + ); + configs.insert( + 2, + make_config( + 2, + "cpu", + "CountMinSketch", + "count", + 300_000, + "tumbling", + &[], + "", + ), + ); + configs.insert( + 9, + make_config( + 9, + "cpu", + "DeltaSetAggregator", + "", + 300_000, + "tumbling", + &[], + "", + ), + ); + + let sum = + find_compatible_aggregation(&configs, &req("cpu", &[Statistic::Sum], 300_000, &[], "")) + .expect("SUM should select the value-weighted sketch"); + assert_eq!(sum.aggregation_id_for_value, 1); + + let count = find_compatible_aggregation( + &configs, + &req("cpu", &[Statistic::Count], 300_000, &[], ""), + ) + .expect("COUNT should select the event-weighted sketch"); + assert_eq!(count.aggregation_id_for_value, 2); + } + #[test] fn quantile_any_value_finds_kll() { let configs = single_config(make_config( @@ -1076,7 +1151,16 @@ mod tests { #[test] fn multi_pop_rejects_tumbling_delta_set_that_cannot_partition_sliding_value_window() { - let mut value = make_config(10, "req", "CountMinSketch", "", 6_000, "sliding", &[], ""); + let mut value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 6_000, + "sliding", + &[], + "", + ); value.slide_interval_ms = 1_000; let delta_keys = make_config( 11, @@ -1102,7 +1186,16 @@ mod tests { #[test] fn multi_pop_accepts_tumbling_delta_set_that_partitions_sliding_value_grid() { - let mut value = make_config(10, "req", "CountMinSketch", "", 6_000, "sliding", &[], ""); + let mut value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 6_000, + "sliding", + &[], + "", + ); value.slide_interval_ms = 1_000; let delta_keys = make_config( 11, @@ -1127,7 +1220,16 @@ mod tests { #[test] fn multi_pop_rejects_tumbling_set_key_on_mismatched_nonzero_grid_step() { - let value = make_config(10, "req", "CountMinSketch", "", 5_000, "tumbling", &[], ""); + let value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 5_000, + "tumbling", + &[], + "", + ); let mut keys = make_config(11, "req", "SetAggregator", "", 5_000, "tumbling", &[], ""); keys.slide_interval_ms = 1_000; let configs = HashMap::from([(10, value), (11, keys)]); @@ -1141,7 +1243,16 @@ mod tests { #[test] fn tumbling_set_pairing_normalizes_zero_slide_to_window_size() { - let mut value = make_config(10, "req", "CountMinSketch", "", 5_000, "tumbling", &[], ""); + let mut value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 5_000, + "tumbling", + &[], + "", + ); let mut key = make_config(11, "req", "SetAggregator", "", 5_000, "tumbling", &[], ""); value.slide_interval_ms = 0; key.slide_interval_ms = 0; @@ -1152,7 +1263,16 @@ mod tests { #[test] fn set_pairing_rejects_each_grid_mismatch_dimension() { - let value = make_config(10, "req", "CountMinSketch", "", 5_000, "sliding", &[], ""); + let value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 5_000, + "sliding", + &[], + "", + ); let mut key = make_config(11, "req", "SetAggregator", "", 5_000, "sliding", &[], ""); key.slide_interval_ms = 1_000; assert!(!key_agg_compatible_with_value(&value, &key)); @@ -1166,7 +1286,16 @@ mod tests { #[test] fn delta_set_pairing_truth_table_checks_both_divisors() { - let mut value = make_config(10, "req", "CountMinSketch", "", 6_000, "sliding", &[], ""); + let mut value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 6_000, + "sliding", + &[], + "", + ); value.slide_interval_ms = 2_000; let key_valid = make_config( 11, @@ -1212,7 +1341,16 @@ mod tests { #[test] fn delta_set_pairing_for_tumbling_values_does_not_apply_sliding_rules() { - let value = make_config(10, "req", "CountMinSketch", "", 6_000, "tumbling", &[], ""); + let value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 6_000, + "tumbling", + &[], + "", + ); let key = make_config( 11, "req", @@ -1228,7 +1366,16 @@ mod tests { #[test] fn matching_skips_incompatible_key_candidate_and_selects_compatible_one() { - let value = make_config(10, "req", "CountMinSketch", "", 6_000, "sliding", &[], ""); + let value = make_config( + 10, + "req", + "CountMinSketch", + "count", + 6_000, + "sliding", + &[], + "", + ); let mut incompatible = make_config(11, "req", "SetAggregator", "", 6_000, "sliding", &[], ""); incompatible.slide_interval_ms = 2_000; @@ -1255,7 +1402,7 @@ mod tests { 2, "cpu", "CountMinSketch", - "", + "count", 300_000, "tumbling", &["job"], diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index abe882b..2fbf2ed 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -38,6 +38,8 @@ impl PrecomputeEngineHandle { &self, config: &StreamingConfig, ) -> Result<(), Box> { + validate_cms_aggregation_configs(config)?; + let agg_configs_map: HashMap> = config .get_all_aggregation_configs() .iter() @@ -155,7 +157,7 @@ impl PrecomputeEngine { /// Start the precompute engine. This spawns worker tasks and all registered /// ingest sources, then blocks until shutdown. pub async fn run(mut self) -> Result<(), Box> { - validate_startup_aggregation_configs(&self.streaming_config)?; + validate_cms_aggregation_configs(&self.streaming_config)?; let num_workers = self.config.num_workers; @@ -256,12 +258,13 @@ impl PrecomputeEngine { } } -/// Validate aggregation configs before starting any worker or ingest task. +/// Validate plain Count-Min Sketch configs before starting workers or ingest +/// tasks, and before accepting runtime replacements. /// /// Plain Count-Min Sketch configs carry their SUM-versus-COUNT contract in /// `aggregation_sub_type`; allowing an invalid value to reach the lazy worker /// path would leave the engine running while silently losing that contract. -fn validate_startup_aggregation_configs( +fn validate_cms_aggregation_configs( streaming_config: &StreamingConfig, ) -> Result<(), Box> { for (&aggregation_id, config) in streaming_config.get_all_aggregation_configs() { @@ -346,4 +349,52 @@ mod tests { assert!(err.to_string().contains("aggregation_id 1")); assert!(err.to_string().contains("sum") && err.to_string().contains("count")); } + + #[tokio::test] + async fn runtime_update_rejects_invalid_cms_subtype_without_replacing_config() { + let mut parameters = HashMap::new(); + parameters.insert("depth".to_string(), json!(3_u64)); + parameters.insert("width".to_string(), json!(128_u64)); + let valid = AggregationConfig::new( + 1, + AggregationType::CountMinSketch, + "sum".to_string(), + parameters, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![ + "host".to_string() + ]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + 1_000, + 1_000, + WindowType::Tumbling, + "requests_total".to_string(), + "requests_total".to_string(), + None, + None, + None, + None, + ); + let engine = PrecomputeEngine::new( + PrecomputeEngineConfig::default(), + Arc::new(StreamingConfig::new(HashMap::from([(1, valid.clone())]))), + Arc::new(NoopOutputSink::new()), + vec![], + ); + let handle = engine.handle(); + + let mut invalid = valid; + invalid.aggregation_sub_type.clear(); + let result = handle + .update_streaming_config(&StreamingConfig::new(HashMap::from([(1, invalid)]))) + .await; + + let err = result.expect_err("invalid runtime CMS subtype must be rejected"); + assert!(err.to_string().contains("aggregation_id 1")); + assert!(err.to_string().contains("sum") && err.to_string().contains("count")); + let current = handle.ingest_agg_configs.load(); + assert_eq!(current.len(), 1); + assert_eq!(current[0].aggregation_sub_type, "sum"); + } } diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index b59a14e..bebc288 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -240,21 +240,26 @@ impl Worker { break; } WorkerMessage::UpdateAggConfigs(new_configs) => { - // Flush and evict group states for agg IDs that are being removed. - // Must happen before swapping agg_configs so GroupState.config is - // still valid during the final window close. - let removed_ids: Vec = self + // Flush and evict group states for removed or materially changed + // aggregations. Must happen before swapping agg_configs so + // GroupState.config is still valid during the final window close. + let reset_ids: Vec = self .agg_configs - .keys() - .filter(|id| !new_configs.contains_key(id)) - .copied() + .iter() + .filter_map(|(&id, old_config)| match new_configs.get(&id) { + None => Some(id), + Some(new_config) if old_config.as_ref() != new_config.as_ref() => { + Some(id) + } + Some(_) => None, + }) .collect(); - if !removed_ids.is_empty() { + if !reset_ids.is_empty() { let mut emit_batch: Vec<(PrecomputedOutput, Box)> = Vec::new(); - for agg_id in &removed_ids { + for agg_id in &reset_ids { // Drain all group states for this agg_id in one move. let Some(inner) = self.group_states.remove(agg_id) else { continue; @@ -307,8 +312,8 @@ impl Worker { if !emit_batch.is_empty() { if let Err(e) = self.output_sink.emit_batch(emit_batch) { warn!( - "Worker {}: error flushing removed agg_ids {:?}: {}", - self.id, removed_ids, e + "Worker {}: error flushing reset agg_ids {:?}: {}", + self.id, reset_ids, e ); } } @@ -316,10 +321,10 @@ impl Worker { self.group_count .store(self.total_groups(), Ordering::Relaxed); info!( - "Worker {}: evicted {} removed agg_id(s) {:?}", + "Worker {}: evicted {} reset agg_id(s) {:?}", self.id, - removed_ids.len(), - removed_ids, + reset_ids.len(), + reset_ids, ); } @@ -2813,6 +2818,97 @@ aggregations: ); } + #[tokio::test] + async fn test_update_agg_configs_recreates_state_when_cms_subtype_changes() { + let mut sum_config = make_agg_config( + 1, + "requests_total", + AggregationType::CountMinSketch, + "sum", + 1_000, + 1_000, + vec![], + ); + sum_config + .parameters + .insert("depth".to_string(), serde_json::json!(3_u64)); + sum_config + .parameters + .insert("width".to_string(), serde_json::json!(128_u64)); + let mut count_config = sum_config.clone(); + count_config.aggregation_sub_type = "count".to_string(); + + let sink = Arc::new(CapturingOutputSink::new()); + let (tx, rx) = tokio::sync::mpsc::channel(32); + let wm = Arc::new(AtomicI64::new(i64::MIN)); + let worker = Worker::new( + 0, + rx, + sink.clone(), + arc_configs(HashMap::from([(1, sum_config)])), + WorkerRuntimeConfig { + max_buffer_per_series: 10_000, + allowed_lateness_ms: 0, + pass_raw_samples: false, + raw_mode_aggregation_id: 0, + late_data_policy: LateDataPolicy::Drop, + wall_clock_grace_period_ms: 0, + }, + Arc::new(AtomicUsize::new(0)), + wm.clone(), + vec![wm], + ); + let handle = tokio::spawn(async move { worker.run().await }); + + tx.send(WorkerMessage::GroupSamples { + agg_id: 1, + group_key: String::new(), + samples: vec![("requests_total".to_string(), 100, 10.0)], + ingest_received_at: std::time::Instant::now(), + }) + .await + .unwrap(); + tx.send(WorkerMessage::UpdateAggConfigs(arc_configs(HashMap::from( + [(1, count_config)], + )))) + .await + .unwrap(); + tx.send(WorkerMessage::GroupSamples { + agg_id: 1, + group_key: String::new(), + samples: vec![("requests_total".to_string(), 1_100, 10.0)], + ingest_received_at: std::time::Instant::now(), + }) + .await + .unwrap(); + tx.send(WorkerMessage::Shutdown).await.unwrap(); + handle.await.unwrap(); + + let mut captured = sink.drain(); + captured.sort_by_key(|(output, _)| output.start_timestamp); + assert_eq!(captured.len(), 2); + + let first = captured[0] + .1 + .as_any() + .downcast_ref::() + .expect("first output should be CMS"); + assert_eq!( + first.query_key(&KeyByLabelValues::new_with_labels(vec![])), + 10.0 + ); + + let second = captured[1] + .1 + .as_any() + .downcast_ref::() + .expect("second output should be CMS"); + assert_eq!( + second.query_key(&KeyByLabelValues::new_with_labels(vec![])), + 1.0 + ); + } + // ----------------------------------------------------------------------- // Test: removing an agg_id force-closes its open windows with a finite // bound — at realistic (epoch-ms) timestamps. From de1f4a16a14da0d700a1ca59173113fc71f78af3 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 2 Sep 2026 08:12:09 -0400 Subject: [PATCH 3/6] fix(query-engine): apply plans transactionally --- asap-query-engine/src/main.rs | 93 +++++++++++++++++++++++++++++------ 1 file changed, 77 insertions(+), 16 deletions(-) diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 95ec12e..e21a9e0 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -1,5 +1,38 @@ mod engine_config; +use std::future::Future; + +#[cfg(test)] +mod tests { + use super::*; + use asap_types::enums::{CleanupPolicy, QueryLanguage}; + use query_engine_rust::planner_client::PlannerResult; + use std::sync::atomic::{AtomicBool, Ordering}; + + #[tokio::test] + async fn rejected_precompute_plan_does_not_apply_other_components() { + let applied = Arc::new(AtomicBool::new(false)); + let applied_by_callback = applied.clone(); + let result = PlannerResult { + streaming_config: StreamingConfig::default(), + inference_config: InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::NoCleanup), + punted_queries: Vec::new(), + }; + + let update_result = apply_plan_if_precompute_succeeds( + result, + |_| async { Err::<(), _>("invalid streaming config".into()) }, + move |_| { + applied_by_callback.store(true, Ordering::SeqCst); + }, + ) + .await; + + assert!(update_result.is_err()); + assert!(!applied.load(Ordering::SeqCst)); + } +} + use clap::Parser; use engine_config::{BackendConfig, EngineConfig, IngestConfig}; use figment::{ @@ -40,6 +73,20 @@ struct Args { overrides: Vec, } +async fn apply_plan_if_precompute_succeeds( + result: query_engine_rust::planner_client::PlannerResult, + update_precompute: F, + apply_other_components: impl FnOnce(query_engine_rust::planner_client::PlannerResult), +) -> Result<()> +where + F: FnOnce(&StreamingConfig) -> Fut, + Fut: Future>, +{ + update_precompute(&result.streaming_config).await?; + apply_other_components(result); + Ok(()) +} + #[tokio::main] async fn main() -> Result<()> { let args = Args::parse(); @@ -441,23 +488,37 @@ async fn main() -> Result<()> { } let result = rx.borrow().clone(); if let Some(result) = result { - if let Some(ref handle) = pe_engine_handle { - if let Err(e) = handle - .update_streaming_config(&result.streaming_config) - .await - { - warn!("Applier: failed to update precompute engine: {}", e); - } + let apply_result = apply_plan_if_precompute_succeeds( + result, + |streaming_config| { + let streaming_config = streaming_config.clone(); + let precompute_handle = pe_engine_handle.as_ref(); + async move { + if let Some(handle) = precompute_handle { + handle.update_streaming_config(&streaming_config).await + } else { + Ok(()) + } + } + }, + |result| { + engine_for_applier + .update_streaming_config(Arc::new(result.streaming_config.clone())); + engine_for_applier + .update_inference_config(result.inference_config.clone()); + store_for_applier + .update_streaming_config(result.streaming_config.clone()); + *streaming_config_ref_for_applier.write().unwrap() = + Arc::new(result.streaming_config); + *inference_config_ref_for_applier.write().unwrap() = + Arc::new(result.inference_config); + info!("Applier: applied new plan from query tracker"); + }, + ) + .await; + if let Err(e) = apply_result { + warn!("Applier: failed to apply plan: {}", e); } - engine_for_applier - .update_streaming_config(Arc::new(result.streaming_config.clone())); - engine_for_applier.update_inference_config(result.inference_config.clone()); - store_for_applier.update_streaming_config(result.streaming_config.clone()); - *streaming_config_ref_for_applier.write().unwrap() = - Arc::new(result.streaming_config); - *inference_config_ref_for_applier.write().unwrap() = - Arc::new(result.inference_config); - info!("Applier: applied new plan from query tracker"); } } }); From de32080eddd632bcaa03120f6216fbc172faadb9 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 2 Sep 2026 09:29:09 -0400 Subject: [PATCH 4/6] refactor(query-engine): separate runtime reconfiguration work --- .../rs/asap_types/src/aggregation_config.rs | 2 +- asap-query-engine/src/main.rs | 93 +++---------- .../src/precompute_engine/engine.rs | 57 +------- .../src/precompute_engine/worker.rs | 124 ++---------------- 4 files changed, 34 insertions(+), 242 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs index 02fb3a2..9351610 100644 --- a/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs +++ b/asap-common/dependencies/rs/asap_types/src/aggregation_config.rs @@ -9,7 +9,7 @@ use crate::utils::normalize_spatial_filter; use promql_utilities::data_model::KeyByLabelNames; use promql_utilities::query_logics::enums::AggregationType; -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct AggregationConfig { pub aggregation_id: u64, pub aggregation_type: AggregationType, diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index e21a9e0..95ec12e 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -1,38 +1,5 @@ mod engine_config; -use std::future::Future; - -#[cfg(test)] -mod tests { - use super::*; - use asap_types::enums::{CleanupPolicy, QueryLanguage}; - use query_engine_rust::planner_client::PlannerResult; - use std::sync::atomic::{AtomicBool, Ordering}; - - #[tokio::test] - async fn rejected_precompute_plan_does_not_apply_other_components() { - let applied = Arc::new(AtomicBool::new(false)); - let applied_by_callback = applied.clone(); - let result = PlannerResult { - streaming_config: StreamingConfig::default(), - inference_config: InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::NoCleanup), - punted_queries: Vec::new(), - }; - - let update_result = apply_plan_if_precompute_succeeds( - result, - |_| async { Err::<(), _>("invalid streaming config".into()) }, - move |_| { - applied_by_callback.store(true, Ordering::SeqCst); - }, - ) - .await; - - assert!(update_result.is_err()); - assert!(!applied.load(Ordering::SeqCst)); - } -} - use clap::Parser; use engine_config::{BackendConfig, EngineConfig, IngestConfig}; use figment::{ @@ -73,20 +40,6 @@ struct Args { overrides: Vec, } -async fn apply_plan_if_precompute_succeeds( - result: query_engine_rust::planner_client::PlannerResult, - update_precompute: F, - apply_other_components: impl FnOnce(query_engine_rust::planner_client::PlannerResult), -) -> Result<()> -where - F: FnOnce(&StreamingConfig) -> Fut, - Fut: Future>, -{ - update_precompute(&result.streaming_config).await?; - apply_other_components(result); - Ok(()) -} - #[tokio::main] async fn main() -> Result<()> { let args = Args::parse(); @@ -488,37 +441,23 @@ async fn main() -> Result<()> { } let result = rx.borrow().clone(); if let Some(result) = result { - let apply_result = apply_plan_if_precompute_succeeds( - result, - |streaming_config| { - let streaming_config = streaming_config.clone(); - let precompute_handle = pe_engine_handle.as_ref(); - async move { - if let Some(handle) = precompute_handle { - handle.update_streaming_config(&streaming_config).await - } else { - Ok(()) - } - } - }, - |result| { - engine_for_applier - .update_streaming_config(Arc::new(result.streaming_config.clone())); - engine_for_applier - .update_inference_config(result.inference_config.clone()); - store_for_applier - .update_streaming_config(result.streaming_config.clone()); - *streaming_config_ref_for_applier.write().unwrap() = - Arc::new(result.streaming_config); - *inference_config_ref_for_applier.write().unwrap() = - Arc::new(result.inference_config); - info!("Applier: applied new plan from query tracker"); - }, - ) - .await; - if let Err(e) = apply_result { - warn!("Applier: failed to apply plan: {}", e); + if let Some(ref handle) = pe_engine_handle { + if let Err(e) = handle + .update_streaming_config(&result.streaming_config) + .await + { + warn!("Applier: failed to update precompute engine: {}", e); + } } + engine_for_applier + .update_streaming_config(Arc::new(result.streaming_config.clone())); + engine_for_applier.update_inference_config(result.inference_config.clone()); + store_for_applier.update_streaming_config(result.streaming_config.clone()); + *streaming_config_ref_for_applier.write().unwrap() = + Arc::new(result.streaming_config); + *inference_config_ref_for_applier.write().unwrap() = + Arc::new(result.inference_config); + info!("Applier: applied new plan from query tracker"); } } }); diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index 2fbf2ed..abe882b 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -38,8 +38,6 @@ impl PrecomputeEngineHandle { &self, config: &StreamingConfig, ) -> Result<(), Box> { - validate_cms_aggregation_configs(config)?; - let agg_configs_map: HashMap> = config .get_all_aggregation_configs() .iter() @@ -157,7 +155,7 @@ impl PrecomputeEngine { /// Start the precompute engine. This spawns worker tasks and all registered /// ingest sources, then blocks until shutdown. pub async fn run(mut self) -> Result<(), Box> { - validate_cms_aggregation_configs(&self.streaming_config)?; + validate_startup_aggregation_configs(&self.streaming_config)?; let num_workers = self.config.num_workers; @@ -258,13 +256,12 @@ impl PrecomputeEngine { } } -/// Validate plain Count-Min Sketch configs before starting workers or ingest -/// tasks, and before accepting runtime replacements. +/// Validate aggregation configs before starting any worker or ingest task. /// /// Plain Count-Min Sketch configs carry their SUM-versus-COUNT contract in /// `aggregation_sub_type`; allowing an invalid value to reach the lazy worker /// path would leave the engine running while silently losing that contract. -fn validate_cms_aggregation_configs( +fn validate_startup_aggregation_configs( streaming_config: &StreamingConfig, ) -> Result<(), Box> { for (&aggregation_id, config) in streaming_config.get_all_aggregation_configs() { @@ -349,52 +346,4 @@ mod tests { assert!(err.to_string().contains("aggregation_id 1")); assert!(err.to_string().contains("sum") && err.to_string().contains("count")); } - - #[tokio::test] - async fn runtime_update_rejects_invalid_cms_subtype_without_replacing_config() { - let mut parameters = HashMap::new(); - parameters.insert("depth".to_string(), json!(3_u64)); - parameters.insert("width".to_string(), json!(128_u64)); - let valid = AggregationConfig::new( - 1, - AggregationType::CountMinSketch, - "sum".to_string(), - parameters, - promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), - promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![ - "host".to_string() - ]), - promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), - String::new(), - 1_000, - 1_000, - WindowType::Tumbling, - "requests_total".to_string(), - "requests_total".to_string(), - None, - None, - None, - None, - ); - let engine = PrecomputeEngine::new( - PrecomputeEngineConfig::default(), - Arc::new(StreamingConfig::new(HashMap::from([(1, valid.clone())]))), - Arc::new(NoopOutputSink::new()), - vec![], - ); - let handle = engine.handle(); - - let mut invalid = valid; - invalid.aggregation_sub_type.clear(); - let result = handle - .update_streaming_config(&StreamingConfig::new(HashMap::from([(1, invalid)]))) - .await; - - let err = result.expect_err("invalid runtime CMS subtype must be rejected"); - assert!(err.to_string().contains("aggregation_id 1")); - assert!(err.to_string().contains("sum") && err.to_string().contains("count")); - let current = handle.ingest_agg_configs.load(); - assert_eq!(current.len(), 1); - assert_eq!(current[0].aggregation_sub_type, "sum"); - } } diff --git a/asap-query-engine/src/precompute_engine/worker.rs b/asap-query-engine/src/precompute_engine/worker.rs index bebc288..b59a14e 100644 --- a/asap-query-engine/src/precompute_engine/worker.rs +++ b/asap-query-engine/src/precompute_engine/worker.rs @@ -240,26 +240,21 @@ impl Worker { break; } WorkerMessage::UpdateAggConfigs(new_configs) => { - // Flush and evict group states for removed or materially changed - // aggregations. Must happen before swapping agg_configs so - // GroupState.config is still valid during the final window close. - let reset_ids: Vec = self + // Flush and evict group states for agg IDs that are being removed. + // Must happen before swapping agg_configs so GroupState.config is + // still valid during the final window close. + let removed_ids: Vec = self .agg_configs - .iter() - .filter_map(|(&id, old_config)| match new_configs.get(&id) { - None => Some(id), - Some(new_config) if old_config.as_ref() != new_config.as_ref() => { - Some(id) - } - Some(_) => None, - }) + .keys() + .filter(|id| !new_configs.contains_key(id)) + .copied() .collect(); - if !reset_ids.is_empty() { + if !removed_ids.is_empty() { let mut emit_batch: Vec<(PrecomputedOutput, Box)> = Vec::new(); - for agg_id in &reset_ids { + for agg_id in &removed_ids { // Drain all group states for this agg_id in one move. let Some(inner) = self.group_states.remove(agg_id) else { continue; @@ -312,8 +307,8 @@ impl Worker { if !emit_batch.is_empty() { if let Err(e) = self.output_sink.emit_batch(emit_batch) { warn!( - "Worker {}: error flushing reset agg_ids {:?}: {}", - self.id, reset_ids, e + "Worker {}: error flushing removed agg_ids {:?}: {}", + self.id, removed_ids, e ); } } @@ -321,10 +316,10 @@ impl Worker { self.group_count .store(self.total_groups(), Ordering::Relaxed); info!( - "Worker {}: evicted {} reset agg_id(s) {:?}", + "Worker {}: evicted {} removed agg_id(s) {:?}", self.id, - reset_ids.len(), - reset_ids, + removed_ids.len(), + removed_ids, ); } @@ -2818,97 +2813,6 @@ aggregations: ); } - #[tokio::test] - async fn test_update_agg_configs_recreates_state_when_cms_subtype_changes() { - let mut sum_config = make_agg_config( - 1, - "requests_total", - AggregationType::CountMinSketch, - "sum", - 1_000, - 1_000, - vec![], - ); - sum_config - .parameters - .insert("depth".to_string(), serde_json::json!(3_u64)); - sum_config - .parameters - .insert("width".to_string(), serde_json::json!(128_u64)); - let mut count_config = sum_config.clone(); - count_config.aggregation_sub_type = "count".to_string(); - - let sink = Arc::new(CapturingOutputSink::new()); - let (tx, rx) = tokio::sync::mpsc::channel(32); - let wm = Arc::new(AtomicI64::new(i64::MIN)); - let worker = Worker::new( - 0, - rx, - sink.clone(), - arc_configs(HashMap::from([(1, sum_config)])), - WorkerRuntimeConfig { - max_buffer_per_series: 10_000, - allowed_lateness_ms: 0, - pass_raw_samples: false, - raw_mode_aggregation_id: 0, - late_data_policy: LateDataPolicy::Drop, - wall_clock_grace_period_ms: 0, - }, - Arc::new(AtomicUsize::new(0)), - wm.clone(), - vec![wm], - ); - let handle = tokio::spawn(async move { worker.run().await }); - - tx.send(WorkerMessage::GroupSamples { - agg_id: 1, - group_key: String::new(), - samples: vec![("requests_total".to_string(), 100, 10.0)], - ingest_received_at: std::time::Instant::now(), - }) - .await - .unwrap(); - tx.send(WorkerMessage::UpdateAggConfigs(arc_configs(HashMap::from( - [(1, count_config)], - )))) - .await - .unwrap(); - tx.send(WorkerMessage::GroupSamples { - agg_id: 1, - group_key: String::new(), - samples: vec![("requests_total".to_string(), 1_100, 10.0)], - ingest_received_at: std::time::Instant::now(), - }) - .await - .unwrap(); - tx.send(WorkerMessage::Shutdown).await.unwrap(); - handle.await.unwrap(); - - let mut captured = sink.drain(); - captured.sort_by_key(|(output, _)| output.start_timestamp); - assert_eq!(captured.len(), 2); - - let first = captured[0] - .1 - .as_any() - .downcast_ref::() - .expect("first output should be CMS"); - assert_eq!( - first.query_key(&KeyByLabelValues::new_with_labels(vec![])), - 10.0 - ); - - let second = captured[1] - .1 - .as_any() - .downcast_ref::() - .expect("second output should be CMS"); - assert_eq!( - second.query_key(&KeyByLabelValues::new_with_labels(vec![])), - 1.0 - ); - } - // ----------------------------------------------------------------------- // Test: removing an agg_id force-closes its open windows with a finite // bound — at realistic (epoch-ms) timestamps. From 69d6ac3147c826dcf223dcf4b4e1cde7e53c0926 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 2 Sep 2026 09:54:58 -0400 Subject: [PATCH 5/6] fix(query-engine): validate heap CMS configuration --- .../precompute_engine/accumulator_factory.rs | 71 ++++++++++++++++--- .../src/precompute_engine/engine.rs | 58 +++++++++++++-- 2 files changed, 117 insertions(+), 12 deletions(-) diff --git a/asap-query-engine/src/precompute_engine/accumulator_factory.rs b/asap-query-engine/src/precompute_engine/accumulator_factory.rs index c53f74e..569e87f 100644 --- a/asap-query-engine/src/precompute_engine/accumulator_factory.rs +++ b/asap-query-engine/src/precompute_engine/accumulator_factory.rs @@ -753,6 +753,17 @@ fn cms_count_events_for_sub_type(sub_type: &str) -> Result { } } +/// Validate the aggregation subtype for a heap-backed Count-Min Sketch. +fn validate_cms_with_heap_sub_type(sub_type: &str) -> Result<(), String> { + if sub_type.eq_ignore_ascii_case("topk") { + Ok(()) + } else { + Err(format!( + "CountMinSketchWithHeap requires aggregation_sub_type 'topk', got '{sub_type}'" + )) + } +} + /// Extract `(row_num, col_num, k)` for HydraKLL configs. fn hydra_kll_params(config: &AggregationConfig) -> Result<(usize, usize, u16), String> { let (row_num, col_num) = cms_params(config)?; @@ -787,12 +798,15 @@ fn cms_heap_params(config: &AggregationConfig) -> Result<(usize, usize, usize), /// Whether a CountMinSketchWithHeap config should count events (weight 1 per /// observation, COUNT semantics) rather than summing the sample value. /// Defaults to `true` so `COUNT(...)` top-k works out of the box. -fn cms_count_events(config: &AggregationConfig) -> bool { - config - .parameters - .get("count_events") - .and_then(|v| v.as_bool()) - .unwrap_or(true) +fn cms_count_events(config: &AggregationConfig) -> Result { + match config.parameters.get("count_events") { + None => Ok(true), + Some(value) => value.as_bool().ok_or_else(|| { + format!( + "CountMinSketchWithHeap parameter 'count_events' must be a boolean, got {value}" + ) + }), + } } /// Extract the HLL `precision` parameter from a config. Falls back to @@ -898,12 +912,13 @@ pub fn create_accumulator_updater( ))) } AggregationType::CountMinSketchWithHeap => { + validate_cms_with_heap_sub_type(sub_type)?; let (row_num, col_num, heap_size) = cms_heap_params(config)?; Ok(Box::new(CmsWithHeapAccumulatorUpdater::new( row_num, col_num, heap_size, - cms_count_events(config), + cms_count_events(config)?, ))) } AggregationType::HydraKLL => { @@ -1559,6 +1574,43 @@ mod tests { ) } + #[test] + fn test_cms_with_heap_rejects_empty_subtype() { + let mut config = cms_heap_config(cms_heap_params_required()); + config.aggregation_sub_type.clear(); + + let err = match create_accumulator_updater(&config) { + Ok(_) => panic!("empty CountMinSketchWithHeap subtype must fail"), + Err(err) => err, + }; + assert!(err.contains("topk")); + } + + #[test] + fn test_cms_with_heap_rejects_unknown_subtype() { + let mut config = cms_heap_config(cms_heap_params_required()); + config.aggregation_sub_type = "count".to_string(); + + let err = match create_accumulator_updater(&config) { + Ok(_) => panic!("unknown CountMinSketchWithHeap subtype must fail"), + Err(err) => err, + }; + assert!(err.contains("count")); + } + + #[test] + fn test_cms_with_heap_rejects_non_boolean_count_events() { + let mut parameters = cms_heap_params_required(); + parameters.insert("count_events".to_string(), serde_json::json!("true")); + let config = cms_heap_config(parameters); + + let err = match create_accumulator_updater(&config) { + Ok(_) => panic!("non-boolean count_events must fail"), + Err(err) => err, + }; + assert!(err.contains("count_events") && err.contains("boolean")); + } + #[test] fn test_cms_with_heap_factory_routes_to_heap_accumulator_and_is_keyed() { // CountMinSketchWithHeap must build a CmsWithHeapAccumulatorUpdater whose @@ -1631,7 +1683,10 @@ mod tests { params.insert("heapsize".to_string(), serde_json::json!(40)); let config = cms_heap_config(params); assert_eq!(cms_heap_params(&config).unwrap(), (4, 2048, 40)); - assert!(cms_count_events(&config), "count_events defaults to true"); + assert!( + cms_count_events(&config).unwrap(), + "count_events defaults to true" + ); } #[test] diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index abe882b..eeab16f 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -258,14 +258,17 @@ impl PrecomputeEngine { /// Validate aggregation configs before starting any worker or ingest task. /// -/// Plain Count-Min Sketch configs carry their SUM-versus-COUNT contract in -/// `aggregation_sub_type`; allowing an invalid value to reach the lazy worker -/// path would leave the engine running while silently losing that contract. +/// Count-Min Sketch configs carry their semantic contract in subtype fields or +/// parameters; allowing an invalid value to reach the lazy worker path would +/// leave the engine running while silently losing that contract. fn validate_startup_aggregation_configs( streaming_config: &StreamingConfig, ) -> Result<(), Box> { for (&aggregation_id, config) in streaming_config.get_all_aggregation_configs() { - if config.aggregation_type != AggregationType::CountMinSketch { + if !matches!( + config.aggregation_type, + AggregationType::CountMinSketch | AggregationType::CountMinSketchWithHeap + ) { continue; } @@ -346,4 +349,51 @@ mod tests { assert!(err.to_string().contains("aggregation_id 1")); assert!(err.to_string().contains("sum") && err.to_string().contains("count")); } + + #[tokio::test] + async fn run_rejects_invalid_cms_with_heap_subtype_before_starting_workers() { + let mut parameters = HashMap::new(); + parameters.insert("depth".to_string(), json!(3_u64)); + parameters.insert("width".to_string(), json!(128_u64)); + parameters.insert("heapsize".to_string(), json!(32_u64)); + let cms = AggregationConfig::new( + 1, + AggregationType::CountMinSketchWithHeap, + String::new(), + parameters, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![ + "host".to_string() + ]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + 1_000, + 1_000, + WindowType::Tumbling, + "requests_total".to_string(), + "requests_total".to_string(), + None, + None, + None, + None, + ); + let engine = PrecomputeEngine::new( + PrecomputeEngineConfig { + num_workers: 1, + late_data_policy: LateDataPolicy::Drop, + ..PrecomputeEngineConfig::default() + }, + Arc::new(StreamingConfig::new(HashMap::from([(1, cms)]))), + Arc::new(NoopOutputSink::new()), + vec![Box::new(ShutdownSource)], + ); + + let result = engine.run().await; + let err = match result { + Ok(()) => panic!("invalid heap CMS subtype must fail before startup"), + Err(err) => err, + }; + assert!(err.to_string().contains("aggregation_id 1")); + assert!(err.to_string().contains("topk")); + } } From e770f082d76e427b411d1fb4cea5e106e933bd13 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Wed, 2 Sep 2026 10:46:06 -0400 Subject: [PATCH 6/6] fix(query-engine): make CMS startup validation deterministic --- .../rs/asap_types/src/capability_matching.rs | 39 ++++++- .../src/precompute_engine/engine.rs | 104 +++++++++++++++++- 2 files changed, 133 insertions(+), 10 deletions(-) diff --git a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs index 30218da..215e5be 100644 --- a/asap-common/dependencies/rs/asap_types/src/capability_matching.rs +++ b/asap-common/dependencies/rs/asap_types/src/capability_matching.rs @@ -198,11 +198,14 @@ fn plain_cms_sub_type_compatible(stat: Statistic, config: &AggregationConfig) -> return true; } - match stat { - Statistic::Sum => config.aggregation_sub_type.eq_ignore_ascii_case("sum"), - Statistic::Count => config.aggregation_sub_type.eq_ignore_ascii_case("count"), - _ => false, - } + let expected_sub_type = match stat { + Statistic::Sum => "sum", + Statistic::Count => "count", + _ => unreachable!("plain CMS matching only supports SUM and COUNT"), + }; + config + .aggregation_sub_type + .eq_ignore_ascii_case(expected_sub_type) } /// Aggregation priority comparator: prefer larger `window_size_ms` (descending). @@ -541,6 +544,32 @@ mod tests { assert_eq!(count.aggregation_id_for_value, 2); } + #[test] + fn plain_cms_with_invalid_subtypes_is_excluded_from_matching() { + for invalid_sub_type in ["", "unknown", " sum "] { + let configs = single_config(make_config( + 1, + "cpu", + "CountMinSketch", + invalid_sub_type, + 300_000, + "tumbling", + &[], + "", + )); + + let result = find_compatible_aggregation( + &configs, + &req("cpu", &[Statistic::Sum], 300_000, &[], ""), + ); + + assert!( + result.is_none(), + "invalid plain CMS subtype {invalid_sub_type:?} must not match SUM" + ); + } + } + #[test] fn quantile_any_value_finds_kll() { let configs = single_config(make_config( diff --git a/asap-query-engine/src/precompute_engine/engine.rs b/asap-query-engine/src/precompute_engine/engine.rs index eeab16f..e3818c5 100644 --- a/asap-query-engine/src/precompute_engine/engine.rs +++ b/asap-query-engine/src/precompute_engine/engine.rs @@ -264,6 +264,8 @@ impl PrecomputeEngine { fn validate_startup_aggregation_configs( streaming_config: &StreamingConfig, ) -> Result<(), Box> { + let mut errors = Vec::new(); + for (&aggregation_id, config) in streaming_config.get_all_aggregation_configs() { if !matches!( config.aggregation_type, @@ -272,13 +274,23 @@ fn validate_startup_aggregation_configs( continue; } - create_accumulator_updater(config).map(|_| ()).map_err( - |err| -> Box { + if let Err(err) = create_accumulator_updater(config) { + errors.push((aggregation_id, err)); + } + } + + if !errors.is_empty() { + errors.sort_by_key(|(aggregation_id, _)| *aggregation_id); + let details = errors + .into_iter() + .map(|(aggregation_id, err)| { format!("invalid aggregation config for aggregation_id {aggregation_id}: {err}") - .into() - }, - )?; + }) + .collect::>() + .join("; "); + return Err(details.into()); } + Ok(()) } @@ -396,4 +408,86 @@ mod tests { assert!(err.to_string().contains("aggregation_id 1")); assert!(err.to_string().contains("topk")); } + + #[tokio::test] + async fn run_reports_all_invalid_cms_configs_in_aggregation_id_order() { + let cms_config = |id, aggregation_type, aggregation_sub_type, parameters| { + AggregationConfig::new( + id, + aggregation_type, + aggregation_sub_type, + parameters, + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![ + "host".to_string(), + ]), + promql_utilities::data_model::key_by_label_names::KeyByLabelNames::new(vec![]), + String::new(), + 1_000, + 1_000, + WindowType::Tumbling, + "requests_total".to_string(), + "requests_total".to_string(), + None, + None, + None, + None, + ) + }; + + let mut heap_parameters = HashMap::new(); + heap_parameters.insert("depth".to_string(), json!(3_u64)); + heap_parameters.insert("width".to_string(), json!(128_u64)); + heap_parameters.insert("heapsize".to_string(), json!(32_u64)); + + let configs = HashMap::from([ + ( + 20, + cms_config( + 20, + AggregationType::CountMinSketchWithHeap, + String::new(), + heap_parameters, + ), + ), + ( + 3, + cms_config( + 3, + AggregationType::CountMinSketch, + String::new(), + HashMap::from([ + ("depth".to_string(), json!(3_u64)), + ("width".to_string(), json!(128_u64)), + ]), + ), + ), + ]); + let engine = PrecomputeEngine::new( + PrecomputeEngineConfig { + num_workers: 1, + late_data_policy: LateDataPolicy::Drop, + ..PrecomputeEngineConfig::default() + }, + Arc::new(StreamingConfig::new(configs)), + Arc::new(NoopOutputSink::new()), + vec![Box::new(ShutdownSource)], + ); + + let result = engine.run().await; + let err = match result { + Ok(()) => panic!("invalid CMS configs must fail before startup"), + Err(err) => err.to_string(), + }; + let id_3 = err + .find("aggregation_id 3") + .expect("CMS error should be reported"); + let id_20 = err + .find("aggregation_id 20") + .expect("heap CMS error should be reported"); + assert!( + id_3 < id_20, + "errors should be ordered by aggregation ID: {err}" + ); + } }