fix(vector): memory-aware WARM offload with a real, reloadable ceiling

CHANGELOG.md

@@ -6,6 +6,34 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed — Vector: memory-aware WARM offload with a real, reloadable ceiling (PR #252)
+
+- **Reloadable byte-cap eviction (A).** `MmapBudget::enforce_budget` (the
+  per-shard `--vec-warm-mmap-budget` cap, default 2gb) evicted a WARM segment by
+  dropping its `Arc` outright with no COLD stub — so a byte-cap eviction silently
+  removed the segment from search until process restart (recall loss), despite
+  the doc claiming reload-on-touch. It now demotes each evicted segment to a
+  reloadable `UnloadedSegment` stub in `unloaded` (every WARM segment is durably
+  disk-backed via `transition_to_warm`, so the stub reloads byte-identically).
+  This makes `--vec-warm-mmap-budget` a real, reloadable memory ceiling. The
+  eviction tick now also takes the holder `reload_lock` so it serializes with
+  the reload/install path instead of relying solely on shard-thread affinity.
+- **WARM memory accounting fix.** `SegmentHolder::resident_bytes()` hardcoded the
+  WARM/COLD tiers to 0, so a shard whose HOT segments had aged into WARM reported
+  ~0 vector memory — blinding both INFO/Prometheus and the memory-pressure
+  trigger below. It now sums the WARM tier (the dominant term for a long-lived
+  shard) plus COLD stubs. (IVF/DiskANN-cold still lack a resident accessor.)
+- **Memory-triggered early offload (C).** Vector segment memory was invisible to
+  every pressure mechanism, so a vector-heavy shard (the primary disk-offload
+  case) only ever offloaded on the wall-clock idle timer
+  (`--engine-offload-idle-secs`, default 3600s), never on RAM pressure; and the
+  pressure cascade's step 2 demoted HOT→WARM (no RSS win). Now the shard's vector
+  resident bytes (HOT immutable + WARM) count toward the pressure trigger, and
+  under pressure step 2 offloads idle vector segments straight to COLD (real RAM
+  reclaim, reloadable) at an aggressive 60s idle floor. Known follow-up:
+  HOT-immutable memory still has no standalone byte budget (only idle/pressure
+  demotion bounds it).
+
 ### 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

src/shard/persistence_tick.rs

@@ -360,7 +360,7 @@ pub(crate) fn run_eviction_tick(
     // Uses the existing lock path (Wave E collapses to slice). Runs every tick
     // so Prometheus and MEMORY DOCTOR never see stale zero values for long.
     // C5 / M4: publish vector/text/graph store-memory atomics via thread-local slice.
-    crate::shard::slice::with_shard(|s| {
+    let vector_resident_bytes = crate::shard::slice::with_shard(|s| {
         use std::sync::atomic::Ordering;
         let (mutable, immutable) = s.vector_store.resident_bytes();
         s.store_memory
@@ -380,10 +380,19 @@ pub(crate) fn run_eviction_tick(
         s.store_memory
             .lua
             .store(script_cache.borrow().resident_bytes(), Ordering::Relaxed);
+        // Return the vector resident total (HOT + WARM) so the pressure check
+        // below can factor it in (memory-triggered vector offload, C).
+        mutable + immutable
     });
 
     if server_config.disk_offload_enabled()
-        && should_run_pressure_cascade(runtime_config, server_config, shard_databases, shard_id)
+        && should_run_pressure_cascade(
+            runtime_config,
+            server_config,
+            shard_databases,
+            shard_id,
+            vector_resident_bytes,
+        )
     {
         handle_memory_pressure(
             page_cache,
@@ -499,6 +508,14 @@ fn apply_completion_vec(
 // Memory pressure cascade (design section 8.5)
 // ---------------------------------------------------------------------------
 
+/// Aggressive idle floor (seconds) used by the memory-pressure cascade to
+/// offload idle vector segments to COLD early (C). Far below the normal
+/// `--engine-offload-idle-secs` (default 3600s): once a shard is over its
+/// memory budget, a segment untouched for a minute is worth shedding to reclaim
+/// RAM rather than waiting out the full idle timeout. Actively-queried segments
+/// (idle < this) stay resident; anything shed reloads on next touch.
+const PRESSURE_OFFLOAD_IDLE_SECS: u64 = 60;
+
 /// Check if memory usage exceeds the disk offload threshold.
 ///
 /// Returns `true` when the pressure cascade should run. Uses actual
@@ -508,6 +525,7 @@ pub(crate) fn should_run_pressure_cascade(
     server_config: &std::sync::Arc<crate::config::ServerConfig>,
     shard_databases: &std::sync::Arc<super::shared_databases::ShardDatabases>,
     shard_id: usize,
+    vector_resident_bytes: usize,
 ) -> bool {
     let rt = runtime_config.read();
     if rt.maxmemory == 0 {
@@ -525,7 +543,16 @@ pub(crate) fn should_run_pressure_cascade(
     let threshold = (budget as f64 * server_config.disk_offload_threshold) as usize;
     // C5 / Phase 3: read the already-published per-shard KV memory (written
     // earlier this same tick by `run_eviction_tick`). Lock-free Relaxed load.
-    let used = shard_databases.published_shard_memory(shard_id);
+    //
+    // Memory-triggered vector offload (C): add the shard's vector resident
+    // bytes (HOT immutable + WARM, computed this same tick) so a vector-heavy
+    // workload — the primary disk-offload use case, where KV is light but
+    // vector segments are the RAM hog — actually triggers the cascade. Without
+    // this, vector memory was invisible to every pressure mechanism and idle
+    // segments only ever offloaded on the wall-clock idle timer.
+    let used = shard_databases
+        .published_shard_memory(shard_id)
+        .saturating_add(vector_resident_bytes);
     used > threshold
 }
 
@@ -564,29 +591,37 @@ pub(crate) fn handle_memory_pressure(
         }
     }
 
-    // Step 2: Force-demote oldest HOT ImmutableSegments to WARM.
-    // Use half the normal warm_after threshold to be more aggressive under pressure.
+    // Step 2: Force-offload idle vector segments straight to COLD (memory-
+    // triggered early offload, C). Previously this demoted HOT->WARM, which
+    // frees NO resident bytes (a WarmSearchSegment is a same-size heap copy of
+    // the HOT segment). Under genuine memory pressure we instead shed the
+    // segments that have been idle beyond an aggressive floor
+    // (`PRESSURE_OFFLOAD_IDLE_SECS`, far below the normal
+    // `--engine-offload-idle-secs`) all the way to COLD (`UnloadedSegment`
+    // stub), which actually returns RAM — and stays reloadable-on-touch. We
+    // pass `warm_after = u64::MAX` so the age-based HOT->WARM path stays
+    // disabled here; only the idle->COLD path fires.
     if let Some(ref mut manifest) = *shard_manifest {
-        let aggressive_threshold = server_config.segment_warm_after / 2;
         let shard_dir = server_config
             .effective_disk_offload_dir()
             .join(format!("shard-{}", shard_id));
         let count = crate::shard::slice::with_shard(|s| {
-            s.vector_store.try_warm_transitions_all(
+            s.vector_store.try_warm_transitions_all_idle(
                 &shard_dir,
                 manifest,
-                aggressive_threshold,
+                u64::MAX,
+                PRESSURE_OFFLOAD_IDLE_SECS,
                 next_file_id,
                 wal_v3,
             )
         });
         if count > 0 {
             tracing::info!(
-                "Shard {}: memory pressure step 2 -- force-demoted {} segment(s) HOT->WARM",
+                "Shard {}: memory pressure step 2 -- offloaded {} idle vector segment(s) to COLD",
                 shard_id,
                 count
             );
-            return; // Freed memory via warm transition; re-evaluate next tick
+            return; // Freed RAM via cold offload; re-evaluate next tick
         }
     }
 
@@ -1515,7 +1550,7 @@ mod tests {
         // Below the 85% threshold (e.g. 50%): must NOT trigger the cascade.
         shared.publish_memory(0, (1024 * 1024) / 2);
         assert!(
-            !should_run_pressure_cascade(&runtime_config, &server_config, &shared, 0),
+            !should_run_pressure_cascade(&runtime_config, &server_config, &shared, 0, 0),
             "50% used_memory must stay below the 85% disk-offload-threshold"
         );
 
@@ -1524,7 +1559,7 @@ mod tests {
         // whole point of WS3 priority 3.
         shared.publish_memory(0, (1024 * 1024 * 90) / 100);
         assert!(
-            should_run_pressure_cascade(&runtime_config, &server_config, &shared, 0),
+            should_run_pressure_cascade(&runtime_config, &server_config, &shared, 0, 0),
             "90% used_memory must cross the 85% disk-offload-threshold and \
              trigger the pressure cascade well before maxmemory is reached"
         );
@@ -1537,9 +1572,42 @@ mod tests {
             let runtime_config2 = Arc::new(parking_lot::RwLock::new(rt2));
             shared.publish_memory(0, usize::MAX / 2);
             assert!(
-                !should_run_pressure_cascade(&runtime_config2, &server_config, &shared, 0),
+                !should_run_pressure_cascade(&runtime_config2, &server_config, &shared, 0, 0),
                 "no memory limit configured => no pressure possible"
             );
         }
     }
+
+    /// C: vector resident bytes must count toward the pressure trigger — a
+    /// vector-heavy shard with near-zero KV memory still fires the cascade
+    /// (previously vector memory was invisible to every pressure mechanism).
+    #[test]
+    fn test_pressure_cascade_triggered_by_vector_memory_alone() {
+        use clap::Parser;
+        let dbs = vec![vec![Database::new()]];
+        let (shared, _inits) = ShardDatabases::new(dbs);
+
+        let mut rt = crate::config::RuntimeConfig::default();
+        rt.maxmemory = 1024 * 1024; // 1 MiB per-shard budget (1 shard)
+        rt.num_shards = 1;
+        let runtime_config = Arc::new(parking_lot::RwLock::new(rt));
+        let server_config = Arc::new(crate::config::ServerConfig::parse_from::<[&str; 0], &str>(
+            [],
+        ));
+
+        // KV memory is trivial (well under the 85% threshold on its own)...
+        shared.publish_memory(0, 1024);
+        assert!(
+            !should_run_pressure_cascade(&runtime_config, &server_config, &shared, 0, 0),
+            "KV alone is far below threshold => no cascade without vector accounting"
+        );
+
+        // ...but the shard is holding ~950 KiB of resident vector segments,
+        // pushing total past the 85% (~892 KiB) threshold.
+        let vec_bytes = (1024 * 1024 * 93) / 100;
+        assert!(
+            should_run_pressure_cascade(&runtime_config, &server_config, &shared, 0, vec_bytes),
+            "vector resident memory must contribute to the pressure trigger"
+        );
+    }
 }

src/vector/persistence/mmap_budget.rs

@@ -13,10 +13,12 @@
 //! `MmapBudget` is a per-shard LRU tracker over `WarmSearchSegment` instances.
 //! On every search the caller records an access (bumps LRU position). The
 //! budget enforcer is called periodically (from the warm-check timer) and
-//! removes the least-recently-accessed `WarmSearchSegment` Arcs from the
-//! `SegmentList.warm` list until the tracked resident bytes fall below the
-//! configured budget. The on-disk .mpf files remain intact; the next search
-//! against an evicted segment causes `from_files` to reload it transparently.
+//! demotes the least-recently-accessed `WarmSearchSegment`s from the
+//! `SegmentList.warm` list to reloadable COLD `UnloadedSegment` stubs (pushed
+//! into `SegmentList.unloaded`) until the tracked resident bytes fall below the
+//! configured budget. The heap copy is freed; the on-disk .mpf files remain
+//! intact, and the next search against an evicted segment reloads it via
+//! `promote_unloaded`/`submit_unloaded_reloads` (the COLD→WARM reload path).
 //!
 //! This mirrors `MADV_DONTNEED` semantics at the segment level: we release the
 //! resident pages (owned Vec<u8>) to the OS and lazily fault them back on
@@ -47,7 +49,8 @@ use crate::vector::segment::SegmentList;
 /// Statistics returned by `enforce_budget`.
 #[derive(Debug, Default, Clone, Copy)]
 pub struct EnforceStats {
-    /// Number of warm segments evicted (Arcs dropped from `SegmentList.warm`).
+    /// Number of warm segments evicted (demoted from `SegmentList.warm` to
+    /// reloadable COLD stubs in `SegmentList.unloaded`).
     pub segments_evicted: u64,
     /// Resident bytes freed by eviction.
     pub bytes_freed: u64,
@@ -234,12 +237,32 @@ impl MmapBudget {
             };
         }
 
-        // Remove evicted segments from segment_list.warm.
-        // Arc drop triggers SegmentHandle refcount decrement; directory removal
-        // only occurs if the handle is tombstoned AND refcount hits zero.
-        segment_list
-            .warm
-            .retain(|arc| !evict_ids.contains(&arc.segment_id()));
+        // Demote evicted WARM segments to reloadable COLD stubs instead of
+        // dropping the Arc outright. Every WARM segment is durably disk-backed
+        // (`transition_to_warm` wrote its .mpf files before `from_files` ever
+        // loaded it), so `UnloadedSegment::from_warm` captures a stub that
+        // reloads byte-identically on next touch (segment_id + a cloned
+        // SegmentHandle keeping the directory alive + any tombstones). The heap
+        // copy is freed when the old `Arc<WarmSearchSegment>` drops. Previously
+        // this dropped the Arc with no stub, so a budget-evicted segment
+        // silently vanished from search until process restart (recall loss).
+        let mut stubs: Vec<
+            std::sync::Arc<crate::vector::persistence::unloaded_segment::UnloadedSegment>,
+        > = Vec::with_capacity(evict_ids.len());
+        segment_list.warm.retain(|arc| {
+            if evict_ids.contains(&arc.segment_id()) {
+                stubs.push(std::sync::Arc::new(
+                    crate::vector::persistence::unloaded_segment::UnloadedSegment::from_warm(
+                        arc.as_ref(),
+                        false,
+                    ),
+                ));
+                false
+            } else {
+                true
+            }
+        });
+        segment_list.unloaded.extend(stubs);
 
         // Update the tracker.
         for seg_id in &evict_ids {
@@ -439,6 +462,51 @@ mod tests {
         );
     }
 
+    /// A: budget eviction must DEMOTE the segment to a reloadable COLD stub in
+    /// `unloaded` — not drop it outright. This is the fix for the documented
+    /// silent-recall-loss bug (evicted segment vanished until restart).
+    #[test]
+    fn test_eviction_demotes_to_reloadable_stub() {
+        distance::init();
+        let tmp = tempfile::tempdir().unwrap();
+        let w1 = make_warm_segment(tmp.path(), 1);
+        let w2 = make_warm_segment(tmp.path(), 2);
+        std::thread::sleep(std::time::Duration::from_millis(2));
+        let w3 = make_warm_segment(tmp.path(), 3);
+        // Make seg 3 the MRU so it is protected; 1 and 2 are the LRU victims.
+        {
+            let query = vec![0.0f32; 128];
+            let mut scratch = crate::vector::hnsw::search::SearchScratch::new(0, 128);
+            let _ = w3.search(&query, 1, 10, &mut scratch);
+        }
+
+        let mut budget = MmapBudget::new(1);
+        budget.register_segment(1, 1_000_000);
+        budget.register_segment(2, 1_000_000);
+        budget.register_segment(3, 1_000_000);
+
+        let mut list = empty_segment_list(vec![w1, w2, w3]);
+        let stats = budget.enforce_budget(&mut list);
+
+        assert!(stats.segments_evicted >= 1, "expected eviction; {stats:?}");
+        // Evicted segments are NOT lost — every one lands in `unloaded` as a stub.
+        assert_eq!(
+            list.unloaded.len() as u64,
+            stats.segments_evicted,
+            "every evicted WARM segment must become a COLD stub, not vanish"
+        );
+        // Warm shrank by exactly the eviction count; MRU (seg 3) survives.
+        assert_eq!(list.warm.len(), 3 - stats.segments_evicted as usize);
+        assert!(
+            list.warm.iter().any(|w| w.segment_id() == 3),
+            "MRU segment 3 must stay resident"
+        );
+        // The stub is genuinely reloadable from its on-disk .mpf files.
+        let stub = &list.unloaded[0];
+        let reloaded = stub.reload().expect("evicted stub reloads from disk");
+        assert_eq!(reloaded.segment_id(), stub.segment_id());
+    }
+
     /// The segment most recently touched by a search is protected from eviction.
     ///
     /// LRU ordering now comes from `WarmSearchSegment::last_access_micros()` which