From be58526fb5b78e5cd2aac8d23520eea45251133d Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 23:10:56 -0400 Subject: [PATCH 1/3] 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. --- .../src/stores/simple_map_store/common.rs | 2 +- .../src/stores/simple_map_store/global.rs | 80 +++++++++++++++++ .../src/stores/simple_map_store/mod.rs | 16 ++++ .../src/stores/simple_map_store/per_key.rs | 87 +++++++++++++++++++ asap-query-engine/src/stores/traits.rs | 19 +++- .../src/tests/query_equivalence_tests.rs | 11 +++ 6 files changed, 213 insertions(+), 2 deletions(-) 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 d77142f6..5a54c60d 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,10 @@ use crate::data_model::{AggregateCore, KeyByLabelValues}; +pub use crate::stores::TimestampRange; use std::collections::{HashMap, HashSet}; use std::sync::{Arc, OnceLock}; 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. 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 0cf2b95a..959f8ae3 100644 --- a/asap-query-engine/src/stores/simple_map_store/global.rs +++ b/asap-query-engine/src/stores/simple_map_store/global.rs @@ -685,6 +685,86 @@ 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; + + let mut data = self.lock.lock().unwrap(); + + 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 mut results: TimestampedBucketsMap = HashMap::new(); + let mut found_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; + } + + // Check current epoch first (newest), then sealed epochs newest-to-oldest. + let entries_opt: Option> = + per_key.current_epoch.exact_query(window).or_else(|| { + per_key + .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 = per_key.intern.resolve(metric_id).clone(); + results.entry(label).or_default().push((window, agg)); + total_entries += 1; + } + found_windows.push(window); + } + } + + // 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; + } + } + + 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 a8a6d321..886d326d 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 5166a632..e8a7cbb3 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 @@ -756,6 +756,93 @@ 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; + + 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()); + } + }; + + // 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 + ) + })?; + + let mut results: TimestampedBucketsMap = HashMap::new(); + let mut found_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: Option)>> = + data.current_epoch.exact_query(window).or_else(|| { + data.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 = data.intern.resolve(metric_id).clone(); + results.entry(label).or_default().push((window, agg)); + total_entries += 1; + } + found_windows.push(window); + } + } + + // 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; + } + } + + 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 679568c3..7503218b 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 add45048..80890a89 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, From 47904b62a0d8ec51ef82f904926b62280a3080e3 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 23:28:44 -0400 Subject: [PATCH 2/3] wip(query-engine): wire batched exact-window query into scan_windows_via_exact (#609) Checkpoint before rebasing onto main (legacy stores removed in #625). --- .../src/engines/simple_engine/mod.rs | 48 ++--- .../src/stores/simple_map_store/global.rs | 5 +- .../src/stores/simple_map_store/per_key.rs | 5 +- .../src/tests/store_correctness_tests.rs | 170 ++++++++++++++++++ 4 files changed, 205 insertions(+), 23 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine/mod.rs b/asap-query-engine/src/engines/simple_engine/mod.rs index be14d90c..33fdd7de 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/global.rs b/asap-query-engine/src/stores/simple_map_store/global.rs index 959f8ae3..19706129 100644 --- a/asap-query-engine/src/stores/simple_map_store/global.rs +++ b/asap-query-engine/src/stores/simple_map_store/global.rs @@ -707,7 +707,10 @@ impl Store for SimpleMapStoreGlobal { let per_key = match data.stores.get(&store_key) { Some(pk) => pk, None => { - debug!("Metric {} not found in store for batched exact query", metric); + debug!( + "Metric {} not found in store for batched exact query", + metric + ); return Ok(HashMap::new()); } }; 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 e8a7cbb3..9b9624da 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 @@ -777,7 +777,10 @@ impl Store for SimpleMapStorePerKey { 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); + debug!( + "Metric {} not found in store for batched exact query", + metric + ); return Ok(HashMap::new()); } }; diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index e913e6a4..1636f2f5 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) { From dd92fcda31d529a15bb413e85eb406df0da18568 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 23:52:31 -0400 Subject: [PATCH 3/3] 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. --- .../src/stores/simple_map_store/common.rs | 51 +++++++++++- .../src/stores/simple_map_store/global.rs | 75 +++++++++-------- .../src/stores/simple_map_store/per_key.rs | 83 ++++++++++++------- 3 files changed, 144 insertions(+), 65 deletions(-) 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 5a54c60d..bab287c5 100644 --- a/asap-query-engine/src/stores/simple_map_store/common.rs +++ b/asap-query-engine/src/stores/simple_map_store/common.rs @@ -1,7 +1,9 @@ use crate::data_model::{AggregateCore, KeyByLabelValues}; pub use crate::stores::TimestampRange; -use std::collections::{HashMap, HashSet}; +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; @@ -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 19706129..6958ed1f 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}; @@ -702,8 +702,27 @@ impl Store for SimpleMapStoreGlobal { 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 => { @@ -715,38 +734,14 @@ impl Store for SimpleMapStoreGlobal { } }; - let mut results: TimestampedBucketsMap = HashMap::new(); - let mut found_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; - } - - // Check current epoch first (newest), then sealed epochs newest-to-oldest. - let entries_opt: Option> = - per_key.current_epoch.exact_query(window).or_else(|| { - per_key - .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 = per_key.intern.resolve(metric_id).clone(); - results.entry(label).or_default().push((window, agg)); - total_entries += 1; - } - found_windows.push(window); - } - } + 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() { @@ -756,6 +751,18 @@ impl Store for SimpleMapStoreGlobal { } } + #[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)", 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 9b9624da..928f1f2c 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; @@ -774,6 +774,9 @@ impl Store for SimpleMapStorePerKey { 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 => { @@ -785,6 +788,21 @@ impl Store for SimpleMapStorePerKey { } }; + #[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| { @@ -794,36 +812,29 @@ impl Store for SimpleMapStorePerKey { ) })?; - let mut results: TimestampedBucketsMap = HashMap::new(); - let mut found_windows: Vec = Vec::new(); - let mut total_entries = 0; + #[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() + ); + } - 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; - } + #[cfg(feature = "lock_profiling")] + let lock_hold_start = Instant::now(); - let entries_opt: Option)>> = - data.current_epoch.exact_query(window).or_else(|| { - data.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 = data.intern.resolve(metric_id).clone(); - results.entry(label).or_default().push((window, agg)); - total_entries += 1; - } - found_windows.push(window); - } - } + 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. @@ -834,6 +845,18 @@ impl Store for SimpleMapStorePerKey { } } + #[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)",