diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index be14d90..33fdd7d 100644 --- a/asap-query-engine/src/engines/simple_engine/mod.rs +++ b/asap-query-engine/src/engines/simple_engine/mod.rs @@ -554,34 +554,32 @@ impl SimpleEngine { let window_size_ms = config.window_size_ms; let step_ms = Self::bucket_step_ms(config); - let mut merged: TimestampedBucketsMap = HashMap::new(); if window_size_ms == 0 || step_ms == 0 || params.start_timestamp > params.end_timestamp { - return Ok(merged); + return Ok(HashMap::new()); } + let mut windows: Vec = Vec::new(); let mut window_start = params.start_timestamp.div_ceil(step_ms) * step_ms; while window_start + window_size_ms <= params.end_timestamp { - let window_end = window_start + window_size_ms; - let partial = self - .store - .query_precomputed_output_exact( - ¶ms.metric, - params.aggregation_id, - window_start, - window_end, - ) - .map_err(|e| { - format!( - "Error querying store for metric {}, agg {}, window [{}, {}]: {}", - params.metric, params.aggregation_id, window_start, window_end, e - ) - })?; - for (key, buckets) in partial { - merged.entry(key).or_default().extend(buckets); - } + windows.push((window_start, window_start + window_size_ms)); window_start += step_ms; } - Ok(merged) + + // #609: one batched store call for the whole grid instead of one + // query_precomputed_output_exact call per window. + self.store + .query_precomputed_output_exact_batch(¶ms.metric, params.aggregation_id, &windows) + .map_err(|e| { + format!( + "Error querying store for metric {}, agg {}, {} windows in [{}, {}]: {}", + params.metric, + params.aggregation_id, + windows.len(), + params.start_timestamp, + params.end_timestamp, + e + ) + }) } /// Executes a single store query based on parameters @@ -2710,6 +2708,14 @@ mod merge_accumulators_regression_tests_596 { ) -> Result> { panic!("NoOpStore should not be called by merge_accumulators tests"); } + fn query_precomputed_output_exact_batch( + &self, + _: &str, + _: u64, + _: &[crate::stores::TimestampRange], + ) -> Result> { + panic!("NoOpStore should not be called by merge_accumulators tests"); + } fn get_earliest_timestamp_per_aggregation_id( &self, ) -> Result, Box> { diff --git a/asap-query-engine/src/stores/simple_map_store/common.rs b/asap-query-engine/src/stores/simple_map_store/common.rs index d77142f..bab287c 100644 --- a/asap-query-engine/src/stores/simple_map_store/common.rs +++ b/asap-query-engine/src/stores/simple_map_store/common.rs @@ -1,10 +1,12 @@ use crate::data_model::{AggregateCore, KeyByLabelValues}; -use std::collections::{HashMap, HashSet}; +pub use crate::stores::TimestampRange; +use crate::stores::TimestampedBucketsMap; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::sync::{Arc, OnceLock}; +use tracing::debug; pub type MetricID = u32; pub type EpochID = u64; -pub type TimestampRange = (u64, u64); pub type MetricBucketMap = HashMap)>>; /// Sorts one key's buckets into chronological (ascending start) order. @@ -418,3 +420,50 @@ impl SealedEpoch { windows } } + +/// Resolves every window in `windows` against `current_epoch`/`sealed_epochs`, merging the +/// results into one map. Shared by `SimpleMapStorePerKey` and `SimpleMapStoreGlobal`'s +/// `query_precomputed_output_exact_batch` (#609) — the only difference between the two +/// backends is how the outer per-aggregation lock is acquired and how `read_counts` is keyed, +/// both handled by the caller. Returns `(results, matched_windows, total_entries)`; +/// `matched_windows` is what the caller bumps read counts for. +pub fn resolve_exact_windows( + current_epoch: &MutableEpoch, + sealed_epochs: &BTreeMap, + intern: &InternTable, + windows: &[TimestampRange], + metric: &str, + aggregation_id: u64, +) -> (TimestampedBucketsMap, Vec, usize) { + let mut results: TimestampedBucketsMap = HashMap::new(); + let mut matched_windows: Vec = Vec::new(); + let mut total_entries = 0; + + for &window in windows { + if window.0 > window.1 { + debug!( + "Invalid exact query range for metric {} agg_id {}: start {} > end {}", + metric, aggregation_id, window.0, window.1 + ); + continue; + } + + let entries_opt = current_epoch.exact_query(window).or_else(|| { + sealed_epochs + .values() + .rev() + .find_map(|epoch| epoch.exact_query(window)) + }); + + if let Some(entries) = entries_opt { + for (metric_id, agg) in entries { + let label = intern.resolve(metric_id).clone(); + results.entry(label).or_default().push((window, agg)); + total_entries += 1; + } + matched_windows.push(window); + } + } + + (results, matched_windows, total_entries) +} diff --git a/asap-query-engine/src/stores/simple_map_store/global.rs b/asap-query-engine/src/stores/simple_map_store/global.rs index 0cf2b95..6958ed1 100644 --- a/asap-query-engine/src/stores/simple_map_store/global.rs +++ b/asap-query-engine/src/stores/simple_map_store/global.rs @@ -2,8 +2,8 @@ use crate::data_model::{ AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, }; use crate::stores::simple_map_store::common::{ - sort_buckets_chronologically, EpochID, InternTable, MetricBucketMap, MutableEpoch, SealedEpoch, - TimestampRange, + resolve_exact_windows, sort_buckets_chronologically, EpochID, InternTable, MetricBucketMap, + MutableEpoch, SealedEpoch, TimestampRange, }; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use std::collections::{BTreeMap, HashMap, HashSet}; @@ -685,6 +685,96 @@ impl Store for SimpleMapStoreGlobal { Ok(results) } + /// Batched exact-window lookup (#609): acquires the process-wide lock once for the + /// whole `windows` slice instead of once per window. Otherwise identical semantics to + /// calling `query_precomputed_output_exact` once per window and merging the results + /// (a window with no exact match simply contributes nothing). + fn query_precomputed_output_exact_batch( + &self, + metric: &str, + aggregation_id: u64, + windows: &[TimestampRange], + ) -> Result> { + if windows.is_empty() { + return Ok(HashMap::new()); + } + + let query_start_time = Instant::now(); + let store_key = aggregation_id; + + // Measure lock acquisition time + #[cfg(feature = "lock_profiling")] + let lock_wait_start = Instant::now(); + + let mut data = self.lock.lock().unwrap(); + + #[cfg(feature = "lock_profiling")] + { + let lock_wait_duration = lock_wait_start.elapsed(); + info!( + "🔒 Batched exact query lock wait time: {:.2}ms (metric: {}, agg_id: {}, windows: {})", + lock_wait_duration.as_secs_f64() * 1000.0, + metric, + aggregation_id, + windows.len() + ); + } + + #[cfg(feature = "lock_profiling")] + let lock_hold_start = Instant::now(); + + let per_key = match data.stores.get(&store_key) { + Some(pk) => pk, + None => { + debug!( + "Metric {} not found in store for batched exact query", + metric + ); + return Ok(HashMap::new()); + } + }; + + let (results, found_windows, total_entries) = resolve_exact_windows( + &per_key.current_epoch, + &per_key.sealed_epochs, + &per_key.intern, + windows, + metric, + aggregation_id, + ); + + // Update read counts (outer Mutex held — no inner Mutex needed) + if !found_windows.is_empty() { + let rc_map = data.read_counts.entry(store_key).or_default(); + for window in &found_windows { + *rc_map.entry(*window).or_insert(0) += 1; + } + } + + #[cfg(feature = "lock_profiling")] + { + let lock_hold_duration = lock_hold_start.elapsed(); + info!( + "🔓 Batched exact query lock hold time: {:.2}ms (metric: {}, agg_id: {}, matched: {})", + lock_hold_duration.as_secs_f64() * 1000.0, + metric, + aggregation_id, + found_windows.len() + ); + } + + let query_duration = query_start_time.elapsed(); + debug!( + "Batched exact timestamp query took: {:.2}ms ({} windows requested, {} matched, {} entries)", + query_duration.as_secs_f64() * 1000.0, + windows.len(), + found_windows.len(), + total_entries + ); + + Ok(results) + } + fn get_earliest_timestamp_per_aggregation_id( &self, ) -> Result, Box> { diff --git a/asap-query-engine/src/stores/simple_map_store/mod.rs b/asap-query-engine/src/stores/simple_map_store/mod.rs index a8a6d32..886d326 100644 --- a/asap-query-engine/src/stores/simple_map_store/mod.rs +++ b/asap-query-engine/src/stores/simple_map_store/mod.rs @@ -130,6 +130,22 @@ impl Store for SimpleMapStore { } } + fn query_precomputed_output_exact_batch( + &self, + metric: &str, + aggregation_id: u64, + windows: &[crate::stores::TimestampRange], + ) -> Result> { + match self { + SimpleMapStore::Global(store) => { + store.query_precomputed_output_exact_batch(metric, aggregation_id, windows) + } + SimpleMapStore::PerKey(store) => { + store.query_precomputed_output_exact_batch(metric, aggregation_id, windows) + } + } + } + fn get_earliest_timestamp_per_aggregation_id( &self, ) -> Result, Box> { diff --git a/asap-query-engine/src/stores/simple_map_store/per_key.rs b/asap-query-engine/src/stores/simple_map_store/per_key.rs index 5166a63..928f1f2 100644 --- a/asap-query-engine/src/stores/simple_map_store/per_key.rs +++ b/asap-query-engine/src/stores/simple_map_store/per_key.rs @@ -2,8 +2,8 @@ use crate::data_model::{ AggregateCore, AggregationType, CleanupPolicy, PrecomputedOutput, StreamingConfig, }; use crate::stores::simple_map_store::common::{ - sort_buckets_chronologically, EpochID, InternTable, MetricBucketMap, MetricID, MutableEpoch, - SealedEpoch, TimestampRange, + resolve_exact_windows, sort_buckets_chronologically, EpochID, InternTable, MetricBucketMap, + MetricID, MutableEpoch, SealedEpoch, TimestampRange, }; use crate::stores::{Store, StoreResult, TimestampedBucketsMap}; use dashmap::DashMap; @@ -756,6 +756,119 @@ impl Store for SimpleMapStorePerKey { Ok(results) } + /// Batched exact-window lookup (#609): acquires the shard's read lock once for the + /// whole `windows` slice instead of once per window, resolving each window against + /// `current_epoch` / `sealed_epochs` in a single pass. Otherwise identical semantics + /// to calling `query_precomputed_output_exact` once per window and merging the + /// results (a window with no exact match simply contributes nothing). + fn query_precomputed_output_exact_batch( + &self, + metric: &str, + aggregation_id: u64, + windows: &[TimestampRange], + ) -> Result> { + if windows.is_empty() { + return Ok(HashMap::new()); + } + + let query_start_time = Instant::now(); + let store_key = aggregation_id; + + #[cfg(feature = "lock_profiling")] + let lock_wait_start = Instant::now(); + + let store_data_lock = match self.store.get(&store_key) { + Some(lock) => lock, + None => { + debug!( + "Metric {} not found in store for batched exact query", + metric + ); + return Ok(HashMap::new()); + } + }; + + #[cfg(feature = "lock_profiling")] + { + let lock_wait_duration = lock_wait_start.elapsed(); + info!( + "🔒 Batched exact query DashMap get time: {:.2}ms (metric: {}, agg_id: {}, windows: {})", + lock_wait_duration.as_secs_f64() * 1000.0, + metric, + aggregation_id, + windows.len() + ); + } + + #[cfg(feature = "lock_profiling")] + let rwlock_wait_start = Instant::now(); + + // Same rationale as query_precomputed_output_exact: exact_query takes &self, so a + // read lock covers the whole batch (issue #607). + let data = store_data_lock.read().map_err(|e| { + format!( + "Failed to acquire read lock for batched exact query aggregation_id {}: {}", + store_key, e + ) + })?; + + #[cfg(feature = "lock_profiling")] + { + let rwlock_wait_duration = rwlock_wait_start.elapsed(); + info!( + "🔒 Batched exact query RwLock wait time: {:.2}ms (metric: {}, agg_id: {}, windows: {})", + rwlock_wait_duration.as_secs_f64() * 1000.0, + metric, + aggregation_id, + windows.len() + ); + } + + #[cfg(feature = "lock_profiling")] + let lock_hold_start = Instant::now(); + + let (results, found_windows, total_entries) = resolve_exact_windows( + &data.current_epoch, + &data.sealed_epochs, + &data.intern, + windows, + metric, + aggregation_id, + ); + + // Batch the read-count update too: one inner-Mutex acquisition for every window + // that hit, instead of one per window. + if !found_windows.is_empty() { + let mut read_counts = data.read_counts.lock().unwrap(); + for window in &found_windows { + *read_counts.entry(*window).or_insert(0) += 1; + } + } + + #[cfg(feature = "lock_profiling")] + { + let lock_hold_duration = lock_hold_start.elapsed(); + info!( + "🔓 Batched exact query lock hold time: {:.2}ms (metric: {}, agg_id: {}, matched: {})", + lock_hold_duration.as_secs_f64() * 1000.0, + metric, + aggregation_id, + found_windows.len() + ); + } + + let query_duration = query_start_time.elapsed(); + debug!( + "Batched exact timestamp query took: {:.2}ms ({} windows requested, {} matched, {} entries)", + query_duration.as_secs_f64() * 1000.0, + windows.len(), + found_windows.len(), + total_entries + ); + + Ok(results) + } + fn get_earliest_timestamp_per_aggregation_id( &self, ) -> Result, Box> { diff --git a/asap-query-engine/src/stores/traits.rs b/asap-query-engine/src/stores/traits.rs index 679568c..7503218 100644 --- a/asap-query-engine/src/stores/traits.rs +++ b/asap-query-engine/src/stores/traits.rs @@ -2,8 +2,11 @@ use crate::data_model::{AggregateCore, KeyByLabelValues, PrecomputedOutput}; use std::collections::HashMap; use std::sync::Arc; +/// A (start_timestamp, end_timestamp) pair identifying one stored window. +pub type TimestampRange = (u64, u64); + /// A bucket with its timestamp range: ((start_timestamp, end_timestamp), aggregate) -pub type TimestampedBucket = ((u64, u64), Arc); +pub type TimestampedBucket = (TimestampRange, Arc); /// Map from key to timestamped buckets (sparse - only contains buckets that exist) pub type TimestampedBucketsMap = HashMap, Vec>; @@ -50,6 +53,20 @@ pub trait Store: Send + Sync { exact_end: u64, ) -> Result>; + /// Batched variant of `query_precomputed_output_exact` (issue #609): resolves every + /// window in `windows` against the same (metric, aggregation_id) shard in one lock + /// acquisition instead of one per window. Results across all windows are merged into + /// a single map exactly as `query_precomputed_output_exact` would merge them one at a + /// time — each `TimestampedBucket` still carries its own window, so no information is + /// lost relative to per-window calls. Windows with no exact match simply contribute + /// nothing (same semantics as the single-window method returning an empty map). + fn query_precomputed_output_exact_batch( + &self, + metric: &str, + aggregation_id: u64, + windows: &[TimestampRange], + ) -> Result>; + /// Get earliest timestamp for each aggregation ID (for monitoring) fn get_earliest_timestamp_per_aggregation_id( &self, diff --git a/asap-query-engine/src/tests/query_equivalence_tests.rs b/asap-query-engine/src/tests/query_equivalence_tests.rs index add4504..80890a8 100644 --- a/asap-query-engine/src/tests/query_equivalence_tests.rs +++ b/asap-query-engine/src/tests/query_equivalence_tests.rs @@ -43,6 +43,17 @@ impl Store for NoOpStore { ); } + fn query_precomputed_output_exact_batch( + &self, + _metric: &str, + _aggregation_id: u64, + _windows: &[crate::stores::TimestampRange], + ) -> Result> { + panic!( + "NoOpStore: query_precomputed_output_exact_batch should not be called in equivalence tests" + ); + } + fn insert_precomputed_output( &self, _output: crate::data_model::PrecomputedOutput, diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index e913e6a..1636f2f 100644 --- a/asap-query-engine/src/tests/store_correctness_tests.rs +++ b/asap-query-engine/src/tests/store_correctness_tests.rs @@ -194,6 +194,13 @@ pub fn run_contract_suite(strategy: LockStrategy) { test_exact_query_correct_across_interleaved_inserts_and_queries(strategy); test_exact_query_correct_after_epoch_rotation(strategy); + // Batched exact-query correctness (#609) + test_batch_exact_query_empty_windows_returns_empty(strategy); + test_batch_exact_query_unknown_metric_returns_empty(strategy); + test_batch_exact_query_invalid_window_is_skipped_not_erroring(strategy); + test_batch_exact_query_equivalent_to_sequential_calls(strategy); + test_batch_exact_query_updates_read_counts_for_cleanup(strategy); + test_batch_insert_full_range_query_returns_all(strategy); test_batch_insert_results_are_chronologically_ordered(strategy); test_range_query_returns_only_windows_within_range(strategy); @@ -548,6 +555,169 @@ fn test_exact_query_correct_after_epoch_rotation(strategy: LockStrategy) { } } +// ── batched exact-query correctness (#609) ─────────────────────────────────── +// +// `query_precomputed_output_exact_batch` must be observationally identical to +// calling `query_precomputed_output_exact` once per window and merging the +// results — the only difference is that it takes the shard lock once for the +// whole slice instead of once per window. + +fn test_batch_exact_query_empty_windows_returns_empty(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let (out, acc) = sum_entry(1, 1_000, 2_000, 1.0); + store.insert_precomputed_output(out, acc).unwrap(); + + let result = store + .query_precomputed_output_exact_batch("cpu_usage", 1, &[]) + .unwrap(); + assert!( + result.is_empty(), + "[{}] an empty windows slice must return an empty map, even with data present", + label(strategy) + ); +} + +fn test_batch_exact_query_unknown_metric_returns_empty(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let result = store + .query_precomputed_output_exact_batch("cpu_usage", 1, &[(1_000, 2_000), (3_000, 4_000)]) + .unwrap(); + assert!( + result.is_empty(), + "[{}] batched exact query against an aggregation_id with no data must return empty", + label(strategy) + ); +} + +fn test_batch_exact_query_invalid_window_is_skipped_not_erroring(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let (out, acc) = sum_entry(1, 1_000, 2_000, 1.0); + store.insert_precomputed_output(out, acc).unwrap(); + + // First window is invalid (start > end); must be skipped, not fail the whole batch. + let result = store + .query_precomputed_output_exact_batch("cpu_usage", 1, &[(2_000, 1_000), (1_000, 2_000)]) + .unwrap(); + assert_eq!( + total_bucket_count(&result), + 1, + "[{}] an invalid window in the batch must be skipped, not fail or drop valid windows", + label(strategy) + ); +} + +fn test_batch_exact_query_equivalent_to_sequential_calls(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let a = key(&["a"]); + let b = key(&["b"]); + store + .insert_precomputed_output( + PrecomputedOutput::new(1_000, 2_000, Some(a.clone()), 1), + Box::new(SumAccumulator::with_sum(10.0)), + ) + .unwrap(); + store + .insert_precomputed_output( + PrecomputedOutput::new(2_000, 3_000, Some(b.clone()), 1), + Box::new(SumAccumulator::with_sum(20.0)), + ) + .unwrap(); + store + .insert_precomputed_output( + PrecomputedOutput::new(3_000, 4_000, Some(a.clone()), 1), + Box::new(SumAccumulator::with_sum(30.0)), + ) + .unwrap(); + + // Windows include a miss (4_000, 5_000) between two hits, mirroring a real + // scan_windows_via_exact grid walk over a range with a gap. + let windows = [ + (1_000, 2_000), + (2_000, 3_000), + (4_000, 5_000), + (3_000, 4_000), + ]; + + let batched = store + .query_precomputed_output_exact_batch("cpu_usage", 1, &windows) + .unwrap(); + + let mut sequential: TimestampedBucketsMap = HashMap::new(); + for &(start, end) in &windows { + let partial = store + .query_precomputed_output_exact("cpu_usage", 1, start, end) + .unwrap(); + for (k, buckets) in partial { + sequential.entry(k).or_default().extend(buckets); + } + } + + assert_eq!( + timestamps_for_key(&batched, &a), + timestamps_for_key(&sequential, &a), + "[{}] batched result for key 'a' must match sequential per-window calls merged together", + label(strategy) + ); + assert_eq!( + timestamps_for_key(&batched, &b), + timestamps_for_key(&sequential, &b), + "[{}] batched result for key 'b' must match sequential per-window calls merged together", + label(strategy) + ); + assert_eq!( + total_bucket_count(&batched), + 3, + "[{}] batched call must find exactly the 3 hits among the 4 requested windows", + label(strategy) + ); +} + +fn test_batch_exact_query_updates_read_counts_for_cleanup(strategy: LockStrategy) { + // read_count_threshold = 2: mirrors test_cleanup_read_based_evicts_after_threshold_reads, + // but drives the reads through the batch call instead of one-at-a-time, to pin that + // batching still updates read_counts per matched window (not e.g. once per batch call). + let store = make_store( + strategy, + CleanupPolicy::ReadBased, + &[(1, AggregationType::Sum, None, Some(2))], + ); + let (out, acc) = sum_entry(1, 1_000, 2_000, 1.0); + store.insert_precomputed_output(out, acc).unwrap(); + + // Read 1 via the batch call — count becomes 1. + store + .query_precomputed_output_exact_batch("cpu_usage", 1, &[(1_000, 2_000)]) + .unwrap(); + let (o2, a2) = sum_entry(1, 3_000, 4_000, 2.0); + store.insert_precomputed_output(o2, a2).unwrap(); + + let still_there = store + .query_precomputed_output_exact("cpu_usage", 1, 1_000, 2_000) + .unwrap(); + assert_eq!( + total_bucket_count(&still_there), + 1, + "[{}] window must survive until read count reaches threshold", + label(strategy) + ); + + // Read 2 via the batch call — count becomes 2, evicted on the next insert. + store + .query_precomputed_output_exact_batch("cpu_usage", 1, &[(1_000, 2_000)]) + .unwrap(); + let (o3, a3) = sum_entry(1, 5_000, 6_000, 3.0); + store.insert_precomputed_output(o3, a3).unwrap(); + + let evicted = store + .query_precomputed_output_exact("cpu_usage", 1, 1_000, 2_000) + .unwrap(); + assert!( + evicted.is_empty(), + "[{}] window must be evicted once batched reads bring its count to threshold", + label(strategy) + ); +} + // ── batch insert correctness ────────────────────────────────────────────────── fn test_batch_insert_full_range_query_returns_all(strategy: LockStrategy) {