fix(optimizer): enforce aggregation_id assignment via OptimizerSolution::register_config (#585)

candidate_gen.rs builds candidates with aggregation_id=0 as a placeholder for
cost evaluation, since the id isn't known until a solver deploys the config.
greedy.rs correctly overwrote it before insertion, but nothing stopped a
future solver (or a change to greedy.rs) from forgetting to.

Make deployed_configs private on OptimizerSolution and route all insertion
through register_config(), which always assigns a fresh id. No caller can
add a deployed config without going through id assignment.

Fixes #564

Author: milindsrivastava1997

asap-planner-rs/src/optimizer/candidate_gen.rs

@@ -172,7 +172,9 @@ fn determine_query_method(
     }
 }
 
-/// Build an AggregationConfig from candidate parameters. aggregation_id = 0 (placeholder).
+/// Build an AggregationConfig from candidate parameters. aggregation_id = 0 is a
+/// placeholder never used past cost evaluation — OptimizerSolution::register_config
+/// overwrites it with a real id when (if) a solver deploys this candidate.
 #[allow(clippy::too_many_arguments)]
 fn build_config(
     aqe: &AQE,
@@ -185,7 +187,7 @@ fn build_config(
     n_windows: u64,
 ) -> AggregationConfig {
     AggregationConfig::new(
-        0, // placeholder; replaced by greedy/MIP solver when deploying
+        0, // placeholder; overwritten by OptimizerSolution::register_config when deployed
         agg_type,
         sub_type.to_string(),
         params.clone(),

asap-planner-rs/src/optimizer/greedy.rs

@@ -1,6 +1,3 @@
-use std::collections::HashMap;
-
-use asap_types::aggregation_config::AggregationConfig;
 use tracing::debug;
 
 use super::atomic_costs::{resolve_atomic_costs, AtomicCostTable};
@@ -29,11 +26,7 @@ pub fn greedy_assign(
     atomic_cost_table: &AtomicCostTable,
     weights: &CostWeights,
 ) -> OptimizerSolution {
-    let mut deployed_configs: HashMap<u64, AggregationConfig> = HashMap::new();
-    let mut assignments = Vec::new();
-    let mut estimated_ingest_cost_per_sec = 0.0;
-    let mut estimated_total_cost_per_sec = 0.0;
-    let mut next_id: u64 = 1;
+    let mut solution = OptimizerSolution::empty();
 
     for aqe in aqes {
         let candidates = enumerate_candidates(&aqe, scrape_interval_ms);
@@ -66,16 +59,7 @@ pub fn greedy_assign(
         let query_rate = aqe.query_frequency_hz * query_cost(&aqe, &best, &costs, weights);
         let query_method = best.query_method.clone();
 
-        let aggregation_id = match best.config {
-            None => None,
-            Some(mut config) => {
-                let id = next_id;
-                next_id += 1;
-                config.aggregation_id = id;
-                deployed_configs.insert(id, config);
-                Some(id)
-            }
-        };
+        let aggregation_id = best.config.map(|config| solution.register_config(config));
 
         debug!(
             metric = %aqe.requirements.metric,
@@ -86,23 +70,18 @@ pub fn greedy_assign(
             "greedy: assigned AQE"
         );
 
-        estimated_ingest_cost_per_sec += ingest;
-        estimated_total_cost_per_sec += ingest + query_rate;
+        solution.estimated_ingest_cost_per_sec += ingest;
+        solution.estimated_total_cost_per_sec += ingest + query_rate;
 
-        assignments.push(AQEAssignment {
+        solution.assignments.push(AQEAssignment {
             aqe,
             aggregation_id,
             query_method,
             estimated_query_cost_per_sec: query_rate,
         });
     }
 
-    OptimizerSolution {
-        deployed_configs,
-        assignments,
-        estimated_ingest_cost_per_sec,
-        estimated_total_cost_per_sec,
-    }
+    solution
 }
 
 #[cfg(test)]
@@ -145,7 +124,7 @@ mod tests {
         );
 
         let mut seen_ids: StdHashMap<u64, ()> = StdHashMap::new();
-        for id in solution.deployed_configs.keys() {
+        for id in solution.deployed_configs().keys() {
             assert!(
                 seen_ids.insert(*id, ()).is_none(),
                 "duplicate aggregation_id"
@@ -178,6 +157,6 @@ mod tests {
             &CostWeights::default(),
         );
         assert_eq!(solution.num_exact_fallback(), 1);
-        assert!(solution.deployed_configs.is_empty());
+        assert!(solution.deployed_configs().is_empty());
     }
 }

asap-planner-rs/src/optimizer/solution.rs

@@ -86,7 +86,14 @@ pub struct AQEAssignment {
 pub struct OptimizerSolution {
     /// Deployed streaming configs (y_g = 1 in the MIP). Keyed by aggregation_id.
     /// Empty for all-EXACT solutions (Phase 1 scaffolding).
-    pub deployed_configs: HashMap<u64, AggregationConfig>,
+    ///
+    /// Private: the only way to add an entry is `register_config`, which
+    /// assigns the id. This keeps candidate_gen.rs's placeholder id (0) from
+    /// ever reaching a deployed config — see ASAPQuery#564.
+    deployed_configs: HashMap<u64, AggregationConfig>,
+
+    /// Next id `register_config` will hand out.
+    next_id: u64,
 
     /// One entry per deduplicated AQE across the full RQE workload.
     pub assignments: Vec<AQEAssignment>,
@@ -100,10 +107,38 @@ pub struct OptimizerSolution {
 }
 
 impl OptimizerSolution {
+    /// An empty solution with no assignments or deployed configs yet, ready
+    /// to be built up incrementally (e.g. by a solver's assignment loop).
+    pub fn empty() -> Self {
+        Self {
+            deployed_configs: HashMap::new(),
+            next_id: 1,
+            assignments: Vec::new(),
+            estimated_ingest_cost_per_sec: 0.0,
+            estimated_total_cost_per_sec: 0.0,
+        }
+    }
+
+    /// Register a candidate config as deployed: assigns it a fresh unique id
+    /// (overwriting whatever placeholder candidate_gen.rs set), stores it,
+    /// and returns the id. The only way to populate `deployed_configs`.
+    pub fn register_config(&mut self, mut config: AggregationConfig) -> u64 {
+        let id = self.next_id;
+        self.next_id += 1;
+        config.aggregation_id = id;
+        self.deployed_configs.insert(id, config);
+        id
+    }
+
+    pub fn deployed_configs(&self) -> &HashMap<u64, AggregationConfig> {
+        &self.deployed_configs
+    }
+
     /// Construct an all-EXACT solution: every AQE falls back to raw data,
     /// no streaming configs are deployed. Used as the Phase 1 scaffolding baseline.
     pub fn all_exact(aqes: Vec<AQE>) -> Self {
-        let assignments = aqes
+        let mut solution = Self::empty();
+        solution.assignments = aqes
             .into_iter()
             .map(|aqe| AQEAssignment {
                 aqe,
@@ -112,13 +147,7 @@ impl OptimizerSolution {
                 estimated_query_cost_per_sec: 0.0,
             })
             .collect();
-
-        Self {
-            deployed_configs: HashMap::new(),
-            assignments,
-            estimated_ingest_cost_per_sec: 0.0,
-            estimated_total_cost_per_sec: 0.0,
-        }
+        solution
     }
 
     /// Number of AQEs served by an approximate sketch (not EXACT fallback).
@@ -137,3 +166,55 @@ impl OptimizerSolution {
             .count()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use asap_types::enums::WindowType;
+    use promql_utilities::data_model::KeyByLabelNames;
+    use promql_utilities::query_logics::enums::AggregationType;
+
+    fn candidate_config() -> AggregationConfig {
+        // aggregation_id: 0, matching candidate_gen.rs's placeholder — the
+        // thing register_config must always overwrite (ASAPQuery#564).
+        AggregationConfig::new(
+            0,
+            AggregationType::CountMinSketch,
+            "sum".into(),
+            HashMap::new(),
+            KeyByLabelNames::empty(),
+            KeyByLabelNames::empty(),
+            KeyByLabelNames::empty(),
+            String::new(),
+            60_000,
+            60_000,
+            WindowType::Tumbling,
+            String::new(),
+            "test_metric".into(),
+            Some(1),
+            None,
+            None,
+            None,
+        )
+    }
+
+    #[test]
+    fn register_config_never_leaves_the_placeholder_id() {
+        let mut solution = OptimizerSolution::empty();
+        let id = solution.register_config(candidate_config());
+        assert_ne!(
+            id, 0,
+            "register_config must not hand out the placeholder id"
+        );
+        assert_eq!(solution.deployed_configs()[&id].aggregation_id, id);
+    }
+
+    #[test]
+    fn register_config_assigns_distinct_ids() {
+        let mut solution = OptimizerSolution::empty();
+        let id1 = solution.register_config(candidate_config());
+        let id2 = solution.register_config(candidate_config());
+        assert_ne!(id1, id2);
+        assert_eq!(solution.deployed_configs().len(), 2);
+    }
+}

asap-planner-rs/src/optimizer/translator.rs

@@ -19,7 +19,7 @@ pub fn translate(solution: &OptimizerSolution) -> (StreamingConfig, InferenceCon
 
 fn build_streaming_config(solution: &OptimizerSolution) -> StreamingConfig {
     // Deployed configs map directly to AggregationConfigs — the types are the same.
-    StreamingConfig::new(solution.deployed_configs.clone())
+    StreamingConfig::new(solution.deployed_configs().clone())
 }
 
 fn build_inference_config(solution: &OptimizerSolution) -> InferenceConfig {
@@ -74,7 +74,7 @@ pub struct TranslationSummary {
 impl TranslationSummary {
     pub fn from_solution(solution: &OptimizerSolution) -> Self {
         Self {
-            num_deployed_configs: solution.deployed_configs.len(),
+            num_deployed_configs: solution.deployed_configs().len(),
             num_sketch_assignments: solution.num_sketch_served(),
             num_exact_fallbacks: solution.num_exact_fallback(),
         }