feat(store): added store interface for batched exact-window queries (#627)

* wip(query-engine): batched exact-window store query (#609)

Checkpoint before merging main (legacy stores gated behind feature
flag in #624). Not yet wired into scan_windows_via_exact; no tests
yet.

* wip(query-engine): wire batched exact-window query into scan_windows_via_exact (#609)

Checkpoint before rebasing onto main (legacy stores removed in #625).

* fix(query-engine): restore lock_profiling on batch exact-query, dedupe per_key/global

Addresses review feedback on #609's batch method:
- per_key.rs and global.rs had dropped the lock_profiling wait/hold-time
  instrumentation that the single-window path has, on exactly the path
  now most affected by longer lock hold times.
- The per-window epoch-resolution loop (current_epoch/sealed_epochs
  lookup, read-count bookkeeping) was copy-pasted between the two
  backends. Extracted into common::resolve_exact_windows, shared by both.

Author: milindsrivastava1997

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<crate::stores::TimestampRange> = 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(
-                    &params.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(&params.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<TimestampedBucketsMap, Box<dyn std::error::Error + Send + Sync>> {
             panic!("NoOpStore should not be called by merge_accumulators tests");
         }
+        fn query_precomputed_output_exact_batch(
+            &self,
+            _: &str,
+            _: u64,
+            _: &[crate::stores::TimestampRange],
+        ) -> Result<TimestampedBucketsMap, Box<dyn std::error::Error + Send + Sync>> {
+            panic!("NoOpStore should not be called by merge_accumulators tests");
+        }
         fn get_earliest_timestamp_per_aggregation_id(
             &self,
         ) -> Result<HashMap<u64, u64>, Box<dyn std::error::Error + Send + Sync>> {

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<MetricID, Vec<(TimestampRange, Arc<dyn AggregateCore>)>>;
 
 /// 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<EpochID, SealedEpoch>,
+    intern: &InternTable,
+    windows: &[TimestampRange],
+    metric: &str,
+    aggregation_id: u64,
+) -> (TimestampedBucketsMap, Vec<TimestampRange>, usize) {
+    let mut results: TimestampedBucketsMap = HashMap::new();
+    let mut matched_windows: Vec<TimestampRange> = 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)
+}

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<TimestampedBucketsMap, Box<dyn std::error::Error + Send + Sync>> {
+        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<HashMap<u64, u64>, Box<dyn std::error::Error + Send + Sync>> {

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<TimestampedBucketsMap, Box<dyn std::error::Error + Send + Sync>> {
+        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<HashMap<u64, u64>, Box<dyn std::error::Error + Send + Sync>> {

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<TimestampedBucketsMap, Box<dyn std::error::Error + Send + Sync>> {
+        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<HashMap<u64, u64>, Box<dyn std::error::Error + Send + Sync>> {