From 2d241c660a2754303f02ce3cfdba3461a376044b Mon Sep 17 00:00:00 2001 From: Ashish Priyadarshi Date: Sun, 31 May 2026 20:09:42 -0700 Subject: [PATCH 1/2] Fixes to locking mechanism --- src/index.rs | 170 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/server.rs | 71 +++++++++------------ 2 files changed, 200 insertions(+), 41 deletions(-) diff --git a/src/index.rs b/src/index.rs index c4f8091..dc912e9 100644 --- a/src/index.rs +++ b/src/index.rs @@ -870,6 +870,51 @@ impl SemanticIndex { } } + /// M29 two-level conversation lookup under a **single** read lock. + /// + /// Replaces the double-lock pattern in `process_query_locally` where the + /// handler acquired `state.index.read()` once for the conversation namespace + /// and then — if that missed — released the lock and re-acquired it for + /// the base namespace fallback. That open window between lock 1 and lock 2 + /// allowed a concurrent write to alter index state, meaning the two reads + /// observed inconsistent snapshots of the index (TOCTOU). + /// + /// By running both namespace lookups inside the same `&self` borrow (the + /// caller holds the `RwLock` read guard for the duration), this function + /// gives a single consistent snapshot: if the conv namespace is empty now, + /// the base namespace result is from the same instant, not a later one. + /// + /// Returns `(hit, hit_namespace, scope)`: + /// - `hit` — the `QueryHit` if found, `None` on miss + /// - `hit_namespace` — the namespace that produced the result; always set + /// so the caller can pass it directly to `record_access` and metrics + /// - `scope` — `Some("conversation")` | `Some("global")` | `None` (miss) + pub fn query_two_level( + &self, + embedding: &[f32], + threshold: f32, + conv_ns: &str, + base_ns: &str, + query_text: Option<&str>, + now: u64, + ) -> Result<(Option, String, Option<&'static str>)> { + // Conversation namespace has priority. + if let Some(ns) = self.namespaces.get(conv_ns) { + if let Some(hit) = ns.query_with_text(embedding, threshold, query_text, now)? { + return Ok((Some(hit), conv_ns.to_string(), Some("conversation"))); + } + } + // Fallback: base namespace. + match self.namespaces.get(base_ns) { + Some(ns) => { + let hit = ns.query_with_text(embedding, threshold, query_text, now)?; + let scope = if hit.is_some() { Some("global") } else { None }; + Ok((hit, base_ns.to_string(), scope)) + } + None => Ok((None, base_ns.to_string(), None)), + } + } + pub fn entry_count(&self) -> usize { self.namespaces.values().map(|n| n.entry_count()).sum() } @@ -2242,4 +2287,129 @@ mod tests { ); assert!(idx.namespaces.contains_key(&conv_ns)); } + + // --- query_two_level unit tests ---------------------------------------- + + #[test] + fn test_query_two_level_conv_hit() { + // Entry in conv namespace → returns scope="conversation", conv ns name. + let mut idx = test_index(); + let base = base_namespace("m::3", None); + let conv = conversation_namespace("m::3", None, "s1"); + idx.insert(vec![1.0_f32, 0.0, 0.0], "base-r".into(), "q".into(), &base) + .unwrap(); + idx.insert(vec![1.0_f32, 0.0, 0.0], "conv-r".into(), "q".into(), &conv) + .unwrap(); + let now = now_unix_secs(); + let (hit, ns, scope) = idx + .query_two_level(&[1.0_f32, 0.0, 0.0], 0.9, &conv, &base, None, now) + .unwrap(); + assert_eq!(hit.unwrap().response, "conv-r"); + assert_eq!(ns, conv); + assert_eq!(scope, Some("conversation")); + } + + #[test] + fn test_query_two_level_base_fallback() { + // Conv namespace is empty → falls through to base, returns scope="global". + let mut idx = test_index(); + let base = base_namespace("m::3", None); + let conv = conversation_namespace("m::3", None, "s1"); + idx.insert(vec![1.0_f32, 0.0, 0.0], "base-r".into(), "q".into(), &base) + .unwrap(); + let now = now_unix_secs(); + let (hit, ns, scope) = idx + .query_two_level(&[1.0_f32, 0.0, 0.0], 0.9, &conv, &base, None, now) + .unwrap(); + assert_eq!(hit.unwrap().response, "base-r"); + assert_eq!(ns, base); + assert_eq!(scope, Some("global")); + } + + #[test] + fn test_query_two_level_both_miss() { + // Nothing in either namespace → None, hit_ns is base, scope is None. + let idx = test_index(); + let base = base_namespace("m::3", None); + let conv = conversation_namespace("m::3", None, "s1"); + let now = now_unix_secs(); + let (hit, ns, scope) = idx + .query_two_level(&[1.0_f32, 0.0, 0.0], 0.9, &conv, &base, None, now) + .unwrap(); + assert!(hit.is_none()); + assert_eq!(ns, base); + assert_eq!(scope, None); + } + + #[test] + fn test_query_two_level_conv_priority_over_base() { + // Distinct responses in both namespaces: conv always wins when both + // have the vector, regardless of insertion order. + let mut idx = test_index(); + let base = base_namespace("m::3", None); + let conv = conversation_namespace("m::3", None, "s1"); + idx.insert(vec![1.0_f32, 0.0, 0.0], "from-base".into(), "q".into(), &base) + .unwrap(); + idx.insert(vec![0.8_f32, 0.6, 0.0], "from-conv".into(), "q2".into(), &conv) + .unwrap(); + let now = now_unix_secs(); + let (hit, _, scope) = idx + .query_two_level(&[1.0_f32, 0.0, 0.0], 0.5, &conv, &base, None, now) + .unwrap(); + // conv has a nearby vector; base has an exact match — conv is checked + // first so "from-conv" is returned. + assert_eq!(hit.unwrap().response, "from-conv"); + assert_eq!(scope, Some("conversation")); + } + + #[test] + fn test_query_two_level_exact_match_in_conv() { + // Exact-match pre-filter works through query_two_level: text match in + // conv namespace hits even with an orthogonal embedding. + let mut idx = test_index(); + let base = base_namespace("m::3", None); + let conv = conversation_namespace("m::3", None, "s1"); + idx.insert( + vec![1.0_f32, 0.0, 0.0], + "Paris".into(), + "Capital of France?".into(), + &conv, + ) + .unwrap(); + let now = now_unix_secs(); + let (hit, _, scope) = idx + .query_two_level( + &[0.0_f32, 0.0, 1.0], // orthogonal — HNSW would miss + 0.99, + &conv, + &base, + Some("Capital of France?"), + now, + ) + .unwrap(); + let h = hit.unwrap(); + assert!(h.exact_match); + assert_eq!(h.response, "Paris"); + assert_eq!(scope, Some("conversation")); + } + + #[test] + fn test_query_two_level_error_propagates() { + // Dimension mismatch in the conv namespace propagates as Err rather + // than silently falling through to the base namespace. + let mut idx = test_index(); + let base = base_namespace("m::3", None); + let conv = conversation_namespace("m::3", None, "s1"); + // Seed conv with a 3-d entry so the namespace knows its dimension. + idx.insert(vec![1.0_f32, 0.0, 0.0], "r".into(), "q".into(), &conv) + .unwrap(); + let now = now_unix_secs(); + // Query with a 4-d vector — dimension mismatch inside conv namespace. + let result = + idx.query_two_level(&[1.0_f32, 0.0, 0.0, 0.0], 0.9, &conv, &base, None, now); + assert!( + result.is_err(), + "dimension mismatch must propagate as Err, not fall through to base" + ); + } } diff --git a/src/server.rs b/src/server.rs index 307c1e4..528b28a 100644 --- a/src/server.rs +++ b/src/server.rs @@ -231,54 +231,43 @@ async fn process_query_locally( let conv_ns = conv_id.map(|cid| conversation_namespace(model_id, req.cache_scope.as_deref(), cid)); - // First-level lookup: conversation namespace if present, else base. - let primary_ns: &str = conv_ns.as_deref().unwrap_or(&base); - let primary_result = { + // Acquire the index read lock exactly once — a single consistent snapshot + // covers both the conversation namespace check and the base fallback. + // The previous double-lock pattern released and re-acquired between the + // two lookups, opening a TOCTOU window where a concurrent WAL flush could + // insert into the conversation namespace and be missed, or where an + // expiry/eviction between the two acquisitions could yield an inconsistent + // result. query_two_level runs both lookups under the same &self borrow, + // eliminating the window and halving the lock-acquire count on the hot path. + let (final_result, hit_ns_owned, hit_scope): ( + Result, anyhow::Error>, + String, + Option<&'static str>, + ) = if let Some(ref cns) = conv_ns { let index = state.index.read().await; - index.query_with_text( + match index.query_two_level( &req.embedding, req.threshold, - primary_ns, + cns, + &base, req.query_text.as_deref(), now, - ) - }; - - // Fallback (only when we had a conversation_id and the conversation - // namespace missed). Errors from the primary still propagate as 400. - let (final_result, hit_ns_owned, hit_scope): ( - Result, anyhow::Error>, - String, - Option<&'static str>, - ) = match primary_result { - Ok(Some(hit)) => { - let scope = if conv_ns.is_some() { - Some("conversation") - } else { - None - }; - (Ok(Some(hit)), primary_ns.to_string(), scope) + ) { + Ok((hit, ns, scope)) => (Ok(hit), ns, scope), + Err(e) => (Err(e), base.clone(), None), } - Ok(None) if conv_ns.is_some() => { - // Two-level fallback: try the base namespace. - let fallback = { - let index = state.index.read().await; - index.query_with_text( - &req.embedding, - req.threshold, - &base, - req.query_text.as_deref(), - now, - ) - }; - match fallback { - Ok(Some(hit)) => (Ok(Some(hit)), base.clone(), Some("global")), - Ok(None) => (Ok(None), base.clone(), None), - Err(e) => (Err(e), base.clone(), None), - } + } else { + let index = state.index.read().await; + match index.query_with_text( + &req.embedding, + req.threshold, + &base, + req.query_text.as_deref(), + now, + ) { + Ok(hit) => (Ok(hit), base.clone(), None), + Err(e) => (Err(e), base.clone(), None), } - Ok(None) => (Ok(None), primary_ns.to_string(), None), - Err(e) => (Err(e), primary_ns.to_string(), None), }; let elapsed = start.elapsed().as_secs_f64(); From 683f5d3b7d4d99debd2be21cd2764844d2fe987f Mon Sep 17 00:00:00 2001 From: Ashish Priyadarshi Date: Sun, 31 May 2026 20:22:47 -0700 Subject: [PATCH 2/2] Fix to formatting --- src/index.rs | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/src/index.rs b/src/index.rs index dc912e9..0dd9262 100644 --- a/src/index.rs +++ b/src/index.rs @@ -2348,10 +2348,20 @@ mod tests { let mut idx = test_index(); let base = base_namespace("m::3", None); let conv = conversation_namespace("m::3", None, "s1"); - idx.insert(vec![1.0_f32, 0.0, 0.0], "from-base".into(), "q".into(), &base) - .unwrap(); - idx.insert(vec![0.8_f32, 0.6, 0.0], "from-conv".into(), "q2".into(), &conv) - .unwrap(); + idx.insert( + vec![1.0_f32, 0.0, 0.0], + "from-base".into(), + "q".into(), + &base, + ) + .unwrap(); + idx.insert( + vec![0.8_f32, 0.6, 0.0], + "from-conv".into(), + "q2".into(), + &conv, + ) + .unwrap(); let now = now_unix_secs(); let (hit, _, scope) = idx .query_two_level(&[1.0_f32, 0.0, 0.0], 0.5, &conv, &base, None, now) @@ -2405,8 +2415,7 @@ mod tests { .unwrap(); let now = now_unix_secs(); // Query with a 4-d vector — dimension mismatch inside conv namespace. - let result = - idx.query_two_level(&[1.0_f32, 0.0, 0.0, 0.0], 0.9, &conv, &base, None, now); + let result = idx.query_two_level(&[1.0_f32, 0.0, 0.0, 0.0], 0.9, &conv, &base, None, now); assert!( result.is_err(), "dimension mismatch must propagate as Err, not fall through to base"