fix(planner): bound MILP retention with CircularBuffer and spill evicted windows

asap-common/dependencies/rs/asap_types/src/streaming_config.rs

@@ -50,20 +50,33 @@ impl StreamingConfig {
         data: &Value,
         inference_config: Option<&InferenceConfig>,
     ) -> Result<Self> {
-        let mut retention_map: HashMap<u64, u64> = HashMap::new();
+        // None = unbounded, which wins over any bounded count from another reference.
+        let mut retention_map: HashMap<u64, Option<u64>> = HashMap::new();
         let mut read_count_threshold_map: HashMap<u64, u64> = 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}");
+    }
 }

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<String>),
     #[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<RollupBlocker>,
 ) -> Result<QueryPlanEntry, MilpOutputError> {
-    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 {

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.

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<String>,
+    /// 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#"

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(),