fix(query-engine): exact-window queries take a read lock, not write

asap-query-engine/src/stores/simple_map_store/INDEX_DESIGN.md

@@ -56,8 +56,8 @@ MutableEpoch {
     // Monotonic ingest fast path (Opt 3)
     last_window:     Option<TimestampRange>
 
-    // Lazy offset index (Opt 1 + 2): built on first exact_query, None after any insert
-    window_to_ids:   Option<HashMap<TimestampRange, Vec<u32>>>
+    // Lazy offset index (Opt 1 + 2): built on first exact_query, cleared on any insert
+    window_to_ids:   OnceLock<HashMap<TimestampRange, Vec<u32>>>
 
     // Epoch bounds for O(1) skip check (updated incrementally on insert)
     min_start:       Option<u64>
@@ -68,14 +68,15 @@ MutableEpoch {
 **Insert** (`O(1)` amortized):
 - Opt 3: if incoming window == `last_window`, skip `windows_set.insert` entirely
 - Three `Vec::push` calls — no secondary index maintenance
-- `window_to_ids = None` — single pointer-width write to invalidate the index
+- `window_to_ids.take()` — through `&mut self`, so no synchronization needed to invalidate the index
 
 **`seal()` → `SealedEpoch`** (`O(M log M)`, paid once at rotation):
 - Zips the three columns into tuples, sorts by `(TimestampRange, MetricID)`, moves `Arc`s without cloning
 
-**`exact_query(&mut self)`** (`O(M)` first call after a write, `O(m)` cached):
-- Opt 1 + 2: if `window_to_ids` is `None`, build it from `windows_col` in one pass storing `u32` offsets
-- Cache is valid until the next `insert`
+**`exact_query(&self)`** (`O(M)` first call after a write, `O(m)` cached):
+- Opt 1 + 2: `window_to_ids.get_or_init(...)` builds the index from `windows_col` in one pass storing `u32` offsets if not already built; otherwise returns the cached index directly
+- `OnceLock` gives build-at-most-once semantics under concurrent callers for free — no external locking needed to call this, and no poisoning risk if the build closure ever panicked
+- Cache is valid until the next `insert` (or `remove_windows`), which calls `window_to_ids.take()`
 
 **`range_query_into`** (`O(M)` mutable epoch):
 - Opt 5: hot loop iterates only `windows_col`; aggregate pointer only chased on match
@@ -200,11 +201,13 @@ No inner `Mutex` for `read_counts` — the outer `Mutex` already serializes all
 
 ### Exact Query `(exact_start, exact_end)`
 
-1. Acquire **write lock** (needed to potentially build the lazy `window_to_ids` index)
-2. Try `current_epoch.exact_query(range)` — builds/uses cached `window_to_ids`
+1. Acquire **read lock** on `StoreKeyData` — same as a range query. `window_to_ids` lives in its
+   own `OnceLock`, independent of this outer lock, so building it (if needed) doesn't require
+   exclusive access here (issue #607).
+2. Try `current_epoch.exact_query(range)` — builds/uses cached `window_to_ids` via `get_or_init`
 3. If not found, iterate `sealed_epochs.values().rev()` calling `SealedEpoch::exact_query`
-4. Return owned `Vec<(MetricID, Arc<dyn AggregateCore>)>`, drop write lock
-5. Re-acquire read lock to resolve MetricIDs → labels
+4. Resolve MetricIDs → labels via `InternTable`, still under the same read lock
+5. Briefly acquire inner `Mutex` to update `read_counts`
 
 ---
 
@@ -239,7 +242,14 @@ No eviction — data accumulates indefinitely.
 |-----------|------|
 | **Insert** | `RwLock::write` for the batch duration |
 | **Range query** | `RwLock::read` → brief `Mutex::lock` on `read_counts` |
-| **Exact query** | `RwLock::write` (lazy index build) → drop → `RwLock::read` for label resolution |
+| **Exact query** | `RwLock::read` (lazy index build no longer needs exclusive access — see below) → brief `Mutex::lock` on `read_counts` |
 | **Cleanup** | Under existing write lock; `Mutex::get_mut()` bypasses inner lock |
 
-Multiple readers per `aggregation_id` run concurrently. Writers only block readers of the same `aggregation_id`.
+Multiple readers per `aggregation_id` run concurrently, including range and exact queries running
+concurrently with each other. Writers only block readers of the same `aggregation_id`.
+
+Before issue #607's fix, exact queries took `RwLock::write` solely because building
+`window_to_ids` required `&mut self`, serializing them against every other concurrent
+reader/writer of the shard even though no queryable data was mutated. Switching `window_to_ids`
+to a `OnceLock` let `exact_query` take `&self`, so its call site only ever needs `RwLock::read`,
+same as a range query.

asap-query-engine/src/stores/simple_map_store/common.rs

@@ -1,6 +1,6 @@
 use crate::data_model::{AggregateCore, KeyByLabelValues};
 use std::collections::{HashMap, HashSet};
-use std::sync::Arc;
+use std::sync::{Arc, OnceLock};
 
 pub type MetricID = u32;
 pub type EpochID = u64;
@@ -77,9 +77,18 @@ impl InternTable {
 ///
 /// **Opt 1 + 2 — Lazy offset index**: `window_to_ids` is built on the *first* `exact_query`
 /// after any write batch and stores u32 column offsets rather than Arc clones.  Any `insert`
-/// simply sets the field to `None` (one pointer-width write); there are no HashMap lookups,
-/// no `HashSet::insert` calls for the index, and no atomic refcount bumps on the hot insert
-/// path.  The index is rebuilt in O(M) on demand from `windows_col` alone.
+/// invalidates it via `OnceLock::take` (called through `&mut self`, so no synchronization is
+/// needed to reset it); there are no HashMap lookups, no `HashSet::insert` calls for the
+/// index, and no atomic refcount bumps on the hot insert path.  The index is rebuilt in O(M)
+/// on demand from `windows_col` alone.
+///
+/// The index lives in its own `OnceLock`, independent of the outer per-shard lock callers
+/// take to reach a `MutableEpoch`. `OnceLock::get_or_init` already gives build-at-most-once
+/// semantics under concurrent callers, so `exact_query` takes `&self` and callers only ever
+/// need shared (read) access to the containing epoch — see `exact_query` below. Unlike
+/// `RwLock`, `OnceLock` has no poisoning: if the init closure ever panicked, the cell would
+/// simply stay uninitialized for the next caller to retry, rather than permanently poisoning
+/// every future exact query on this epoch. (Issue #607.)
 ///
 /// **Opt 3 — Monotonic ingest fast path**: `last_window` tracks the most recently inserted
 /// window.  Consecutive inserts to the same window (multiple label combinations for one time
@@ -102,7 +111,9 @@ pub struct MutableEpoch {
 
     // Lazy offset index: built on first exact_query, invalidated on any insert (Opt 1 + 2).
     // Stores column indices (u32) instead of Arc clones — zero atomic ops during insert.
-    window_to_ids: Option<HashMap<TimestampRange, Vec<u32>>>,
+    // OnceLock (not RwLock) so exact_query can take &self with build-once-per-invalidation
+    // semantics for free, and with no poisoning risk (issue #607).
+    window_to_ids: OnceLock<HashMap<TimestampRange, Vec<u32>>>,
 
     /// Epoch time bounds for O(1) skip check, updated incrementally on insert.
     min_start: Option<u64>,
@@ -123,7 +134,7 @@ impl MutableEpoch {
             aggregates_col: Vec::with_capacity(cap),
             windows_set: HashSet::new(),
             last_window: None,
-            window_to_ids: None,
+            window_to_ids: OnceLock::new(),
             min_start: None,
             max_end: None,
         }
@@ -171,8 +182,9 @@ impl MutableEpoch {
         self.metric_ids_col.push(metric_id);
         self.aggregates_col.push(agg);
 
-        // Opt 1: invalidate lazy index at zero cost
-        self.window_to_ids = None;
+        // Opt 1: invalidate lazy index at zero cost. take() needs &mut self, which we
+        // already have — no synchronization on this path.
+        self.window_to_ids.take();
 
         self.min_start = Some(self.min_start.map_or(range.0, |m| m.min(range.0)));
         self.max_end = Some(self.max_end.map_or(range.1, |m| m.max(range.1)));
@@ -231,29 +243,42 @@ impl MutableEpoch {
     /// The offset index (`HashMap<TimestampRange, Vec<u32>>`) is constructed from `windows_col`
     /// on the first call after any write batch, then cached.  Building it scans `windows_col`
     /// once with no Arc clones (only integer offsets are stored).  The index remains valid
-    /// until the next `insert`, which sets `window_to_ids = None`.
+    /// until the next `insert`, which invalidates it via `OnceLock::take`.
     ///
-    /// Takes `&mut self` because building the index mutates `window_to_ids`.
-    /// Callers must hold exclusive (write) access to the containing epoch.
+    /// Takes `&self`: the index lives in its own `OnceLock`, independent of the outer
+    /// per-shard lock callers hold to reach this epoch. `get_or_init` already builds the
+    /// index at most once per invalidation under concurrent callers (and without the
+    /// poisoning risk a hand-rolled `RwLock`-based cache would have), so callers only need
+    /// shared (read) access to the epoch — see the `window_to_ids` field doc (issue #607).
     pub fn exact_query(
-        &mut self,
+        &self,
         range: TimestampRange,
     ) -> Option<Vec<(MetricID, Arc<dyn AggregateCore>)>> {
-        if self.window_to_ids.is_none() {
+        let index = self.window_to_ids.get_or_init(|| {
             let mut idx: HashMap<TimestampRange, Vec<u32>> =
                 HashMap::with_capacity(self.windows_set.len());
             for (i, &tr) in self.windows_col.iter().enumerate() {
                 idx.entry(tr).or_default().push(i as u32);
             }
-            self.window_to_ids = Some(idx);
-        }
-        let offsets = self.window_to_ids.as_ref().unwrap().get(&range)?;
+            idx
+        });
+        Self::lookup_offsets(index, range, &self.metric_ids_col, &self.aggregates_col)
+    }
+
+    /// Resolves a window's offsets (from the lazy index) into (MetricID, aggregate) pairs.
+    fn lookup_offsets(
+        index: &HashMap<TimestampRange, Vec<u32>>,
+        range: TimestampRange,
+        metric_ids_col: &[MetricID],
+        aggregates_col: &[Arc<dyn AggregateCore>],
+    ) -> Option<Vec<(MetricID, Arc<dyn AggregateCore>)>> {
+        let offsets = index.get(&range)?;
         Some(
             offsets
                 .iter()
                 .map(|&i| {
                     let i = i as usize;
-                    (self.metric_ids_col[i], Arc::clone(&self.aggregates_col[i]))
+                    (metric_ids_col[i], Arc::clone(&aggregates_col[i]))
                 })
                 .collect(),
         )
@@ -281,7 +306,7 @@ impl MutableEpoch {
         }
 
         // Invalidate lazy index and monotonic fast-path hint.
-        self.window_to_ids = None;
+        self.window_to_ids.take();
         self.last_window = None;
 
         // Recompute bounds (cleanup is rare; linear scan is fine).

asap-query-engine/src/stores/simple_map_store/global.rs

@@ -605,38 +605,33 @@ impl Store for SimpleMapStoreGlobal {
 
         let timestamp_range = (exact_start, exact_end);
 
-        // Opt 1: exact_query now takes &mut self (lazy index build).
-        // Call it inside a scoped block so the &mut borrow on data.stores ends before we
-        // re-borrow data.stores immutably to resolve MetricIDs → labels.
-        let entries_opt: Option<Vec<_>> = {
-            let per_key = match data.stores.get_mut(&store_key) {
-                Some(pk) => pk,
-                None => {
-                    debug!("Metric {} not found in store for exact query", metric);
-                    return Ok(HashMap::new());
-                }
-            };
-            // Check current epoch first (newest). exact_query returns an owned Vec so the
-            // &mut borrow of per_key ends immediately — no lifetime overlap with the
-            // sealed_epochs scan below.
-            per_key
-                .current_epoch
-                .exact_query(timestamp_range)
-                .or_else(|| {
-                    per_key
-                        .sealed_epochs
-                        .values()
-                        .rev()
-                        .find_map(|epoch| epoch.exact_query(timestamp_range))
-                })
-        }; // &mut borrow of data.stores ends here
+        // exact_query takes &self (its lazy index build is behind its own inner lock,
+        // see MutableEpoch::exact_query / issue #607) — no &mut borrow of data.stores
+        // needed, so per_key can be looked up once and reused below.
+        let per_key = match data.stores.get(&store_key) {
+            Some(pk) => pk,
+            None => {
+                debug!("Metric {} not found in store for exact query", metric);
+                return Ok(HashMap::new());
+            }
+        };
+        // Check current epoch first (newest), then sealed epochs newest-to-oldest.
+        let entries_opt: Option<Vec<_>> = per_key
+            .current_epoch
+            .exact_query(timestamp_range)
+            .or_else(|| {
+                per_key
+                    .sealed_epochs
+                    .values()
+                    .rev()
+                    .find_map(|epoch| epoch.exact_query(timestamp_range))
+            });
 
         let mut results: TimestampedBucketsMap = HashMap::new();
         let mut total_entries = 0;
         let found_match = entries_opt.is_some();
 
         if let Some(entries) = entries_opt {
-            let per_key = data.stores.get(&store_key).unwrap();
             for (metric_id, agg) in entries {
                 let label = per_key.intern.resolve(metric_id).clone();
                 results

asap-query-engine/src/stores/simple_map_store/per_key.rs

@@ -662,11 +662,11 @@ impl Store for SimpleMapStorePerKey {
         #[cfg(feature = "lock_profiling")]
         let rwlock_wait_start = Instant::now();
 
-        // Opt 1: exact_query takes &mut self (lazy index build), so we need a write lock.
-        // Range queries still use a read lock — only exact queries pay the write-lock cost.
-        let mut data = store_data_lock.write().map_err(|e| {
+        // exact_query takes &self (its lazy index build is behind its own inner lock,
+        // see MutableEpoch::exact_query / issue #607), so a read lock suffices here too.
+        let data = store_data_lock.read().map_err(|e| {
             format!(
-                "Failed to acquire write lock for exact query aggregation_id {}: {}",
+                "Failed to acquire read lock for exact query aggregation_id {}: {}",
                 store_key, e
             )
         })?;
@@ -688,7 +688,7 @@ impl Store for SimpleMapStorePerKey {
         let timestamp_range = (exact_start, exact_end);
 
         // Opt 1: exact_query on the mutable epoch builds the lazy offset index if absent,
-        // then looks up the window in O(m). Returns an owned Vec — the &mut borrow ends here.
+        // then looks up the window in O(m). Takes &self (issue #607) — no write lock needed.
         let entries_opt: Option<Vec<(MetricID, Arc<dyn AggregateCore>)>> =
             data.current_epoch.exact_query(timestamp_range).or_else(|| {
                 data.sealed_epochs
@@ -727,7 +727,8 @@ impl Store for SimpleMapStorePerKey {
             );
         }
 
-        // Update read count — write lock already held, no inner Mutex needed
+        // Update read count. Outer lock is now only a read lock (issue #607), so this
+        // inner Mutex is what actually serializes concurrent read-count updates.
         if found_match {
             let mut read_counts = data.read_counts.lock().unwrap();
             *read_counts.entry(timestamp_range).or_insert(0) += 1;