diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d3b22e2..8754e40c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,36 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Performance — Vector: COLD-segment reload moved off the shard event loop (PR #251) + +- Under `--disk-offload`, an idle vector segment is demoted to the COLD tier + (`UnloadedSegment` stub, everything in-memory dropped). The first FT.SEARCH to + touch it must reload it — `WarmSearchSegment::from_files`: mmap + page-in (+ + optional mlock) of a whole segment. That reload ran **inline on the + single-threaded shard event loop** (`promote_unloaded` at the top of the + yielding search capture), so it stalled EVERY other connection on the shard for + the full reload (hundreds of ms–seconds on real NVMe/EBS with a cold page + cache) — a multi-tenancy violation (audit finding 18). +- A new process-global `SegmentReloadPool` (`src/vector/reload_pool.rs`, auto-on; + `MOON_VEC_RELOAD_WORKERS=0` disables) now performs the reload I/O on dedicated + worker threads. The yielding capture path calls `submit_unloaded_reloads` + (non-blocking): it submits each COLD stub off-loop and stashes the `flume` + receivers in the `SearchSnapshot`; the handler `await`s them + (`await_pending_reloads`) before scanning, which parks only the triggering + query's **task** (not the OS thread) and then splices the reloaded segments + into that query's own snapshot. Result: the dense-KNN query that triggered the + reload still sees **full recall**, while sibling connections on the shard keep + running. Reloads are **single-flight** (concurrent first-touches of one segment + share a job) and cached until the next capture installs them into WARM + (reclaiming the stub's memory). +- Design-for-failure: a reload that errors or whose worker panics/dies replies + `Err`/disconnect to every waiter — the query answers with degraded recall + (segment stays COLD, retried next touch) rather than hanging or crashing. When + the pool is disabled, `submit_unloaded_reloads` falls back to the blocking + `promote_unloaded`, preserving exact pre-#18 behavior. Non-yielding sync search + paths are unchanged (still block); HYBRID/SPARSE/RANGE/non-default-field + queries do not take the off-loop path (accepted follow-up). + ### Docs — engineering roadmap suite under `docs/roadmap/` (PR #254) - New planning docs (excluded from the published docs site via `mkdocs.yml` diff --git a/src/command/vector_search/ft_search/dispatch.rs b/src/command/vector_search/ft_search/dispatch.rs index 5e96d2ff..235123b8 100644 --- a/src/command/vector_search/ft_search/dispatch.rs +++ b/src/command/vector_search/ft_search/dispatch.rs @@ -667,11 +667,15 @@ fn capture_dense_knn_snapshot( let filter_strategy = crate::vector::filter::selectivity::select_strategy(filter_bitmap.as_ref(), total_vectors); - // WS3 round 2: the yielding (worker-pool) search path captures its + // WS3 round 2 + #18: the yielding (worker-pool) search path captures its // segment snapshot here, separately from `SegmentHolder::search_filtered`'s - // own promote-before-search call -- must reload any COLD segments before - // this capture too, or a query on this path would silently skip them. - idx.segments.promote_unloaded(); + // own promote-before-search call. Rather than BLOCK the shard event loop on + // COLD-segment reloads, submit them off-loop and stash the receivers in the + // snapshot; the handler awaits them (parking only this query's task) before + // scanning, so siblings never stall. When the reload pool is disabled this + // falls back to the blocking `promote_unloaded` (no receivers), exactly as + // before, so a COLD segment is never silently skipped. + let pending_reloads = idx.segments.submit_unloaded_reloads(); let segments = idx.segments.load_full(); let mutable_len = segments.mutable.len(); @@ -696,6 +700,7 @@ fn capture_dense_knn_snapshot( rerank_mult: idx.meta.rerank_mult, exact_beam: idx.meta.exact_beam, }, + pending_reloads, }) } diff --git a/src/main.rs b/src/main.rs index 2019c7f6..d739554e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -488,6 +488,19 @@ fn main() -> anyhow::Result<()> { info!("FT.SEARCH worker pool: {ft_workers} threads"); } + // #18: off-loop COLD-segment reload pool. Moves the mmap+page-in of an + // idle-demoted vector segment off the shard event loop so the first + // FT.SEARCH that touches it doesn't stall sibling connections. Auto-on with + // a small worker count (reloads are rare + single-flight); disable with + // MOON_VEC_RELOAD_WORKERS=0 to fall back to the blocking reload. + let reload_workers = moon::vector::reload_pool::default_workers(num_shards); + moon::vector::reload_pool::init_global(reload_workers); + if reload_workers > 0 { + info!("vector COLD-segment reload pool: {reload_workers} threads"); + } else { + info!("vector COLD-segment reload pool: disabled (blocking reload)"); + } + // Checked cast: --shards is bounded by clap's value_parser, but `as u16` // would silently wrap for values > 65535. Fail loudly instead. // ALLOW: panic is appropriate here — this is `main`, not library code. diff --git a/src/server/conn/handler_monoio/ft.rs b/src/server/conn/handler_monoio/ft.rs index eda64704..a343c414 100644 --- a/src/server/conn/handler_monoio/ft.rs +++ b/src/server/conn/handler_monoio/ft.rs @@ -641,6 +641,10 @@ pub(super) async fn try_handle_ft_command( offset, count, } => { + // #18: await any off-loop COLD→WARM reloads submitted at + // capture (parks this task, not the shard thread) so the + // scan below sees full recall while siblings keep running. + snapshot.await_pending_reloads().await; let results = crate::vector::segment::holder::SegmentHolder::search_mvcc_yielding( &mut *snapshot, diff --git a/src/server/conn/handler_sharded/ft.rs b/src/server/conn/handler_sharded/ft.rs index b1127bc9..11d5d6f6 100644 --- a/src/server/conn/handler_sharded/ft.rs +++ b/src/server/conn/handler_sharded/ft.rs @@ -523,6 +523,10 @@ pub(super) async fn try_handle_ft_command( offset, count, } => { + // #18: await any off-loop COLD→WARM reloads submitted at + // capture (parks this task, not the shard thread) so the scan + // below sees full recall while siblings keep running. + snapshot.await_pending_reloads().await; let results = crate::vector::segment::holder::SegmentHolder::search_mvcc_yielding( &mut *snapshot, crate::vector::segment::holder::ft_search_yield_budget(), diff --git a/src/vector/mod.rs b/src/vector/mod.rs index 8abe0ccb..da90252a 100644 --- a/src/vector/mod.rs +++ b/src/vector/mod.rs @@ -13,6 +13,7 @@ pub mod keymap; pub mod metrics; pub mod mvcc; pub mod persistence; +pub mod reload_pool; pub mod search_pool; pub mod segment; pub mod sparse; diff --git a/src/vector/reload_pool.rs b/src/vector/reload_pool.rs new file mode 100644 index 00000000..2a112d67 --- /dev/null +++ b/src/vector/reload_pool.rs @@ -0,0 +1,334 @@ +//! Off-loop cold-segment reload pool (prod-hardening #18). +//! +//! When an idle vector segment has been demoted to the COLD tier +//! (`UnloadedSegment` stub) under `--disk-offload`, the first FT.SEARCH that +//! touches it must reload it (`WarmSearchSegment::from_files`: mmap + page-in + +//! optional mlock of a megabytes-to-gigabytes segment). Doing that inline on +//! the single-threaded shard event loop stalls EVERY other connection on the +//! shard for the whole reload (hundreds of ms–seconds on real NVMe/EBS with a +//! cold page cache) — a multi-tenancy violation. +//! +//! This pool moves the reload I/O onto dedicated worker threads. The shard +//! thread `submit`s an `Arc` (non-blocking) and gets back a +//! `flume::Receiver`; the awaiting **async** connection task +//! `recv_async().await`s it, which parks the *task*, never the OS thread, so +//! sibling connections keep running (the same cross-thread-reply idiom the +//! graph-tier `search_pool` and the cross-shard coordinator already rely on — +//! flume's receiver is polled from *inside* the task, sidestepping the monoio +//! `!Send`-waker gotcha). +//! +//! Guarantees: +//! - **Single-flight:** N concurrent first-touches of the same segment share +//! ONE reload job (deduped by `segment_id`); every waiter gets the result. +//! - **Completed cache:** a finished reload is held until the next capture on +//! that index installs it into the holder's WARM tier (see +//! `SegmentHolder::submit_unloaded_reloads`), so the segment is materialized +//! exactly once for future queries and the stub's memory is reclaimed. +//! - **Design-for-failure:** a panicking or erroring reload replies `Err` to +//! every waiter (never a hang); a disabled/uninitialized pool makes callers +//! fall back to the synchronous blocking reload — identical to pre-#18. + +use std::collections::HashMap; +use std::collections::hash_map::Entry; +use std::sync::Arc; + +use parking_lot::Mutex; + +use crate::vector::persistence::unloaded_segment::UnloadedSegment; +use crate::vector::persistence::warm_search::WarmSearchSegment; + +/// Result delivered to every waiter for one reload job. `Arc` keeps +/// the outcome `Clone` (io::Error is not) so a single reload can answer many +/// deduped waiters. +pub type ReloadOutcome = Result, Arc>; + +/// One in-flight reload: the stub to reload plus every waiter's reply sender. +struct InFlight { + stub: Arc, + waiters: Vec>, +} + +#[derive(Default)] +struct ReloadState { + /// Segments currently being reloaded, keyed by `segment_id` (single-flight). + in_flight: HashMap, + /// Finished reloads awaiting install into the holder's WARM tier. + completed: HashMap>, +} + +/// Worker-thread pool for cold-segment reloads. Create once per process +/// (see [`init_global`]); dropping it disconnects the channel and workers exit. +pub struct SegmentReloadPool { + job_tx: flume::Sender, + state: Arc>, + num_workers: usize, + _workers: Vec>, +} + +impl SegmentReloadPool { + /// Spawn `num_workers` (≥1) reload threads. + pub fn new(num_workers: usize) -> Self { + assert!(num_workers >= 1, "at least one worker required"); + // Unbounded: job volume is (distinct cold segments touched), naturally + // bounded by index count; the shard thread must NEVER block on submit. + let (job_tx, job_rx) = flume::unbounded::(); + let state = Arc::new(Mutex::new(ReloadState::default())); + let workers: Vec<_> = (0..num_workers) + .map(|i| { + let rx = job_rx.clone(); + let st = Arc::clone(&state); + std::thread::Builder::new() + .name(format!("moon-vec-reload-{i}")) + .spawn(move || worker_loop(&rx, &st)) + .expect("failed to spawn vector reload worker thread") + }) + .collect(); + Self { + job_tx, + state, + num_workers, + _workers: workers, + } + } + + /// Submit `stub` for reload and return a receiver for its outcome. + /// + /// Non-blocking. If the segment is already reloaded (completed cache) the + /// receiver resolves immediately; if a reload is already in flight this + /// waiter joins it (single-flight); otherwise a new job is enqueued. + pub fn submit(&self, stub: Arc) -> flume::Receiver { + let sid = stub.segment_id(); + let (wtx, wrx) = flume::bounded(1); + let mut enqueue = false; + { + let mut st = self.state.lock(); + if let Some(seg) = st.completed.get(&sid) { + // Already reloaded — answer immediately, keep it in `completed` + // so the next capture still installs it into the WARM tier. + let _ = wtx.send(Ok(Arc::clone(seg))); + return wrx; + } + match st.in_flight.entry(sid) { + Entry::Occupied(mut e) => e.get_mut().waiters.push(wtx), + Entry::Vacant(e) => { + e.insert(InFlight { + stub, + waiters: vec![wtx], + }); + enqueue = true; + } + } + } + if enqueue { + // Channel is unbounded; send only fails if the pool is shutting down, + // in which case the waiter's sender is dropped → RecvError, not a hang. + let _ = self.job_tx.send(sid); + } + wrx + } + + /// Take a finished reload for `segment_id`, if any, removing it from the + /// completed cache. Called by the holder to install the reloaded segment + /// into its WARM tier for future queries. + pub fn take_completed(&self, segment_id: u64) -> Option> { + self.state.lock().completed.remove(&segment_id) + } + + pub fn num_workers(&self) -> usize { + self.num_workers + } +} + +fn worker_loop(rx: &flume::Receiver, state: &Arc>) { + while let Ok(sid) = rx.recv() { + // Clone the stub out from under the lock so the (slow) reload I/O runs + // lock-free — concurrent submits for OTHER segments never block on it. + let stub = { + let st = state.lock(); + st.in_flight.get(&sid).map(|f| Arc::clone(&f.stub)) + }; + let Some(stub) = stub else { + continue; // Raced with a shutdown / already drained. + }; + + // Contain a panicking reload to this job (keep the worker alive). + let outcome: ReloadOutcome = + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| stub.reload())) { + Ok(Ok(seg)) => Ok(Arc::new(seg)), + Ok(Err(e)) => Err(Arc::new(e)), + Err(_) => Err(Arc::new(std::io::Error::other( + "cold segment reload panicked", + ))), + }; + + let waiters = { + let mut st = state.lock(); + if let Ok(ref seg) = outcome { + // Hold the reloaded segment for the next capture to install. + st.completed.insert(sid, Arc::clone(seg)); + } + st.in_flight + .remove(&sid) + .map(|f| f.waiters) + .unwrap_or_default() + }; + for w in waiters { + // Receiver may be gone (client disconnected mid-reload) — fine. + let _ = w.send(outcome.clone()); + } + } + // job_tx dropped → channel disconnected → exit cleanly. +} + +// ── Process-global pool ───────────────────────────────────────────────────── + +/// `None` inside = explicitly disabled. Uninitialized = disabled: callers then +/// fall back to the synchronous blocking reload (pre-#18 behavior), and unit +/// tests see that exact path unless they opt in. +static GLOBAL_POOL: std::sync::OnceLock> = std::sync::OnceLock::new(); + +/// Initialize the process-global reload pool. `num_workers == 0` disables it +/// (callers fall back to blocking reload). First call wins (idempotent). +pub fn init_global(num_workers: usize) { + let _ = GLOBAL_POOL.set(if num_workers == 0 { + None + } else { + Some(SegmentReloadPool::new(num_workers)) + }); +} + +/// The process-global pool, or `None` when disabled/uninitialized. +pub fn global() -> Option<&'static SegmentReloadPool> { + GLOBAL_POOL.get().and_then(|p| p.as_ref()) +} + +/// Resolve the default reload-worker count. +/// +/// `MOON_VEC_RELOAD_WORKERS=` overrides (0 = disable, falling back to the +/// blocking reload). Otherwise a small default (2 on multicore, 1 on ≤2-core +/// boxes) — reloads are rare and single-flight, so a couple of workers cover +/// concurrent cold touches without competing with shard threads. Never fewer +/// than 1 unless explicitly disabled. +pub fn default_workers(cores: usize) -> usize { + if let Ok(v) = std::env::var("MOON_VEC_RELOAD_WORKERS") { + if let Ok(n) = v.trim().parse::() { + return n; + } + } + if cores <= 2 { 1 } else { 2 } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::tiered::SegmentHandle; + use crate::vector::distance; + use crate::vector::hnsw::graph::HnswGraph; + use crate::vector::persistence::warm_segment::{ + write_codes_mpf, write_graph_mpf, write_mvcc_mpf, + }; + use crate::vector::turbo_quant::collection::{CollectionMetadata, QuantizationConfig}; + use crate::vector::types::DistanceMetric; + use std::path::Path; + + fn make_collection() -> Arc { + Arc::new(CollectionMetadata::new( + 1, + 128, + DistanceMetric::L2, + QuantizationConfig::TurboQuant4, + 42, + )) + } + + fn write_minimal_segment(seg_dir: &Path, seg_id: u64) { + std::fs::create_dir_all(seg_dir).unwrap(); + let empty_graph = HnswGraph::new( + 0, + 16, + 32, + 0, + 0, + crate::vector::aligned_buffer::AlignedBuffer::new(0), + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + 68, + ); + let graph_bytes = empty_graph.to_bytes(); + write_codes_mpf(&seg_dir.join("codes.mpf"), seg_id, &[]).unwrap(); + write_graph_mpf(&seg_dir.join("graph.mpf"), seg_id, &graph_bytes).unwrap(); + write_mvcc_mpf(&seg_dir.join("mvcc.mpf"), seg_id, &[]).unwrap(); + } + + /// Build a real `UnloadedSegment` stub backed by on-disk `.mpf` files. + fn make_stub(tmp: &Path, seg_id: u64) -> Arc { + let seg_dir = tmp.join(format!("seg-{seg_id}")); + write_minimal_segment(&seg_dir, seg_id); + let handle = SegmentHandle::new(seg_id, seg_dir.clone()); + let warm = + WarmSearchSegment::from_files(&seg_dir, seg_id, make_collection(), handle, false) + .unwrap(); + Arc::new(UnloadedSegment::from_warm(&warm, false)) + } + + #[test] + fn test_submit_reloads_off_thread_and_delivers() { + distance::init(); + let tmp = tempfile::tempdir().unwrap(); + let stub = make_stub(tmp.path(), 1); + let pool = SegmentReloadPool::new(1); + + let rx = pool.submit(Arc::clone(&stub)); + let outcome = rx.recv().expect("worker replies"); + let seg = outcome.expect("reload succeeds"); + assert_eq!(seg.segment_id(), 1); + + // Completed cache holds it for install until taken exactly once. + assert!(pool.take_completed(1).is_some()); + assert!(pool.take_completed(1).is_none()); + } + + #[test] + fn test_single_flight_dedupes_concurrent_submits() { + distance::init(); + let tmp = tempfile::tempdir().unwrap(); + let stub = make_stub(tmp.path(), 7); + let pool = SegmentReloadPool::new(2); + + // Three concurrent waiters for the SAME segment must all get the result + // from ONE reload job. + let r1 = pool.submit(Arc::clone(&stub)); + let r2 = pool.submit(Arc::clone(&stub)); + let r3 = pool.submit(Arc::clone(&stub)); + for r in [r1, r2, r3] { + let seg = r.recv().expect("reply").expect("reload ok"); + assert_eq!(seg.segment_id(), 7); + } + } + + #[test] + fn test_completed_cache_answers_later_submit_immediately() { + distance::init(); + let tmp = tempfile::tempdir().unwrap(); + let stub = make_stub(tmp.path(), 3); + let pool = SegmentReloadPool::new(1); + + // First reload populates the completed cache. + pool.submit(Arc::clone(&stub)).recv().unwrap().unwrap(); + // A later submit (before anyone installs) resolves from cache — the + // receiver is already ready. + let rx = pool.submit(Arc::clone(&stub)); + let seg = rx.try_recv().expect("cache hit is immediate").unwrap(); + assert_eq!(seg.segment_id(), 3); + } + + #[test] + fn test_default_workers_env_override() { + // Deterministic branch (no env set in the suite): ≤2 cores → 1, else 2. + assert_eq!(default_workers(1), 1); + assert_eq!(default_workers(2), 1); + assert_eq!(default_workers(8), 2); + } +} diff --git a/src/vector/search_pool.rs b/src/vector/search_pool.rs index 3863fe8a..53ca70fe 100644 --- a/src/vector/search_pool.rs +++ b/src/vector/search_pool.rs @@ -302,6 +302,7 @@ mod tests { key_hash_to_key: idx.key_hash_to_key.clone(), ef_defaulted: false, tuning: crate::vector::types::SearchTuning::default(), + pending_reloads: Vec::new(), } } @@ -387,6 +388,7 @@ mod tests { key_hash_to_key: key_map.clone(), ef_defaulted: false, tuning: crate::vector::types::SearchTuning::default(), + pending_reloads: Vec::new(), }; let results = futures::executor::block_on( SegmentHolder::search_mvcc_yielding_with_pool( diff --git a/src/vector/segment/holder.rs b/src/vector/segment/holder.rs index 95e8b348..edbf5374 100644 --- a/src/vector/segment/holder.rs +++ b/src/vector/segment/holder.rs @@ -217,6 +217,72 @@ pub struct SearchSnapshot { /// Per-index recall/QPS knobs (FT.CONFIG RERANK_MULT / EXACT_BEAM), /// copied from `IndexMeta` at capture. pub tuning: crate::vector::types::SearchTuning, + /// #18 (off-loop reload): receivers for COLD→WARM reloads that were + /// SUBMITTED (not blocked-on) at capture. Empty on the common path (no + /// unloaded segments) and whenever the reload pool is disabled (then the + /// capture blocked on `promote_unloaded` exactly as before). The yielding + /// handler `await`s these before scanning (`await_pending_reloads`), so the + /// triggering dense-KNN query sees FULL recall while sibling connections on + /// the shard keep running (the reload I/O is off the event loop). + pub pending_reloads: Vec>, +} + +impl SearchSnapshot { + /// #18: await every off-loop COLD→WARM reload submitted at capture, then + /// splice the reloaded segments into THIS snapshot's WARM tier (and drop the + /// matching stubs from `unloaded`) so the scan sees full recall. + /// + /// Called by the yielding FT.SEARCH handler BEFORE `search_mvcc_yielding`. + /// Awaiting a `flume` receiver parks the connection *task* (not the OS + /// thread), so sibling connections on the shard keep running while the + /// reload I/O completes on a pool worker. Design-for-failure: a reload that + /// errors or whose worker died is logged and skipped — the segment stays in + /// `unloaded` (retried at the next capture) and this query answers with + /// degraded recall rather than hanging or crashing. No-op (no await) when + /// nothing was submitted, i.e. the common path. + pub async fn await_pending_reloads(&mut self) { + if self.pending_reloads.is_empty() { + return; + } + let receivers = std::mem::take(&mut self.pending_reloads); + let mut reloaded: Vec> = + Vec::with_capacity(receivers.len()); + for rx in receivers { + match rx.recv_async().await { + Ok(Ok(seg)) => reloaded.push(seg), + Ok(Err(e)) => tracing::warn!( + error = %e, + "off-loop COLD segment reload failed -- query answers with degraded recall" + ), + Err(_) => tracing::warn!( + "off-loop COLD segment reload worker dropped -- query answers with degraded recall" + ), + } + } + if reloaded.is_empty() { + return; + } + + let cur = &self.segments; + let mut new_warm = cur.warm.clone(); + let reloaded_ids: std::collections::HashSet = + reloaded.iter().map(|s| s.segment_id()).collect(); + new_warm.extend(reloaded); + let new_unloaded: Vec<_> = cur + .unloaded + .iter() + .filter(|stub| !reloaded_ids.contains(&stub.segment_id())) + .cloned() + .collect(); + self.segments = Arc::new(SegmentList { + mutable: Arc::clone(&cur.mutable), + immutable: cur.immutable.clone(), + ivf: cur.ivf.clone(), + warm: new_warm, + cold: cur.cold.clone(), + unloaded: new_unloaded, + }); + } } /// Lock-free segment holder. Searches load() once at query start and hold @@ -1121,6 +1187,145 @@ mod tests { v } + /// #18: build a real COLD `UnloadedSegment` stub backed by minimal on-disk + /// `.mpf` files (0 vectors) — enough for `reload()` to round-trip a + /// `WarmSearchSegment` with the given `seg_id`. + fn make_unloaded_stub( + dir: &std::path::Path, + seg_id: u64, + ) -> Arc { + use crate::storage::tiered::SegmentHandle; + use crate::vector::hnsw::graph::HnswGraph; + use crate::vector::persistence::unloaded_segment::UnloadedSegment; + use crate::vector::persistence::warm_search::WarmSearchSegment; + use crate::vector::persistence::warm_segment::{ + write_codes_mpf, write_graph_mpf, write_mvcc_mpf, + }; + + let seg_dir = dir.join(format!("seg-{seg_id}")); + std::fs::create_dir_all(&seg_dir).unwrap(); + let empty_graph = HnswGraph::new( + 0, + 16, + 32, + 0, + 0, + crate::vector::aligned_buffer::AlignedBuffer::new(0), + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + 68, + ); + let graph_bytes = empty_graph.to_bytes(); + write_codes_mpf(&seg_dir.join("codes.mpf"), seg_id, &[]).unwrap(); + write_graph_mpf(&seg_dir.join("graph.mpf"), seg_id, &graph_bytes).unwrap(); + write_mvcc_mpf(&seg_dir.join("mvcc.mpf"), seg_id, &[]).unwrap(); + + let handle = SegmentHandle::new(seg_id, seg_dir.clone()); + let warm = WarmSearchSegment::from_files( + &seg_dir, + seg_id, + make_test_collection(128), + handle, + false, + ) + .unwrap(); + Arc::new(UnloadedSegment::from_warm(&warm, false)) + } + + /// #18 (holder half): with the reload pool enabled, `submit_unloaded_reloads` + /// must NOT block — it submits the stub off-loop (returns a receiver) and a + /// later capture installs the finished reload into WARM, emptying `unloaded`. + #[test] + fn test_submit_unloaded_reloads_installs_off_loop() { + distance::init(); + crate::vector::reload_pool::init_global(1); + // If a prior test disabled the global pool, this exercises the blocking + // fallback instead — still correct, asserted the same way below. + let pool_on = crate::vector::reload_pool::global().is_some(); + + let tmp = tempfile::tempdir().unwrap(); + let collection = make_test_collection(128); + let holder = SegmentHolder::new(128, collection.clone()); + holder.swap(SegmentList { + mutable: Arc::new(MutableSegment::new(128, collection)), + immutable: Vec::new(), + ivf: Vec::new(), + warm: Vec::new(), + cold: Vec::new(), + unloaded: vec![make_unloaded_stub(tmp.path(), 1)], + }); + + let receivers = holder.submit_unloaded_reloads(); + if pool_on { + assert_eq!(receivers.len(), 1, "one off-loop reload submitted"); + let outcome = receivers[0].recv().expect("worker replies"); + assert!(outcome.is_ok(), "reload succeeds"); + // Next capture installs the completed reload into WARM. + let again = holder.submit_unloaded_reloads(); + assert!(again.is_empty(), "nothing left to submit after install"); + } + // Fallback path already promoted synchronously; either way, converged: + let snap = holder.load(); + assert_eq!(snap.warm.len(), 1, "segment now resident in WARM"); + assert!(snap.unloaded.is_empty(), "no COLD stubs remain"); + } + + /// #18 (snapshot half): `await_pending_reloads` splices a reloaded segment + /// into the snapshot's OWN warm tier and drops the matching cold stub, so + /// the scan that follows sees full recall. + #[test] + fn test_await_pending_reloads_splices_into_snapshot() { + distance::init(); + let tmp = tempfile::tempdir().unwrap(); + let collection = make_test_collection(128); + let stub = make_unloaded_stub(tmp.path(), 7); + + // Local pool (no global-state coupling) to produce a real receiver. + let pool = crate::vector::reload_pool::SegmentReloadPool::new(1); + let rx = pool.submit(Arc::clone(&stub)); + + let segments = Arc::new(SegmentList { + mutable: Arc::new(MutableSegment::new(128, collection)), + immutable: Vec::new(), + ivf: Vec::new(), + warm: Vec::new(), + cold: Vec::new(), + unloaded: vec![stub], + }); + let mut snap = SearchSnapshot { + segments, + query_f32: vec![0.0f32; 128], + k: 1, + ef_search: 16, + filter_bitmap: None, + filter_strategy: crate::vector::filter::selectivity::FilterStrategy::Unfiltered, + snapshot_lsn: 0, + my_txn_id: 0, + committed: std::sync::Arc::new(roaring::RoaringTreemap::new()), + dimension: 128, + mutable_len: 0, + scratch: crate::vector::hnsw::search::SearchScratch::new(0, padded_dimension(128)), + key_hash_to_key: crate::vector::keymap::BucketedKeyMap::new(), + ef_defaulted: false, + tuning: crate::vector::types::SearchTuning::default(), + pending_reloads: vec![rx], + }; + + assert_eq!(snap.segments.unloaded.len(), 1); + assert_eq!(snap.segments.warm.len(), 0); + + futures::executor::block_on(snap.await_pending_reloads()); + + assert_eq!(snap.segments.warm.len(), 1, "reloaded segment spliced in"); + assert!( + snap.segments.unloaded.is_empty(), + "cold stub dropped after splice" + ); + assert!(snap.pending_reloads.is_empty(), "receivers drained"); + } + #[test] fn test_holder_new_has_empty_immutable() { let collection = make_test_collection(128); diff --git a/src/vector/segment/holder/promote.rs b/src/vector/segment/holder/promote.rs index 67eca324..60baea1a 100644 --- a/src/vector/segment/holder/promote.rs +++ b/src/vector/segment/holder/promote.rs @@ -88,4 +88,80 @@ impl SegmentHolder { } promoted } + + /// #18 (off-loop reload): the non-blocking counterpart to + /// [`Self::promote_unloaded`], for the yielding FT.SEARCH capture path. + /// + /// Instead of blocking the shard event loop on `UnloadedSegment::reload` + /// (mmap + page-in of a whole segment) it: + /// 1. installs any reload that ALREADY completed (submitted by an earlier + /// query) into `warm` — a cheap in-memory swap, the memory-reclaim step; + /// 2. SUBMITS every still-cold stub to the process reload pool and returns + /// the receivers so the caller can stash them in the `SearchSnapshot` + /// and `await` them off-loop (see `SearchSnapshot::await_pending_reloads`). + /// + /// Falls back to the blocking `promote_unloaded` (returning no receivers) + /// when the pool is disabled/uninitialized — identical to pre-#18 behavior, + /// which is also exactly what non-yielding sync search paths still do. + /// + /// Fast path: a single `is_empty()` check with no lock and no allocation + /// when there are no COLD segments (the overwhelmingly common case). + pub fn submit_unloaded_reloads( + &self, + ) -> Vec> { + { + let snap = self.segments.load(); + if snap.unloaded.is_empty() { + return Vec::new(); + } + } + let Some(pool) = crate::vector::reload_pool::global() else { + // Pool disabled: preserve the blocking pre-#18 contract exactly. + self.promote_unloaded(); + return Vec::new(); + }; + + let _guard = self.reload_lock.lock(); + let snap = self.segments.load(); + if snap.unloaded.is_empty() { + return Vec::new(); + } + + let mut new_warm = snap.warm.clone(); + let mut still_unloaded: Vec< + Arc, + > = Vec::new(); + let mut receivers = Vec::new(); + let mut installed = 0usize; + for stub in snap.unloaded.iter() { + if let Some(seg) = pool.take_completed(stub.segment_id()) { + // A prior query already reloaded this off-loop — install it now + // so future queries hit WARM and the stub's memory is freed. + tracing::info!( + segment_id = stub.segment_id(), + "COLD segment installed from completed off-loop reload" + ); + new_warm.push(seg); + installed += 1; + } else { + // Not ready yet: submit (single-flight in the pool) and keep the + // stub cold; this query awaits the receiver for full recall. + receivers.push(pool.submit(Arc::clone(stub))); + still_unloaded.push(Arc::clone(stub)); + } + } + + if installed > 0 { + let new_list = SegmentList { + mutable: Arc::clone(&snap.mutable), + immutable: snap.immutable.clone(), + ivf: snap.ivf.clone(), + warm: new_warm, + cold: snap.cold.clone(), + unloaded: still_unloaded, + }; + self.segments.store(Arc::new(new_list)); + } + receivers + } }