fix(query-engine): invalidate changed aggregation state

Author: milindsrivastava1997

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,

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"],

asap-query-engine/src/precompute_engine/engine.rs

@@ -38,6 +38,8 @@ impl PrecomputeEngineHandle {
         &self,
         config: &StreamingConfig,
     ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+        validate_cms_aggregation_configs(config)?;
+
         let agg_configs_map: HashMap<u64, Arc<AggregationConfig>> = 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<dyn std::error::Error + Send + Sync>> {
-        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<dyn std::error::Error + Send + Sync>> {
     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");
+    }
 }