From 69b8786c1824d2cbe8f7e5837a07502d3582d81b Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Sun, 23 Aug 2026 17:27:05 -0400 Subject: [PATCH] fix(optimizer): enforce aggregation_id assignment via OptimizerSolution::register_config 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 --- .../src/optimizer/candidate_gen.rs | 6 +- asap-planner-rs/src/optimizer/greedy.rs | 37 ++----- asap-planner-rs/src/optimizer/solution.rs | 99 +++++++++++++++++-- asap-planner-rs/src/optimizer/translator.rs | 4 +- 4 files changed, 104 insertions(+), 42 deletions(-) diff --git a/asap-planner-rs/src/optimizer/candidate_gen.rs b/asap-planner-rs/src/optimizer/candidate_gen.rs index db62b30..b3b023c 100644 --- a/asap-planner-rs/src/optimizer/candidate_gen.rs +++ b/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(), diff --git a/asap-planner-rs/src/optimizer/greedy.rs b/asap-planner-rs/src/optimizer/greedy.rs index 3275a32..086bfd7 100644 --- a/asap-planner-rs/src/optimizer/greedy.rs +++ b/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 = 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,10 +70,10 @@ 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, @@ -97,12 +81,7 @@ pub fn greedy_assign( }); } - 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 = 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()); } } diff --git a/asap-planner-rs/src/optimizer/solution.rs b/asap-planner-rs/src/optimizer/solution.rs index 646a458..3451622 100644 --- a/asap-planner-rs/src/optimizer/solution.rs +++ b/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, + /// + /// 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, + + /// Next id `register_config` will hand out. + next_id: u64, /// One entry per deduplicated AQE across the full RQE workload. pub assignments: Vec, @@ -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 { + &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) -> 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); + } +} diff --git a/asap-planner-rs/src/optimizer/translator.rs b/asap-planner-rs/src/optimizer/translator.rs index 453d096..8f5a07a 100644 --- a/asap-planner-rs/src/optimizer/translator.rs +++ b/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(), }