diff --git a/crates/uteke-core/src/consolidate.rs b/crates/uteke-core/src/consolidate.rs index 22f2bf3..e2a8a0f 100644 --- a/crates/uteke-core/src/consolidate.rs +++ b/crates/uteke-core/src/consolidate.rs @@ -248,6 +248,15 @@ impl crate::Uteke { let mut removed_ids = Vec::new(); let mut kept_ids = Vec::new(); let mut already_removed = std::collections::HashSet::new(); + let mut index_dirty = false; + + // Acquire the write lock ONCE before the loop to avoid repeated + // lock contention. Save once after all removals. + let mut index = self + .index + .write() + .map_err(|_| Error::lock("index write lock during consolidate"))?; + for pair in &pairs { if already_removed.contains(&pair.id_a) || already_removed.contains(&pair.id_b) { continue; @@ -286,25 +295,25 @@ impl crate::Uteke { .map_err(|e| Error::db("consolidate delete", e))?; } // SQLite first (source of truth), then vector index. - let mut index = self - .index - .write() - .map_err(|_| Error::lock("index write lock during consolidate"))?; if !index.remove(to_remove) { tracing::warn!( "Vector index entry not found during consolidate for id={}", to_remove ); } + index_dirty = true; + removed_ids.push(to_remove.clone()); + kept_ids.push(to_keep.clone()); + already_removed.insert(to_remove.clone()); + } + // Persist vector index once after all removals. + if index_dirty { if let Err(e) = index.save() { tracing::warn!( "Failed to persist vector index after consolidate: {e}. \ Orphan entries will be cleaned up by verify/repair." ); } - removed_ids.push(to_remove.clone()); - kept_ids.push(to_keep.clone()); - already_removed.insert(to_remove.clone()); } // Invalidate recall cache — deleted memories affect search results. if !removed_ids.is_empty() { diff --git a/crates/uteke-core/src/memory/vector.rs b/crates/uteke-core/src/memory/vector.rs index eaf53e5..3db5682 100644 --- a/crates/uteke-core/src/memory/vector.rs +++ b/crates/uteke-core/src/memory/vector.rs @@ -80,12 +80,12 @@ impl VectorIndex { /// identical — `save_to_buffer` and `restore_from_buffer` produce/consume /// the same byte stream as the native file-based methods. pub fn load_or_create(path: &Path, dims: usize) -> Result { - // Ensure the file exists so we can open + lock it. - if !path.exists() { - // Create a zero-byte placeholder; usearch will overwrite on save. - std::fs::write(path, []).map_err(|e| Error::embed("create usearch file", e))?; - } - + // Atomically create the file if it doesn't exist (avoids TOCTOU race + // where another process creates the file between our exists() and write()). + // O_CREAT | O_EXCL ensures only one writer wins; failure is harmless. + use std::fs::OpenOptions; + let _ = OpenOptions::new().create_new(true).write(true).open(path); + // Regardless of who created it, the file now exists — open + lock it. let mut lock_file = acquire_file_lock(path)?; let mut idx = if lock_file @@ -129,22 +129,26 @@ impl VectorIndex { let mut next_key = 0u64; let mapping_path = path.with_extension("keys"); - if mapping_path.exists() { - let data = std::fs::read_to_string(&mapping_path) - .map_err(|e| Error::embed("read key mapping", e))?; - for line in data.lines() { - let line = line.trim(); - if line.is_empty() { - continue; - } - if let Some((key_str, id)) = line.split_once('\t') { - if let Ok(key) = key_str.parse::() { - key_to_id.insert(key, id.to_string()); - id_to_key.insert(id.to_string(), key); - next_key = next_key.max(key + 1); + match std::fs::read_to_string(&mapping_path) { + Ok(data) => { + for line in data.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + if let Some((key_str, id)) = line.split_once('\t') { + if let Ok(key) = key_str.parse::() { + key_to_id.insert(key, id.to_string()); + id_to_key.insert(id.to_string(), key); + next_key = next_key.max(key.saturating_add(1)); + } } } } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + // No key mapping sidecar — fresh index, start from key 0. + } + Err(e) => return Err(Error::embed("read key mapping", e)), } Ok(Self { @@ -323,7 +327,7 @@ impl VectorIndex { } let key = self.next_key; - self.next_key += 1; + self.next_key = self.next_key.saturating_add(1); self.key_to_id.insert(key, id.to_string()); self.id_to_key.insert(id.to_string(), key); diff --git a/crates/uteke-server/src/main.rs b/crates/uteke-server/src/main.rs index c2b60cf..8d35084 100644 --- a/crates/uteke-server/src/main.rs +++ b/crates/uteke-server/src/main.rs @@ -466,12 +466,14 @@ fn main() { // Request loop — spawn each request in a thread for concurrent handling. // Arc> allows safe shared access across threads. - // Cap concurrent threads to prevent thread explosion under load: - // an atomic counter plus spin-yield provides simple backpressure. + // Cap concurrent threads via Condvar-based semaphore: park instead of spin. let max_threads = std::thread::available_parallelism() .map(|n| n.get() * 2) .unwrap_or(8); - let active = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let pair = Arc::new(( + std::sync::Mutex::new(0usize), // active count + std::sync::Condvar::new(), + )); for mut req in server.incoming_requests() { if SHUTDOWN.load(Ordering::SeqCst) { @@ -479,12 +481,13 @@ fn main() { break; } - // Backpressure: wait until a thread slot is available. - while active.load(Ordering::Acquire) >= max_threads { - if SHUTDOWN.load(Ordering::SeqCst) { - break; + // Backpressure: wait until a thread slot is available (parked, not spinning). + { + let (lock, cvar) = &*pair; + let mut active = lock.lock().unwrap(); + while *active >= max_threads && !SHUTDOWN.load(Ordering::SeqCst) { + active = cvar.wait(active).unwrap(); } - std::thread::yield_now(); } if SHUTDOWN.load(Ordering::SeqCst) { break; @@ -496,21 +499,32 @@ fn main() { let uteke = Arc::clone(&uteke); let ctx = ctx.clone(); - let active = Arc::clone(&active); - active.fetch_add(1, Ordering::AcqRel); + let pair = Arc::clone(&pair); + let pair_err = Arc::clone(&pair); + + { + let (lock, _) = &*pair; + *lock.lock().unwrap() += 1; + } - let active_clone = Arc::clone(&active); let result = std::thread::Builder::new().spawn(move || { let response = handlers::route(&uteke, &ctx, &mut req); if let Err(e) = req.respond(response) { warn!("Response error: {e}"); } - active.fetch_sub(1, Ordering::AcqRel); + // Release slot and notify the waiting accept loop. + let (lock, cvar) = &*pair; + let mut active = lock.lock().unwrap(); + *active -= 1; + cvar.notify_one(); }); if let Err(e) = result { // Spawn failed — release the slot we reserved. - active_clone.fetch_sub(1, Ordering::AcqRel); + let (lock, cvar) = &*pair_err; + let mut active = lock.lock().unwrap(); + *active -= 1; + cvar.notify_one(); warn!("Failed to spawn request thread: {e}"); } }