diff --git a/asap-query-engine/src/stores/sketch_db/simple_map_store/global.rs b/asap-query-engine/src/stores/sketch_db/simple_map_store/global.rs index adabdf5a..ac551e00 100644 --- a/asap-query-engine/src/stores/sketch_db/simple_map_store/global.rs +++ b/asap-query-engine/src/stores/sketch_db/simple_map_store/global.rs @@ -695,4 +695,34 @@ impl Store for SimpleMapStoreGlobal { info!("SimpleMapStoreGlobal closed"); Ok(()) } + + fn drop_agg_id(&self, agg_id: u64) -> StoreResult { + let mut data = self.lock.lock().unwrap(); + // Count windows before eviction so callers get an accurate + // "records dropped" number, useful for audit logging. + let evicted = data + .stores + .get(&agg_id) + .map(|per_key| { + per_key.current_epoch.len() + + per_key + .sealed_epochs + .values() + .map(|e| e.entries.len()) + .sum::() + }) + .unwrap_or(0); + data.stores.remove(&agg_id); + data.read_counts.remove(&agg_id); + data.earliest_timestamp_per_aggregation_id.remove(&agg_id); + // `metrics` and `items_inserted` are keyed by metric name, + // not agg_id, so we don't touch them — other agg_ids for the + // same metric (e.g. historical schemas) may still exist. + info!( + agg_id, + evicted_windows = evicted, + "SimpleMapStoreGlobal::drop_agg_id" + ); + Ok(evicted) + } } diff --git a/asap-query-engine/src/stores/sketch_db/simple_map_store/mod.rs b/asap-query-engine/src/stores/sketch_db/simple_map_store/mod.rs index 5caab2cf..3319a991 100644 --- a/asap-query-engine/src/stores/sketch_db/simple_map_store/mod.rs +++ b/asap-query-engine/src/stores/sketch_db/simple_map_store/mod.rs @@ -161,4 +161,158 @@ impl Store for SimpleMapStore { SimpleMapStore::PerKey(store) => store.close(), } } + + fn drop_agg_id(&self, agg_id: u64) -> StoreResult { + match self { + SimpleMapStore::Global(store) => store.drop_agg_id(agg_id), + SimpleMapStore::PerKey(store) => store.drop_agg_id(agg_id), + } + } +} + +#[cfg(test)] +mod drop_agg_id_tests { + use super::*; + use crate::data_model::AggregationType; + use crate::precompute_operators::SumAccumulator; + use asap_types::aggregation_config::AggregationConfig; + use asap_types::enums::WindowType; + use promql_utilities::data_model::key_by_label_names::KeyByLabelNames; + + fn two_agg_streaming_config() -> Arc { + let cfg = |id: u64| AggregationConfig { + aggregation_id: id, + aggregation_type: AggregationType::Sum, + aggregation_sub_type: String::new(), + parameters: HashMap::new(), + grouping_labels: KeyByLabelNames::empty(), + aggregated_labels: KeyByLabelNames::empty(), + rollup_labels: KeyByLabelNames::empty(), + original_yaml: String::new(), + window_size: 1, + slide_interval: 1, + window_type: WindowType::Tumbling, + spatial_filter: String::new(), + spatial_filter_normalized: String::new(), + metric: format!("metric_{id}"), + num_aggregates_to_retain: None, + read_count_threshold: None, + table_name: None, + value_column: None, + }; + let mut map = HashMap::new(); + map.insert(1u64, cfg(1)); + map.insert(2u64, cfg(2)); + Arc::new(StreamingConfig::new(map)) + } + + fn write_one(store: &SimpleMapStore, agg_id: u64, value: f64, ts: u64) { + let acc = SumAccumulator::with_sum(value); + let output = PrecomputedOutput::new(ts, ts + 1000, None, agg_id); + store + .insert_precomputed_output(output, Box::new(acc)) + .expect("insert ok"); + } + + fn total_buckets(store: &SimpleMapStore, metric: &str, agg_id: u64) -> usize { + let map = store + .query_precomputed_output(metric, agg_id, 0, u64::MAX / 2) + .expect("query ok"); + map.values().map(|v| v.len()).sum() + } + + fn make_store(strategy: LockStrategy) -> SimpleMapStore { + SimpleMapStore::new_with_strategy( + two_agg_streaming_config(), + CleanupPolicy::NoCleanup, + strategy, + ) + } + + #[test] + fn global_drop_removes_only_target_agg() { + let store = make_store(LockStrategy::Global); + write_one(&store, 1, 10.0, 0); + write_one(&store, 1, 20.0, 1_000); + write_one(&store, 1, 30.0, 2_000); + write_one(&store, 2, 99.0, 5_000); + + assert_eq!(total_buckets(&store, "metric_1", 1), 3); + assert_eq!(total_buckets(&store, "metric_2", 2), 1); + + let evicted = store.drop_agg_id(1).expect("drop ok"); + assert_eq!(evicted, 3); + + assert_eq!(total_buckets(&store, "metric_1", 1), 0); + assert_eq!(total_buckets(&store, "metric_2", 2), 1); + } + + #[test] + fn global_drop_unknown_agg_is_noop_returns_zero() { + let store = make_store(LockStrategy::Global); + let evicted = store.drop_agg_id(999).expect("drop ok"); + assert_eq!(evicted, 0); + } + + #[test] + fn global_drop_clears_earliest_timestamp_index() { + let store = make_store(LockStrategy::Global); + write_one(&store, 1, 10.0, 100); + let before = store.get_earliest_timestamp_per_aggregation_id().unwrap(); + assert!(before.contains_key(&1)); + + store.drop_agg_id(1).unwrap(); + + let after = store.get_earliest_timestamp_per_aggregation_id().unwrap(); + assert!(!after.contains_key(&1)); + } + + #[test] + fn per_key_drop_removes_only_target_agg() { + let store = make_store(LockStrategy::PerKey); + write_one(&store, 1, 10.0, 0); + write_one(&store, 1, 20.0, 1_000); + write_one(&store, 2, 99.0, 5_000); + + assert_eq!(total_buckets(&store, "metric_1", 1), 2); + assert_eq!(total_buckets(&store, "metric_2", 2), 1); + + let evicted = store.drop_agg_id(1).expect("drop ok"); + assert_eq!(evicted, 2); + + assert_eq!(total_buckets(&store, "metric_1", 1), 0); + assert_eq!(total_buckets(&store, "metric_2", 2), 1); + } + + #[test] + fn per_key_drop_unknown_agg_is_noop() { + let store = make_store(LockStrategy::PerKey); + let evicted = store.drop_agg_id(999).expect("drop ok"); + assert_eq!(evicted, 0); + } + + #[test] + fn per_key_drop_clears_earliest_timestamp_index() { + let store = make_store(LockStrategy::PerKey); + write_one(&store, 1, 10.0, 100); + let before = store.get_earliest_timestamp_per_aggregation_id().unwrap(); + assert!(before.contains_key(&1)); + + store.drop_agg_id(1).unwrap(); + + let after = store.get_earliest_timestamp_per_aggregation_id().unwrap(); + assert!(!after.contains_key(&1)); + } + + #[test] + fn drop_then_reinsert_behaves_as_fresh_agg() { + let store = make_store(LockStrategy::Global); + write_one(&store, 1, 10.0, 100); + store.drop_agg_id(1).unwrap(); + + write_one(&store, 1, 42.0, 500); + assert_eq!(total_buckets(&store, "metric_1", 1), 1); + let ts_map = store.get_earliest_timestamp_per_aggregation_id().unwrap(); + assert_eq!(ts_map.get(&1).copied(), Some(500)); + } } diff --git a/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs b/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs index 6e618811..dde41d7e 100644 --- a/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/sketch_db/simple_map_store/per_key.rs @@ -1008,6 +1008,40 @@ impl Store for SimpleMapStorePerKey { info!("SimpleMapStorePerKey closed"); Ok(()) } + + fn drop_agg_id(&self, agg_id: u64) -> StoreResult { + // Count windows before eviction for the return value. Grab + // the shard's read lock to snapshot the count, then upgrade + // to a removal. DashMap::remove is atomic on the key — no + // global lock needed, so a concurrent insert for a + // different agg_id keeps running unblocked. + let evicted = match self.inner.store.get(&agg_id) { + Some(entry) => { + let data = entry.value().read().unwrap(); + let count = data.current_epoch.len() + + data + .sealed_epochs + .values() + .map(|e| e.entries.len()) + .sum::(); + drop(data); + drop(entry); + count + } + None => 0, + }; + self.inner.store.remove(&agg_id); + self.inner.earliest_timestamps.remove(&agg_id); + // `metrics` and `items_inserted` are keyed by metric name, + // not agg_id, so we don't touch them — other agg_ids under + // the same metric stay live. + info!( + agg_id, + evicted_windows = evicted, + "SimpleMapStorePerKey::drop_agg_id" + ); + Ok(evicted) + } } // ================================================================= diff --git a/asap-query-engine/src/stores/traits.rs b/asap-query-engine/src/stores/traits.rs index 679568c3..dbd1de2e 100644 --- a/asap-query-engine/src/stores/traits.rs +++ b/asap-query-engine/src/stores/traits.rs @@ -57,6 +57,33 @@ pub trait Store: Send + Sync { /// Close the store and clean up resources fn close(&self) -> Result<(), Box>; + + /// Drop every precompute record whose `aggregation_id` equals + /// `agg_id`, regardless of window or subpopulation key. Used by + /// the §6-driven `SchemaEvictionService` to reclaim space when + /// a retired schema passes its `expires_at_ms`. + /// + /// Returns the number of records removed (best-effort — a store + /// that can't easily count still returns 0 and logs, never + /// fails). + /// + /// Contract: + /// * Must be idempotent — calling on an unknown `agg_id` is a + /// no-op that returns `Ok(0)`. + /// * Must be atomic AT LEAST with respect to concurrent reads + /// for the same agg_id: a reader either sees all of the + /// pre-drop records or none, not a half-dropped state. Stores + /// backed by a single RwLock get this for free; stores with + /// finer-grained locking may need to grab a global lock + /// briefly. + /// * Does NOT touch any registry (backfill, schema); the caller + /// is responsible for post-drop cleanup of those. + /// + /// Default impl returns `Ok(0)` so existing stores stay + /// compiling — the impl must override to actually delete. + fn drop_agg_id(&self, _agg_id: u64) -> Result> { + Ok(0) + } } /// Result type for store operations