From 280244f9445401967bd13c98028e1995a5496476 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 14:45:38 -0400 Subject: [PATCH 1/3] fix(query-engine): let exact-window queries take a read lock, not write MutableEpoch::exact_query needed &mut self solely to lazily build/cache its window_to_ids offset index, forcing per_key.rs's exact-query call site to take a full write lock on the per-aggregation shard even though the queried data itself isn't mutated. That serialized exact queries against every other concurrent reader/writer of the shard, unlike range queries, which only ever needed a read lock. window_to_ids now lives behind its own RwLock, independent of the outer shard lock, so exact_query can take &self and build the index lazily under double-checked locking. insert/remove_windows still invalidate it at zero cost via get_mut() (no lock acquired, since &mut self already proves exclusive access). per_key.rs's call site drops to .read(); global.rs is updated to compile against the new signature but has no locking change since it uses a single Mutex with no read/write split. Test coverage was designed by a separate agent given only the problem description (not this fix), working from an isolated worktree at the pre-fix commit, to avoid biasing the tests toward this specific implementation: - Correctness tests confirm exact_query's lazy cache never returns stale/incorrect results across inserts, repeated calls, misses, and epoch rotation. - A concurrency test forces an index rebuild on a large shard and measures whether a concurrent range query on the same shard blocks for the rebuild's full duration; pre-fix it reliably did (ratio ~1.00 across repeated runs), post-fix it does not. Fixes #607. Co-Authored-By: Claude Sonnet 5 --- .../src/stores/simple_map_store/common.rs | 68 +++- .../src/stores/simple_map_store/global.rs | 47 +-- .../src/stores/simple_map_store/per_key.rs | 13 +- .../src/tests/store_correctness_tests.rs | 384 ++++++++++++++++++ 4 files changed, 464 insertions(+), 48 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 4da34a5..48a3cdd 100644 --- a/asap-query-engine/src/stores/simple_map_store/common.rs +++ b/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, RwLock}; pub type MetricID = u32; pub type EpochID = u64; @@ -77,10 +77,17 @@ 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 +/// invalidates it via `RwLock::get_mut` (one pointer-width write, no lock acquisition since +/// `&mut self` already proves exclusive access); 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 behind its own inner `RwLock`, independent of the outer per-shard lock +/// callers take to reach a `MutableEpoch`. That lets `exact_query` take `&self` and use +/// double-checked locking to build the index at most once per invalidation, so concurrent +/// exact queries only ever need shared (read) access to the containing epoch — see +/// `exact_query` below. (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 /// bucket — the common case in ordered TSDB ingestion) skip the `windows_set` HashSet probe @@ -102,7 +109,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>>, + // Behind its own RwLock (independent of the outer per-shard lock) so exact_query can + // take &self and build it under double-checked locking (issue #607). + window_to_ids: RwLock>>>, /// Epoch time bounds for O(1) skip check, updated incrementally on insert. min_start: Option, @@ -123,7 +132,7 @@ impl MutableEpoch { aggregates_col: Vec::with_capacity(cap), windows_set: HashSet::new(), last_window: None, - window_to_ids: None, + window_to_ids: RwLock::new(None), min_start: None, max_end: None, } @@ -171,8 +180,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. get_mut() needs &mut RwLock, which + // &mut self already grants exclusively — no lock acquisition on this path. + *self.window_to_ids.get_mut().unwrap() = None; 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 +241,55 @@ impl MutableEpoch { /// The offset index (`HashMap>`) 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 `RwLock::get_mut`. /// - /// 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 behind its own inner `RwLock`, independent of the outer + /// per-shard lock callers hold to reach this epoch. Double-checked locking (read-check, + /// then write-and-recheck) means the index is built at most once per invalidation even + /// under concurrent callers, and callers only need shared (read) access to the epoch — + /// see the `window_to_ids` field doc for why this is safe (issue #607). pub fn exact_query( - &mut self, + &self, range: TimestampRange, ) -> Option)>> { - if self.window_to_ids.is_none() { + // Fast path: index already built (the common case after the first call). + if let Some(map) = self.window_to_ids.read().unwrap().as_ref() { + return Self::lookup_offsets(map, range, &self.metric_ids_col, &self.aggregates_col); + } + + // Slow path: build under the inner write lock. Re-check after acquiring it — + // another thread may have built the index while we were waiting. + let mut guard = self.window_to_ids.write().unwrap(); + if guard.is_none() { let mut idx: HashMap> = 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); + *guard = Some(idx); } - let offsets = self.window_to_ids.as_ref().unwrap().get(&range)?; + Self::lookup_offsets( + guard.as_ref().unwrap(), + 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>, + range: TimestampRange, + metric_ids_col: &[MetricID], + aggregates_col: &[Arc], + ) -> Option)>> { + 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 +317,7 @@ impl MutableEpoch { } // Invalidate lazy index and monotonic fast-path hint. - self.window_to_ids = None; + *self.window_to_ids.get_mut().unwrap() = None; self.last_window = None; // Recompute bounds (cleanup is rare; linear scan is fine). 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 d34181f..0cf2b95 100644 --- a/asap-query-engine/src/stores/simple_map_store/global.rs +++ b/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> = { - 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> = 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 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 98bf722..5166a63 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 @@ -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)>> = 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; diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index 1afbada..1716733 100644 --- a/asap-query-engine/src/tests/store_correctness_tests.rs +++ b/asap-query-engine/src/tests/store_correctness_tests.rs @@ -186,6 +186,14 @@ pub fn run_contract_suite(strategy: LockStrategy) { test_single_insert_exact_query_hit(strategy); test_single_insert_exact_query_wrong_start_returns_empty(strategy); test_single_insert_exact_query_wrong_end_returns_empty(strategy); + + // Exact-query cache correctness (Issue: query_precomputed_output_exact over-locking) + test_exact_query_is_stable_across_repeated_calls(strategy); + test_exact_query_sees_window_inserted_after_a_prior_exact_query(strategy); + test_exact_query_miss_then_hit_after_insert(strategy); + test_exact_query_correct_across_interleaved_inserts_and_queries(strategy); + test_exact_query_correct_after_epoch_rotation(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); @@ -345,6 +353,201 @@ fn test_single_insert_exact_query_wrong_end_returns_empty(strategy: LockStrategy ); } +// ── exact-query cache correctness ───────────────────────────────────────────── +// +// `query_precomputed_output_exact` is used by sliding-window instant queries to +// fetch a precompute with an exactly-matching timestamp range, no merging. Some +// implementations (PerKey's `MutableEpoch`) maintain a lazy internal lookup +// index for this path, built on first use and invalidated on the next insert. +// These tests pin the observable contract that must hold regardless of how +// (or whether) that caching is implemented: exact queries must always reflect +// the store's true current contents, never a stale snapshot from before the +// most recent write, and must remain correct across repeated calls. + +fn test_exact_query_is_stable_across_repeated_calls(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let (out, acc) = sum_entry(1, 1_000, 2_000, 42.0); + store.insert_precomputed_output(out, acc).unwrap(); + + let mut jsons = Vec::new(); + for call in 0..5 { + let result = store + .query_precomputed_output_exact("cpu_usage", 1, 1_000, 2_000) + .unwrap(); + assert_eq!( + total_bucket_count(&result), + 1, + "[{}] call #{call}: repeated exact query must keep finding the inserted window", + label(strategy) + ); + jsons.push(result.get(&None).unwrap()[0].1.serialize_to_json()); + } + assert!( + jsons.iter().all(|j| j == &jsons[0]), + "[{}] repeated exact queries for the same window must return identical values \ + across calls (any internal lookup cache must not corrupt results)", + label(strategy) + ); +} + +fn test_exact_query_sees_window_inserted_after_a_prior_exact_query(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let (out1, acc1) = sum_entry(1, 1_000, 2_000, 1.0); + store.insert_precomputed_output(out1, acc1).unwrap(); + + // First exact query — on implementations with a lazy lookup index (e.g. PerKey's + // `MutableEpoch::exact_query`), this call is what builds/populates that index. + let first = store + .query_precomputed_output_exact("cpu_usage", 1, 1_000, 2_000) + .unwrap(); + assert_eq!( + total_bucket_count(&first), + 1, + "[{}] sanity: first exact query must find the inserted window", + label(strategy) + ); + + // A brand-new window inserted after that must be visible to exact queries — + // whatever cache the query above populated must not shadow it. + let (out2, acc2) = sum_entry(1, 3_000, 4_000, 2.0); + store.insert_precomputed_output(out2, acc2).unwrap(); + + let second = store + .query_precomputed_output_exact("cpu_usage", 1, 3_000, 4_000) + .unwrap(); + assert_eq!( + total_bucket_count(&second), + 1, + "[{}] exact query must find a window inserted after a previous exact query \ + populated any internal cache — must not return stale/empty results", + label(strategy) + ); + + // The original window must still be correctly retrievable too. + let first_again = store + .query_precomputed_output_exact("cpu_usage", 1, 1_000, 2_000) + .unwrap(); + assert_eq!( + total_bucket_count(&first_again), + 1, + "[{}] original window must remain correctly retrievable after a later insert", + label(strategy) + ); +} + +fn test_exact_query_miss_then_hit_after_insert(strategy: LockStrategy) { + let store = make_store_simple(strategy); + + // Query for a window that doesn't exist yet. On PerKey this still builds + // (an empty-for-this-range) lazy index as a side effect. + let miss = store + .query_precomputed_output_exact("cpu_usage", 1, 5_000, 6_000) + .unwrap(); + assert!( + miss.is_empty(), + "[{}] query for a nonexistent window must be empty", + label(strategy) + ); + + // Now insert exactly that window. + let (out, acc) = sum_entry(1, 5_000, 6_000, 9.0); + store.insert_precomputed_output(out, acc).unwrap(); + + let hit = store + .query_precomputed_output_exact("cpu_usage", 1, 5_000, 6_000) + .unwrap(); + assert_eq!( + total_bucket_count(&hit), + 1, + "[{}] a window inserted after a prior miss must be found on the next exact query \ + — any cache built by the miss must be invalidated by the insert", + label(strategy) + ); +} + +fn test_exact_query_correct_across_interleaved_inserts_and_queries(strategy: LockStrategy) { + let store = make_store_simple(strategy); + let n = 30u64; + for i in 0..n { + let (out, acc) = sum_entry(1, i * 1_000, (i + 1) * 1_000, i as f64); + store.insert_precomputed_output(out, acc).unwrap(); + + // After each insert, re-verify every window inserted so far — including this + // one — is retrievable with its own correct value. This interleaves cache + // invalidation (insert) with cache use (exact query) on every iteration, and + // checks that a rebuilt lookup index never mixes up offsets across windows. + for j in 0..=i { + let result = store + .query_precomputed_output_exact("cpu_usage", 1, j * 1_000, (j + 1) * 1_000) + .unwrap(); + assert_eq!( + total_bucket_count(&result), + 1, + "[{}] window {j} must be retrievable after inserting window {i}", + label(strategy) + ); + let expected = SumAccumulator::with_sum(j as f64).serialize_to_json(); + let actual = result.get(&None).unwrap()[0].1.serialize_to_json(); + assert_eq!( + actual, + expected, + "[{}] window {j} must return its own value, not another window's, \ + after inserting window {i}", + label(strategy) + ); + } + } +} + +fn test_exact_query_correct_after_epoch_rotation(strategy: LockStrategy) { + // capacity=2, max_epochs=4 (hardcoded in StoreKeyData/PerKeyState) => + // retention_limit = 8. Inserting 10 windows evicts exactly the oldest 2 + // (windows 0 and 1, both in the oldest sealed epoch), leaving windows + // 2..10 spread across sealed epochs and the current (still-open) epoch. + let store = make_store( + strategy, + CleanupPolicy::CircularBuffer, + &[(1, AggregationType::Sum, Some(2), None)], + ); + let n = 10u64; + for i in 0..n { + let (out, acc) = sum_entry(1, i * 60_000, (i + 1) * 60_000, i as f64); + store.insert_precomputed_output(out, acc).unwrap(); + } + + for i in 0u64..2 { + let evicted = store + .query_precomputed_output_exact("cpu_usage", 1, i * 60_000, (i + 1) * 60_000) + .unwrap(); + assert!( + evicted.is_empty(), + "[{}] window {i} must have been evicted by circular-buffer rotation", + label(strategy) + ); + } + + for i in 2u64..n { + let result = store + .query_precomputed_output_exact("cpu_usage", 1, i * 60_000, (i + 1) * 60_000) + .unwrap(); + assert_eq!( + total_bucket_count(&result), + 1, + "[{}] window {i} must be retrievable via exact query after epoch rotation", + label(strategy) + ); + let expected = SumAccumulator::with_sum(i as f64).serialize_to_json(); + let actual = result.get(&None).unwrap()[0].1.serialize_to_json(); + assert_eq!( + actual, + expected, + "[{}] window {i} must return its own value after epoch rotation, \ + not a value from a sealed epoch's stale offset", + label(strategy) + ); + } +} + // ── batch insert correctness ────────────────────────────────────────────────── fn test_batch_insert_full_range_query_returns_all(strategy: LockStrategy) { @@ -1046,6 +1249,187 @@ fn test_concurrent_reads_return_complete_results(strategy: LockStrategy) { } } +// ── lock-contention characterization (PerKey only) ───────────────────────────── +// +// `SimpleMapStoreGlobal` has no lock granularity to have a bug in (one giant +// `Mutex` for everything), so this section targets `LockStrategy::PerKey` only. + +/// Characterizes a known over-locking bug in `SimpleMapStorePerKey`: +/// `query_precomputed_output_exact` takes the shard's `RwLock` as a *write* +/// lock (see `per_key.rs`, driven by `MutableEpoch::exact_query` taking +/// `&mut self` to lazily build/cache an internal lookup index) even though it +/// never mutates any queryable data. `query_precomputed_output` correctly +/// takes only a `.read()` lock on the same shard. +/// +/// Consequently, a long-running exact query currently blocks — rather than +/// runs concurrently with — a cheap, unrelated range query on the same +/// aggregation shard. +/// +/// # Method +/// +/// 1. Insert a large number of distinct windows into aggregation_id=1's +/// current (still-open) epoch, then insert one more to guarantee the +/// lazy `window_to_ids` lookup index is invalidated. This makes the next +/// exact query pay a full O(current-epoch-size) index rebuild under +/// whatever lock it takes — a large, unambiguous, easily measured +/// critical section (tens of milliseconds), instead of a cheap cache hit. +/// 2. Spawn reader threads that continuously issue range queries for a time +/// range with **no overlap with any inserted data**. Per `per_key.rs`, +/// this still requires acquiring the shard's lock (the `DashMap` entry +/// exists), but the scan itself is skipped via an O(1) time-bounds check +/// — so each call's *uncontended* cost is on the order of microseconds, +/// regardless of how much data the shard holds. +/// 3. While readers are looping, issue the single expensive exact query and +/// record its duration. +/// 4. Track the maximum single-call latency observed by any reader across +/// the whole run. +/// +/// A cheap, non-blocked reader call should never take anywhere near as long +/// as the exact query's own multi-millisecond lock hold — regardless of that +/// duration — because a read lock only excludes writers, not other readers. +/// If exact queries instead exclude readers (the bug), at least one reader +/// call will be observed stalled for a duration comparable to the exact +/// query's, since it has to wait out the writer's turn before proceeding. +/// +/// # Expected result on the current (pre-fix) implementation +/// +/// FAILS: `max_reader_latency` comes out a large fraction of `exact_duration` +/// (in practice, close to 100% — some reader gets stuck waiting the entire +/// time) instead of staying near the reader's own uncontended cost. This +/// assertion encodes the *desired* post-fix behavior, so it is expected to +/// start passing once `query_precomputed_output_exact` no longer requires a +/// write lock for this path. +/// +/// # Flakiness note +/// +/// This is a timing-based test. `n_windows` is chosen large enough that the +/// forced index rebuild takes tens of milliseconds — comfortably above +/// normal OS scheduling jitter (typically sub-millisecond to a few ms) — so +/// the 20% threshold has a wide margin in both directions. Extreme host +/// contention (e.g. a heavily oversubscribed CI runner) could in principle +/// still perturb timing; if this test flakes, prefer raising `n_windows` or +/// loosening the threshold over deleting it. +#[test] +fn test_exact_query_does_not_block_concurrent_range_queries_per_key() { + use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + use std::time::{Duration, Instant}; + + let n_windows = 500_000u64; + let store = Arc::new(make_store( + LockStrategy::PerKey, + CleanupPolicy::NoCleanup, + &[(1, AggregationType::Sum, None, None)], + )); + let batch: Vec<_> = (0..n_windows) + .map(|i| sum_entry(1, i * 10, i * 10 + 1, i as f64)) + .collect(); + store.insert_precomputed_output_batch(batch).unwrap(); + + // One more insert to guarantee the next exact query's lazy index is + // invalidated and must be rebuilt from scratch under its lock. + let extra_ts = n_windows * 10; + let (out, acc) = sum_entry(1, extra_ts, extra_ts + 1, 0.0); + store.insert_precomputed_output(out, acc).unwrap(); + + // Sanity: the exact query must still find its window (also serves as a + // warm-up call, though it is not the one that gets timed below). + let sanity = store + .query_precomputed_output_exact("cpu_usage", 1, extra_ts, extra_ts + 1) + .unwrap(); + assert_eq!(total_bucket_count(&sanity), 1); + + // A second insert to invalidate the index again for the timed call. + let extra_ts2 = extra_ts + 1; + let (out2, acc2) = sum_entry(1, extra_ts2, extra_ts2 + 1, 0.0); + store.insert_precomputed_output(out2, acc2).unwrap(); + + // Readers query a range far outside all inserted data (max inserted + // timestamp is extra_ts2 + 1), so the per-epoch time-bounds check skips + // any real scan — this call is cheap purely from lock acquisition + + // bookkeeping, independent of shard size. + let far_start = u64::MAX - 1_000; + let far_end = u64::MAX; + + let stop = Arc::new(AtomicBool::new(false)); + let max_reader_latency_ns = Arc::new(AtomicU64::new(0)); + let n_readers = 4; + + let reader_handles: Vec<_> = (0..n_readers) + .map(|_| { + let store = store.clone(); + let stop = stop.clone(); + let max_latency = max_reader_latency_ns.clone(); + std::thread::spawn(move || { + let mut iters = 0u64; + while !stop.load(Ordering::Relaxed) { + let call_start = Instant::now(); + let result = store + .query_precomputed_output("cpu_usage", 1, far_start, far_end) + .unwrap(); + debug_assert!(result.is_empty()); + let elapsed_ns = call_start.elapsed().as_nanos() as u64; + max_latency.fetch_max(elapsed_ns, Ordering::Relaxed); + iters += 1; + } + iters + }) + }) + .collect(); + + // Let readers start looping before the exact query fires. + std::thread::sleep(Duration::from_millis(20)); + + let exact_start = Instant::now(); + let result = store + .query_precomputed_output_exact("cpu_usage", 1, extra_ts2, extra_ts2 + 1) + .unwrap(); + let exact_duration = exact_start.elapsed(); + assert_eq!( + total_bucket_count(&result), + 1, + "sanity: the timed exact query must still find its window" + ); + + // Let readers keep looping briefly after, then stop them. + std::thread::sleep(Duration::from_millis(20)); + stop.store(true, Ordering::Relaxed); + let total_reader_iters: u64 = reader_handles.into_iter().map(|h| h.join().unwrap()).sum(); + + let max_reader_latency = Duration::from_nanos(max_reader_latency_ns.load(Ordering::Relaxed)); + + eprintln!( + "[PerKey exact-query lock contention] exact_duration={:?}, \ + max_reader_latency={:?}, total_reader_iters={total_reader_iters} \ + (ratio max_reader_latency/exact_duration = {:.3})", + exact_duration, + max_reader_latency, + max_reader_latency.as_secs_f64() / exact_duration.as_secs_f64().max(1e-12) + ); + + // Sanity floor: make sure the forced rebuild actually took long enough + // for this to be a meaningful signal (not swallowed by noise). + assert!( + exact_duration >= Duration::from_millis(1), + "exact query completed in {:?}, too fast to reliably characterize lock \ + contention — increase n_windows", + exact_duration + ); + + assert!( + max_reader_latency.as_nanos() * 5 <= exact_duration.as_nanos(), + "a concurrent range query was stalled for {:?} while a single exact query \ + ran for {:?} (ratio {:.3}, allowed <= 0.20) — this indicates \ + query_precomputed_output_exact is blocking concurrent range queries on \ + the same PerKey shard for the duration of its lock hold. This failure is \ + expected on the current implementation (which takes a write lock for \ + exact queries even though it doesn't mutate queryable data); it should \ + start passing once that lock is narrowed to a read lock.", + max_reader_latency, + exact_duration, + max_reader_latency.as_secs_f64() / exact_duration.as_secs_f64().max(1e-12) + ); +} + // ── test entry points ───────────────────────────────────────────────────────── /// Contract suite against `SimpleMapStore` with [`LockStrategy::PerKey`]. From 34658856d65f90c7df60414c9c5ce857cf04eb8c Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 15:23:34 -0400 Subject: [PATCH 2/3] fix(query-engine): use OnceLock instead of hand-rolled RwLock cache Code review on #611 flagged two issues with the RwLock> approach: - It reimplements double-checked locking that std::sync::OnceLock already provides via get_or_init. - It introduces a poisoning-cascade not present before the PR: a panic while holding the inner write lock would poison window_to_ids specifically, and every subsequent exact_query on that epoch would then panic on .unwrap() forever (until the epoch reseals), unlike the outer per-shard lock's poisoning, which call sites already convert gracefully into a Result::Err. OnceLock has no poisoning concept: if the init closure ever panicked, the cell would simply stay uninitialized for the next caller to retry. Switching to it removes both the reimplemented locking logic and the new failure mode in one change. Invalidation on insert/remove_windows now uses OnceLock::take() (still through &mut self, still no synchronization needed) instead of writing None through a lock guard. Co-Authored-By: Claude Sonnet 5 --- .../src/stores/simple_map_store/common.rs | 71 ++++++++----------- 1 file changed, 30 insertions(+), 41 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 48a3cdd..d77142f 100644 --- a/asap-query-engine/src/stores/simple_map_store/common.rs +++ b/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, RwLock}; +use std::sync::{Arc, OnceLock}; pub type MetricID = u32; pub type EpochID = u64; @@ -77,16 +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` -/// invalidates it via `RwLock::get_mut` (one pointer-width write, no lock acquisition since -/// `&mut self` already proves exclusive access); 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 behind its own inner `RwLock`, independent of the outer per-shard lock -/// callers take to reach a `MutableEpoch`. That lets `exact_query` take `&self` and use -/// double-checked locking to build the index at most once per invalidation, so concurrent -/// exact queries only ever need shared (read) access to the containing epoch — see -/// `exact_query` below. (Issue #607.) +/// 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 @@ -109,9 +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. - // Behind its own RwLock (independent of the outer per-shard lock) so exact_query can - // take &self and build it under double-checked locking (issue #607). - window_to_ids: RwLock>>>, + // 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>>, /// Epoch time bounds for O(1) skip check, updated incrementally on insert. min_start: Option, @@ -132,7 +134,7 @@ impl MutableEpoch { aggregates_col: Vec::with_capacity(cap), windows_set: HashSet::new(), last_window: None, - window_to_ids: RwLock::new(None), + window_to_ids: OnceLock::new(), min_start: None, max_end: None, } @@ -180,9 +182,9 @@ impl MutableEpoch { self.metric_ids_col.push(metric_id); self.aggregates_col.push(agg); - // Opt 1: invalidate lazy index at zero cost. get_mut() needs &mut RwLock, which - // &mut self already grants exclusively — no lock acquisition on this path. - *self.window_to_ids.get_mut().unwrap() = 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))); @@ -241,39 +243,26 @@ impl MutableEpoch { /// The offset index (`HashMap>`) 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 invalidates it via `RwLock::get_mut`. + /// until the next `insert`, which invalidates it via `OnceLock::take`. /// - /// Takes `&self`: the index lives behind its own inner `RwLock`, independent of the outer - /// per-shard lock callers hold to reach this epoch. Double-checked locking (read-check, - /// then write-and-recheck) means the index is built at most once per invalidation even - /// under concurrent callers, and callers only need shared (read) access to the epoch — - /// see the `window_to_ids` field doc for why this is safe (issue #607). + /// 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( &self, range: TimestampRange, ) -> Option)>> { - // Fast path: index already built (the common case after the first call). - if let Some(map) = self.window_to_ids.read().unwrap().as_ref() { - return Self::lookup_offsets(map, range, &self.metric_ids_col, &self.aggregates_col); - } - - // Slow path: build under the inner write lock. Re-check after acquiring it — - // another thread may have built the index while we were waiting. - let mut guard = self.window_to_ids.write().unwrap(); - if guard.is_none() { + let index = self.window_to_ids.get_or_init(|| { let mut idx: HashMap> = 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); } - *guard = Some(idx); - } - Self::lookup_offsets( - guard.as_ref().unwrap(), - range, - &self.metric_ids_col, - &self.aggregates_col, - ) + 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. @@ -317,7 +306,7 @@ impl MutableEpoch { } // Invalidate lazy index and monotonic fast-path hint. - *self.window_to_ids.get_mut().unwrap() = None; + self.window_to_ids.take(); self.last_window = None; // Recompute bounds (cleanup is rare; linear scan is fine). From 91f323ae3ca44fe50bfb6735a9cec936346ac096 Mon Sep 17 00:00:00 2001 From: Milind Srivastava Date: Tue, 25 Aug 2026 16:17:08 -0400 Subject: [PATCH 3/3] fix(query-engine): address review findings on #611 - INDEX_DESIGN.md still documented the pre-fix write-lock/&mut-self exact_query design (Option, write lock in the concurrency table and query-mechanics walkthrough). Updated to describe the OnceLock-based design and the read-lock-only concurrency table, with a note on why exact queries no longer need a write lock. - test_exact_query_does_not_block_concurrent_range_queries_per_key computed max_reader_latency_ns but never asserted it was actually touched. If every reader thread got starved of scheduler time during the timed window, the value would stay at its initial 0 and the ratio assertion (0 <= anything) would pass vacuously, having measured no contention at all. Added an explicit assertion that readers recorded at least one iteration each and that max_reader_latency_ns is nonzero before trusting the ratio. Co-Authored-By: Claude Sonnet 5 --- .../stores/simple_map_store/INDEX_DESIGN.md | 34 ++++++++++++------- .../src/tests/store_correctness_tests.rs | 12 +++++++ 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/asap-query-engine/src/stores/simple_map_store/INDEX_DESIGN.md b/asap-query-engine/src/stores/simple_map_store/INDEX_DESIGN.md index 594e6a9..288af58 100644 --- a/asap-query-engine/src/stores/simple_map_store/INDEX_DESIGN.md +++ b/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 - // Lazy offset index (Opt 1 + 2): built on first exact_query, None after any insert - window_to_ids: Option>> + // Lazy offset index (Opt 1 + 2): built on first exact_query, cleared on any insert + window_to_ids: OnceLock>> // Epoch bounds for O(1) skip check (updated incrementally on insert) min_start: Option @@ -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)>`, 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. diff --git a/asap-query-engine/src/tests/store_correctness_tests.rs b/asap-query-engine/src/tests/store_correctness_tests.rs index 1716733..e913e6a 100644 --- a/asap-query-engine/src/tests/store_correctness_tests.rs +++ b/asap-query-engine/src/tests/store_correctness_tests.rs @@ -1406,6 +1406,18 @@ fn test_exact_query_does_not_block_concurrent_range_queries_per_key() { max_reader_latency.as_secs_f64() / exact_duration.as_secs_f64().max(1e-12) ); + // Sanity floor: make sure reader threads actually ran and recorded a sample. + // Without this, a run where every reader got starved of scheduler time during + // the timed window would leave max_reader_latency_ns at its initial 0, and the + // ratio assertion below would pass vacuously (0 <= anything) without having + // measured contention at all. + assert!( + total_reader_iters >= n_readers as u64 && max_reader_latency_ns.load(Ordering::Relaxed) > 0, + "no reader thread recorded a latency sample (total_reader_iters={total_reader_iters}) \ + — readers may have been starved of scheduler time during this run, so it \ + doesn't validate anything; rerun, or investigate scheduler contention on this host" + ); + // Sanity floor: make sure the forced rebuild actually took long enough // for this to be a meaningful signal (not swallowed by noise). assert!(