added logic to remove stale states from precompute engine after reconfiguration

Author: milindsrivastava1997

.design_docs/issue-242-design-decisions.md

@@ -105,9 +105,30 @@ This means: if a metric was being precomputed and the planner decides it still s
 
 ---
 
-## Decision 8: Stale precomputed data after config replace
+## Decision 8: Stale precomputed data and in-flight state after config replace
 
-**Question:** When `streaming_config` is replaced, the `SimpleMapStore` may contain precomputed sketch data for aggregations that are no longer in the new config. Should this data be purged?
+**Question:** When `streaming_config` is replaced, how is in-flight precompute data handled, and what happens to stale state in both the store and the workers?
+
+### In-flight data in precompute engine workers
+
+`UpdateAggConfigs` replaces each worker's `agg_configs` map but **does not touch `group_states`**. Existing `GroupState` entries (keyed by `(agg_id, group_key)`) for old aggregation IDs are left in the worker's HashMap.
+
+Consequences:
+- **Data is not silently dropped.** The flush loop iterates over *all* `group_states`, so old groups still close their open windows and emit accumulated sketch data to the `SimpleMapStore` — under the old aggregation IDs.
+- **No new samples reach old group states.** The ingest router only routes samples to agg IDs present in the new `agg_configs`, so old groups stop receiving input but still flush whatever they have accumulated.
+- **Memory is never reclaimed.** Old `GroupState` entries are never evicted from the worker's HashMap. For the one-shot first-plan scenario this is a fixed-size overhead. For repeated reconfigurations this is a latent memory leak — stale group states accumulate in worker memory until process restart.
+
+The data committed to the store under old agg IDs becomes **orphaned**: the query engine now uses new inference/streaming configs pointing to new agg IDs, so those store entries are never queried and eventually expire via the cleanup policy.
+
+**Implementation:** On `UpdateAggConfigs`, workers now prune `group_states` entries whose `agg_id` is absent from the new config, after a final forced flush. Specifically:
+1. Identify removed agg IDs (old `agg_configs` keys not present in `new_configs`).
+2. For each removed group state, run the window-close logic with `effective_wm = i64::MAX` — since no new samples will arrive, all open windows are treated as due regardless of watermark.
+3. Emit accumulated data from those windows to the store, then remove the entry from `group_states` and update `group_count`.
+4. Swap in `new_configs`.
+
+The final flush happens **before** `agg_configs` is swapped so that `GroupState.config` (an `Arc<AggregationConfig>`) is still valid during window-bound calculations. The emitted data is still written under old agg IDs (and will not be queried), but this avoids silently discarding partially-accumulated windows and correctly frees the worker memory.
+
+### Stale data already committed to `SimpleMapStore`
 
 **Decision:** Let it expire naturally via the existing cleanup policy. No active purge.
 

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

@@ -170,6 +170,85 @@ 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<u64> = self
+                        .agg_configs
+                        .keys()
+                        .filter(|id| !new_configs.contains_key(id))
+                        .copied()
+                        .collect();
+
+                    if !removed_ids.is_empty() {
+                        let mut emit_batch: Vec<(PrecomputedOutput, Box<dyn AggregateCore>)> =
+                            Vec::new();
+
+                        for agg_id in &removed_ids {
+                            // Drain all group states for this agg_id.
+                            let removed_keys: Vec<_> = self
+                                .group_states
+                                .keys()
+                                .filter(|(id, _)| id == agg_id)
+                                .cloned()
+                                .collect();
+
+                            for key in removed_keys {
+                                let Some(state) = self.group_states.remove(&key) else {
+                                    continue;
+                                };
+                                if state.previous_watermark_ms == i64::MIN {
+                                    continue; // No samples received — nothing to emit.
+                                }
+                                // Force-close all open windows by advancing the watermark
+                                // to i64::MAX. No new samples will arrive for this group.
+                                let (group_key_str, mut active_panes) =
+                                    (key.1.clone(), state.active_panes);
+                                let closed = state
+                                    .window_manager
+                                    .closed_windows(state.previous_watermark_ms, i64::MAX);
+
+                                for window_start in &closed {
+                                    let (_, window_end) =
+                                        state.window_manager.window_bounds(*window_start);
+                                    let pane_starts =
+                                        state.window_manager.panes_for_window(*window_start);
+                                    if let Some(accumulator) =
+                                        merge_panes_for_window(&mut active_panes, &pane_starts)
+                                    {
+                                        let group_key_lv =
+                                            build_group_key_label_values(&group_key_str);
+                                        let output = PrecomputedOutput::new(
+                                            *window_start as u64,
+                                            window_end as u64,
+                                            Some(group_key_lv),
+                                            *agg_id,
+                                        );
+                                        emit_batch.push((output, accumulator));
+                                    }
+                                }
+                            }
+                        }
+
+                        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
+                                );
+                            }
+                        }
+
+                        self.group_count
+                            .store(self.group_states.len(), Ordering::Relaxed);
+                        info!(
+                            "Worker {}: evicted {} removed agg_id(s) {:?}",
+                            self.id,
+                            removed_ids.len(),
+                            removed_ids,
+                        );
+                    }
+
                     let added = new_configs.len().saturating_sub(self.agg_configs.len());
                     self.agg_configs = new_configs;
                     info!(