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

Fixes #677

Co-Authored-By: Leo <noreply@khive.ai>

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;

crates/ruvector-diskann/tests/search_allocations.rs

@@ -0,0 +1,91 @@
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use ruvector_diskann::{DiskAnnConfig, DiskAnnIndex};
+use std::alloc::{GlobalAlloc, Layout, System};
+use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
+
+struct CountingAllocator;
+
+static TRACKING: AtomicBool = AtomicBool::new(false);
+static LARGE_ALLOCATION_THRESHOLD: AtomicUsize = AtomicUsize::new(usize::MAX);
+static LARGE_ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
+static TOTAL_ALLOCATED_BYTES: AtomicUsize = AtomicUsize::new(0);
+
+unsafe impl GlobalAlloc for CountingAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        if TRACKING.load(Ordering::Relaxed) {
+            TOTAL_ALLOCATED_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
+            if layout.size() >= LARGE_ALLOCATION_THRESHOLD.load(Ordering::Relaxed) {
+                LARGE_ALLOCATED_BYTES.fetch_add(layout.size(), Ordering::Relaxed);
+            }
+        }
+        System.alloc(layout)
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        System.dealloc(ptr, layout);
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        if TRACKING.load(Ordering::Relaxed) {
+            TOTAL_ALLOCATED_BYTES.fetch_add(new_size, Ordering::Relaxed);
+            if new_size >= LARGE_ALLOCATION_THRESHOLD.load(Ordering::Relaxed) {
+                LARGE_ALLOCATED_BYTES.fetch_add(new_size, Ordering::Relaxed);
+            }
+        }
+        System.realloc(ptr, layout, new_size)
+    }
+}
+
+#[global_allocator]
+static ALLOCATOR: CountingAllocator = CountingAllocator;
+
+#[test]
+fn pooled_search_does_not_allocate_visited_set_storage() {
+    const N: usize = 250_000;
+    const DIM: usize = 8;
+    const SEARCHES: usize = 100;
+    const MAX_HOT_SEARCH_ALLOCATED_BYTES: usize = 300_000;
+
+    let mut rng = StdRng::seed_from_u64(0x677A_110C);
+    let mut index = DiskAnnIndex::new(DiskAnnConfig {
+        dim: DIM,
+        max_degree: 4,
+        build_beam: 8,
+        search_beam: 32,
+        alpha: 1.0,
+        ..Default::default()
+    });
+    for id in 0..N {
+        let vector = (0..DIM).map(|_| rng.gen()).collect();
+        index.insert(id.to_string(), vector).unwrap();
+    }
+    index.build().unwrap();
+
+    let query: Vec<f32> = (0..DIM).map(|_| rng.gen()).collect();
+    index.search(&query, 10).unwrap();
+
+    // VisitedSet's generation vector alone requests N * sizeof(u64) bytes.
+    // Count every allocation at least that large while allowing the unrelated,
+    // much smaller candidate and result buffers used by the search itself.
+    let visited_generation_bytes = N * std::mem::size_of::<u64>();
+    LARGE_ALLOCATION_THRESHOLD.store(visited_generation_bytes, Ordering::Relaxed);
+    LARGE_ALLOCATED_BYTES.store(0, Ordering::Relaxed);
+    TOTAL_ALLOCATED_BYTES.store(0, Ordering::Relaxed);
+    TRACKING.store(true, Ordering::Relaxed);
+    for _ in 0..SEARCHES {
+        std::hint::black_box(index.search(&query, 10).unwrap());
+    }
+    TRACKING.store(false, Ordering::Relaxed);
+
+    let allocated = LARGE_ALLOCATED_BYTES.load(Ordering::Relaxed);
+    let total_allocated = TOTAL_ALLOCATED_BYTES.load(Ordering::Relaxed);
+    assert_eq!(
+        allocated, 0,
+        "{SEARCHES} hot searches allocated {allocated} bytes in visited-set-sized blocks"
+    );
+    assert!(
+        total_allocated < MAX_HOT_SEARCH_ALLOCATED_BYTES,
+        "{SEARCHES} hot searches allocated {total_allocated} total bytes, expected less than {MAX_HOT_SEARCH_ALLOCATED_BYTES} bytes and far below the {visited_generation_bytes}-byte generation array"
+    );
+}