perf(diskann): reuse VisitedSet across searches, add search_with and allocation regression test

crates/ruvector-diskann/examples/bench_search_hot.rs

@@ -0,0 +1,83 @@
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use ruvector_diskann::{DiskAnnConfig, DiskAnnIndex};
+use std::hint::black_box;
+use std::time::Instant;
+
+const DIM: usize = 128;
+const QUERY_COUNT: usize = 2_000;
+const WARMUP_COUNT: usize = 100;
+const K: usize = 10;
+
+fn random_unit_vector(rng: &mut StdRng) -> Vec<f32> {
+    let mut vector: Vec<f32> = (0..DIM).map(|_| rng.gen_range(-1.0..1.0)).collect();
+    let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
+    for value in &mut vector {
+        *value /= norm;
+    }
+    vector
+}
+
+fn parse_n() -> usize {
+    let mut args = std::env::args().skip(1);
+    match args.next() {
+        Some(flag) if flag == "--n" => args
+            .next()
+            .expect("--n requires a value")
+            .parse()
+            .expect("n must be a positive integer"),
+        Some(value) => value.parse().expect("n must be a positive integer"),
+        None => 100_000,
+    }
+}
+
+fn main() {
+    let n = parse_n();
+    assert!(n > 1, "n must be greater than one");
+
+    let mut data_rng = StdRng::seed_from_u64(0x677D_15CA_77);
+    let mut index = DiskAnnIndex::new(DiskAnnConfig {
+        dim: DIM,
+        max_degree: 4,
+        build_beam: 8,
+        search_beam: 64,
+        alpha: 1.0,
+        ..Default::default()
+    });
+
+    let build_started = Instant::now();
+    for id in 0..n {
+        index
+            .insert(id.to_string(), random_unit_vector(&mut data_rng))
+            .expect("benchmark insert failed");
+    }
+    index.build().expect("benchmark build failed");
+    let build_seconds = build_started.elapsed().as_secs_f64();
+
+    let mut query_rng = StdRng::seed_from_u64(0x6770_0A11_CE);
+    let queries: Vec<Vec<f32>> = (0..QUERY_COUNT)
+        .map(|_| random_unit_vector(&mut query_rng))
+        .collect();
+
+    for query in queries.iter().take(WARMUP_COUNT) {
+        black_box(index.search(black_box(query), K).expect("warmup failed"));
+    }
+
+    let mut latencies_ns = Vec::with_capacity(QUERY_COUNT);
+    let run_started = Instant::now();
+    for query in &queries {
+        let query_started = Instant::now();
+        black_box(index.search(black_box(query), K).expect("search failed"));
+        latencies_ns.push(query_started.elapsed().as_nanos() as u64);
+    }
+    let elapsed = run_started.elapsed();
+
+    latencies_ns.sort_unstable();
+    let p50_us = latencies_ns[QUERY_COUNT / 2] as f64 / 1_000.0;
+    let p99_us = latencies_ns[(QUERY_COUNT * 99 / 100).min(QUERY_COUNT - 1)] as f64 / 1_000.0;
+    let qps = QUERY_COUNT as f64 / elapsed.as_secs_f64();
+
+    println!(
+        "RESULT n={n} d={DIM} k={K} queries={QUERY_COUNT} build_s={build_seconds:.3} p50_us={p50_us:.3} p99_us={p99_us:.3} qps={qps:.3}"
+    );
+}

crates/ruvector-diskann/src/distance.rs

@@ -199,7 +199,30 @@ impl VisitedSet {
     /// Reset for a new search — O(1) via generation counter
     #[inline]
     pub fn clear(&mut self) {
-        self.generation += 1;
+        if self.generation == u64::MAX {
+            self.bits.fill(0);
+            self.gens.fill(0);
+            self.generation = 1;
+        } else {
+            self.generation += 1;
+        }
+    }
+
+    /// Prepare this set for an index of `n` nodes.
+    ///
+    /// A size mismatch reinitializes the backing storage. Repeated use with the
+    /// same index size takes the O(1) generation-counter path in [`Self::clear`].
+    #[inline]
+    pub(crate) fn prepare(&mut self, n: usize) {
+        if self.gens.len() != n {
+            self.bits.resize((n + 63) / 64, 0);
+            self.bits.fill(0);
+            self.gens.resize(n, 0);
+            self.gens.fill(0);
+            self.generation = 1;
+        } else {
+            self.clear();
+        }
     }
 
     /// Mark node as visited
@@ -341,6 +364,31 @@ mod tests {
         assert!(vs.contains(43));
     }
 
+    #[test]
+    fn test_visited_set_generation_wrap() {
+        let mut vs = VisitedSet::new(100);
+        vs.generation = u64::MAX;
+        vs.insert(42);
+
+        vs.clear();
+
+        assert_eq!(vs.generation, 1);
+        assert!(!vs.contains(42));
+    }
+
+    #[test]
+    fn test_visited_set_size_mismatch_reinitializes() {
+        let mut vs = VisitedSet::new(2);
+        vs.insert(1);
+
+        vs.prepare(100);
+
+        assert_eq!(vs.gens.len(), 100);
+        assert!(!vs.contains(1));
+        vs.insert(99);
+        assert!(vs.contains(99));
+    }
+
     #[test]
     fn test_pq_flat_table() {
         // 2 subspaces, 4 centroids each (k=4 for test)

crates/ruvector-diskann/src/graph.rs

@@ -141,7 +141,7 @@ impl VamanaGraph {
         beam_width: usize,
         visited: &mut VisitedSet,
     ) -> (Vec<u32>, usize) {
-        visited.clear();
+        visited.prepare(vectors.len());
 
         let mut candidates = BinaryHeap::<Candidate>::new();
         let mut best = BinaryHeap::<MaxCandidate>::new();

crates/ruvector-diskann/src/index.rs

@@ -9,6 +9,9 @@ use std::collections::HashMap;
 use std::fs::{self, File};
 use std::io::{BufWriter, Write};
 use std::path::{Path, PathBuf};
+use std::sync::{Mutex, TryLockError};
+
+const MAX_VISITED_POOL_SIZE: usize = 4;
 
 /// Search result
 #[derive(Debug, Clone)]
@@ -70,8 +73,8 @@ pub struct DiskAnnIndex {
     pq_codes: Vec<Vec<u8>>,
     /// Whether index has been built
     built: bool,
-    /// Reusable visited set for search (avoids per-query allocation)
-    visited: Option<VisitedSet>,
+    /// Small pool of reusable visited sets for shared-reference searches.
+    visited_pool: Mutex<Vec<VisitedSet>>,
     /// Memory-mapped vector data (for large datasets)
     mmap: Option<Mmap>,
 }
@@ -89,7 +92,7 @@ impl DiskAnnIndex {
             pq: None,
             pq_codes: Vec::new(),
             built: false,
-            visited: None,
+            visited_pool: Mutex::new(Vec::new()),
             mmap: None,
         }
     }
@@ -154,8 +157,12 @@ impl DiskAnnIndex {
         graph.build(&self.vectors)?;
         self.graph = Some(graph);
 
-        // Pre-allocate visited set for search
-        self.visited = Some(VisitedSet::new(n));
+        // Seed the single-threaded fast path. Concurrent searches allocate a
+        // fallback set when every pooled set is checked out.
+        *self
+            .visited_pool
+            .get_mut()
+            .unwrap_or_else(|poisoned| poisoned.into_inner()) = vec![VisitedSet::new(n)];
         self.built = true;
 
         if let Some(ref path) = self.config.storage_path {
@@ -167,20 +174,49 @@ impl DiskAnnIndex {
 
     /// Search for k nearest neighbors
     pub fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchResult>> {
-        if !self.built {
-            return Err(DiskAnnError::NotBuilt);
-        }
-        if query.len() != self.config.dim {
-            return Err(DiskAnnError::DimensionMismatch {
-                expected: self.config.dim,
-                actual: query.len(),
-            });
+        self.validate_search(query)?;
+
+        let mut visited = match self.visited_pool.try_lock() {
+            Ok(mut pool) => pool
+                .pop()
+                .unwrap_or_else(|| VisitedSet::new(self.vectors.len())),
+            Err(TryLockError::Poisoned(poisoned)) => poisoned
+                .into_inner()
+                .pop()
+                .unwrap_or_else(|| VisitedSet::new(self.vectors.len())),
+            Err(TryLockError::WouldBlock) => VisitedSet::new(self.vectors.len()),
+        };
+
+        let result = self.search_with(query, k, &mut visited);
+
+        let mut pool = self
+            .visited_pool
+            .lock()
+            .unwrap_or_else(|poisoned| poisoned.into_inner());
+        if pool.len() < MAX_VISITED_POOL_SIZE {
+            pool.push(visited);
         }
 
+        result
+    }
+
+    /// Search with caller-owned visited-state storage.
+    ///
+    /// Reuse the same [`VisitedSet`] for steady-state searches without visited
+    /// set allocation. If its size differs from this index, the set is resized
+    /// and fully reset before the search.
+    pub fn search_with(
+        &self,
+        query: &[f32],
+        k: usize,
+        visited: &mut VisitedSet,
+    ) -> Result<Vec<SearchResult>> {
+        self.validate_search(query)?;
+
         let graph = self.graph.as_ref().unwrap();
         let beam = self.config.search_beam.max(k);
 
-        let (candidates, _) = graph.greedy_search(&self.vectors, query, beam);
+        let (candidates, _) = graph.greedy_search_fast(&self.vectors, query, beam, visited);
 
         // Re-rank candidates with exact distance
         let mut scored: Vec<(u32, f32)> = candidates
@@ -199,6 +235,20 @@ impl DiskAnnIndex {
             .collect())
     }
 
+    fn validate_search(&self, query: &[f32]) -> Result<()> {
+        if !self.built {
+            return Err(DiskAnnError::NotBuilt);
+        }
+        if query.len() != self.config.dim {
+            return Err(DiskAnnError::DimensionMismatch {
+                expected: self.config.dim,
+                actual: query.len(),
+            });
+        }
+
+        Ok(())
+    }
+
     /// Get the number of vectors in the index
     pub fn count(&self) -> usize {
         self.vectors.len()
@@ -407,7 +457,7 @@ impl DiskAnnIndex {
             pq,
             pq_codes,
             built: true,
-            visited: Some(VisitedSet::new(n)),
+            visited_pool: Mutex::new(vec![VisitedSet::new(n)]),
             mmap: Some(mmap),
         })
     }
@@ -459,6 +509,64 @@ mod tests {
         assert!(results[0].distance < 1e-6); // Exact match
     }
 
+    #[test]
+    fn search_paths_return_identical_ids() {
+        use rand::prelude::*;
+
+        let dim = 32;
+        let n = 500;
+        let mut index = DiskAnnIndex::new(DiskAnnConfig {
+            dim,
+            max_degree: 16,
+            build_beam: 32,
+            search_beam: 32,
+            alpha: 1.2,
+            ..Default::default()
+        });
+        index.insert_batch(random_vectors(n, dim)).unwrap();
+        index.build().unwrap();
+
+        let mut query_rng = rand::rngs::StdRng::seed_from_u64(0x6775_EA2C);
+        let mut visited = VisitedSet::new(n);
+
+        for _ in 0..100 {
+            let query: Vec<f32> = (0..dim).map(|_| query_rng.gen()).collect();
+            let pooled_ids: Vec<String> = index
+                .search(&query, 10)
+                .unwrap()
+                .into_iter()
+                .map(|result| result.id)
+                .collect();
+            let caller_owned_ids: Vec<String> = index
+                .search_with(&query, 10, &mut visited)
+                .unwrap()
+                .into_iter()
+                .map(|result| result.id)
+                .collect();
+
+            // This mirrors the pre-fix search shape: allocate a VisitedSet in
+            // graph.greedy_search, then perform the same exact-distance rerank.
+            let graph = index.graph.as_ref().unwrap();
+            let beam = index.config.search_beam.max(10);
+            let (candidates, _) = graph.greedy_search(&index.vectors, &query, beam);
+            let mut scored: Vec<(u32, f32)> = candidates
+                .into_iter()
+                .map(|id| (id, l2_squared(index.vectors.get(id as usize), &query)))
+                .collect();
+            scored.sort_unstable_by(|a, b| {
+                a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)
+            });
+            let old_path_ids: Vec<String> = scored
+                .into_iter()
+                .take(10)
+                .map(|(id, _)| index.id_map[id as usize].clone())
+                .collect();
+
+            assert_eq!(pooled_ids, old_path_ids);
+            assert_eq!(caller_owned_ids, old_path_ids);
+        }
+    }
+
     #[test]
     fn test_diskann_with_pq() {
         let mut index = DiskAnnIndex::new(DiskAnnConfig {

crates/ruvector-diskann/src/lib.rs

@@ -19,6 +19,7 @@ pub mod pq;
 #[cfg(feature = "reuse-under-drift")]
 pub mod reuse;
 
+pub use distance::VisitedSet;
 pub use error::{DiskAnnError, Result};
 pub use index::{DiskAnnConfig, DiskAnnIndex};
 pub use pq::ProductQuantizer;