diff --git a/CHANGELOG.md b/CHANGELOG.md index 3fc4f29d1..caf851479 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,52 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Changed — Memory: vector tiering accounting spine (M1, tiering-v2 D3/D9) + +- **`--maxmemory` now counts vector segment memory.** The background eviction + check compared only KV bytes against the per-shard budget, so a pure-vector + workload could drive RSS to OOM while eviction reported "under budget". + `timers::run_eviction` now gates on the shard AGGREGATE — Σ all dbs' KV + + the shard's published vector bytes, computed once per 100ms tick — and + evicts across dbs only until the aggregate is back under budget (adversarial + review caught the initial per-db formulation, which both under-detected with + KV spread across dbs and over-evicted sibling dbs). When the un-evictable + vector term alone exceeds the budget, KV drains then errors OOM + (shared-budget semantics; per-db quotas remain the tenant-isolation + mechanism); the pressure cascade (which shrinks vectors via offload) fires + earlier at `--disk-offload-threshold`, so with disk-offload enabled vectors + shed first. **Known limitation:** the on-write eviction gate still checks + KV-only, so under `noeviction` a vector-heavy shard over the vector-aware + budget does not yet reject client writes — RSS is bounded by the pressure + cascade and the RSS watchdog instead (write-gate consistency is a tracked + M1 follow-up). +- **Elastic budget classification is vector-aware.** A vector-heavy/KV-light + shard was misclassified as an idle donor, lending headroom to siblings while + its true footprint was over base — and the pressure cascade compared a + vector-inclusive used-term against a budget inflated by that donation. The + donor/hot snapshot now sums KV + vector per shard; a vector-heavy shard is + classified hot and borrows instead of donating. +- **`--vec-warm-mmap-budget` is now an instance-total cap divided across + shards** (matching `--maxmemory` semantics). Each shard previously applied + the full value — an N-shard instance silently allowed N× the configured WARM + memory. **Behavior change:** multi-shard deployments relying on the old + per-shard meaning should multiply their flag value by the shard count. A + nonzero total floors at 1 byte/shard ("0" still disables enforcement). +- **IVF and DiskANN-cold segments report real `resident_bytes()`.** Both tiers + contributed a hardcoded 0 to the roll-up (untracked RAM: IVF centroids + + posting lists; DiskANN PQ codes + codebook). They now feed the pressure + trigger, MEMORY DOCTOR, and Prometheus like every other tier. +- **`INFO reclamation_mmap_warm_bytes` reports the live WARM counter.** The + field read a never-incremented local static (permanent 0) while the real + `MmapBudget` counter was write-only. It now reads the live counter, and a + new `reclamation_mmap_budget_evictions_total` field exposes cumulative + byte-cap evictions. +- **D9 (DiskANN retention) quarantine:** config docs now disambiguate + COLD-stub (`unloaded`, exact reload-on-touch default valve) vs COLD-ann + (`cold`, DiskANN serve-from-disk, inert behind `MOON_VEC_COLD_TIER`); the + dead knobs are marked `[reserved: M3/M5]`. Keep-vs-delete is decided at the + M3 exit-review on real per-index query-frequency telemetry. + ### Fixed — Vector: memory-aware WARM offload with a real, reloadable ceiling (PR #252) - **Reloadable byte-cap eviction (A).** `MmapBudget::enforce_budget` (the diff --git a/src/command/info_reclamation.rs b/src/command/info_reclamation.rs index e0d78111c..61a3b59db 100644 --- a/src/command/info_reclamation.rs +++ b/src/command/info_reclamation.rs @@ -104,10 +104,6 @@ pub static RECL_MANIFEST_ACTIVE: AtomicU64 = AtomicU64::new(0); /// P1 wires this alongside RECL_MANIFEST_ACTIVE. pub static RECL_MANIFEST_TOMBSTONES: AtomicU64 = AtomicU64::new(0); -/// Bytes mapped by mmap'd warm vector segment files (OS page-cache backed). -/// TODO(P10→Wave2): wire from WarmSearchSegment on map/unmap. -pub static RECL_MMAP_WARM_BYTES: AtomicU64 = AtomicU64::new(0); - /// Number of committed transactions in the MVCC treemap (RoaringTreemap::len()). /// MA2 wires this; emits 0 until pruning is implemented. pub static RECL_MVCC_COMMITTED: AtomicU64 = AtomicU64::new(0); @@ -273,10 +269,15 @@ pub fn write_reclamation_section(buf: &mut String) { ); // -- Mmap warm bytes -- + // Re-pointed at the LIVE warm-tier resident counter maintained by + // `MmapBudget` (add/sub on register/evict). The former local + // `RECL_MMAP_WARM_BYTES` static here was never incremented (permanent 0). let _ = write!( buf, - "reclamation_mmap_warm_bytes:{}\r\n", - RECL_MMAP_WARM_BYTES.load(Ordering::Relaxed) + "reclamation_mmap_warm_bytes:{}\r\n\ + reclamation_mmap_budget_evictions_total:{}\r\n", + crate::admin::recl_atomics::warm_resident_bytes(), + crate::admin::recl_atomics::budget_evictions_total() ); // -- MVCC -- @@ -494,6 +495,41 @@ mod tests { "RECL_WAL_BYTES store must be visible in section output" ); } + + /// The `reclamation_mmap_warm_bytes` INFO field must reflect the LIVE + /// warm-tier resident-bytes counter maintained by `MmapBudget` + /// (`admin::recl_atomics`), not a dead never-incremented local static. + /// + /// RED until the emit is re-pointed at `recl_atomics::warm_resident_bytes()`: + /// the pre-existing `info_reclamation::RECL_MMAP_WARM_BYTES` is never + /// incremented anywhere, so the field is a permanent `0`. + /// + /// Robust to parallel `MmapBudget` tests touching the same process-global + /// counter: we hold a known `DELTA` across the emit so the live counter has + /// a provable floor (`>= DELTA`); balanced add/sub in sibling tests cannot + /// erase our contribution before our own `sub`. + #[test] + fn info_reclamation_mmap_warm_bytes_reflects_live_counter() { + use crate::admin::recl_atomics; + const DELTA: u64 = 4_096; + recl_atomics::add_warm_resident(DELTA); + let mut buf = String::new(); + write_reclamation_section(&mut buf); + recl_atomics::sub_warm_resident(DELTA); // restore + + let emitted = buf + .lines() + .find_map(|l| l.strip_prefix("reclamation_mmap_warm_bytes:")) + .and_then(|v| v.trim().parse::().ok()) + .expect("reclamation_mmap_warm_bytes field must be present and numeric"); + + assert!( + emitted >= DELTA, + "INFO reclamation_mmap_warm_bytes ({emitted}) must reflect the live \ + warm-resident counter (>= {DELTA} while we hold a delta); the dead \ + never-incremented static emits a permanent 0" + ); + } } // --------------------------------------------------------------------------- diff --git a/src/config.rs b/src/config.rs index 993eb5562..09557c3ba 100644 --- a/src/config.rs +++ b/src/config.rs @@ -429,20 +429,31 @@ pub struct ServerConfig { #[arg(long = "vec-codes-mlock", default_value = "enable")] pub vec_codes_mlock: String, - /// Maximum resident bytes allowed across all warm-tier vector segments on - /// this shard (e.g. "2gb", "512mb", "0"). When the total exceeds this - /// limit the budget enforcer drops LRU warm segments from memory; they - /// are reloaded from disk on next access. Set to "0" to disable. + /// Maximum resident bytes allowed across all warm-tier vector segments + /// on this INSTANCE (e.g. "2gb", "512mb", "0"), divided evenly across + /// shards — matching `--maxmemory` semantics (A5, tiering-v2 D3). When a + /// shard's share is exceeded the budget enforcer demotes LRU warm + /// segments to reloadable COLD stubs; they reload from disk on next + /// access. Set to "0" to disable. /// /// Default: "2gb". Tune down for cgroup-constrained containers. #[arg(long = "vec-warm-mmap-budget", default_value = "2gb")] pub vec_warm_mmap_budget: String, // ── Cold-tier / DiskANN (EXPERIMENTAL — gated by MOON_VEC_COLD_TIER) ── + // Two distinct COLD concepts exist (tiering-v2 decision D9): + // COLD-stub (`SegmentList.unloaded`, `UnloadedSegment`) — the DEFAULT + // valve: exact, ~0 RAM, reload-on-touch. Always available. + // COLD-ann (`SegmentList.cold`, `DiskAnnSegment`) — THIS section: + // approximate serve-from-disk (PQ in RAM + Vamana on NVMe), kept + // inert behind MOON_VEC_COLD_TIER until the M5 production gate + // (promote-back, delete, restart recovery, recall gate); an M3-exit + // review decides productionize-vs-delete on real EWMA telemetry. // The DiskANN cold tier is incomplete (no cold-segment deletion, ADC-only // recall, restart reload of PQ codebooks unfinished). The WARM->COLD // transition is a NO-OP unless an operator sets `MOON_VEC_COLD_TIER=1`; - // warm-tier mmap + LRU eviction handles out-of-RAM indexes by default. + // warm-tier byte-budget LRU demotion to COLD-stub handles out-of-RAM + // indexes by default. /// Seconds after last access before a WARM segment is promoted to COLD. /// Consumed by the cold-transition timer ONLY when the experimental cold /// tier is enabled (`MOON_VEC_COLD_TIER=1`); otherwise inert. @@ -450,7 +461,8 @@ pub struct ServerConfig { pub segment_cold_after: u64, /// Minimum queries-per-second threshold; segments below this are COLD candidates. - /// Not yet consumed — reserved for the experimental cold-tier heuristic. + /// [reserved: M3] — becomes the per-index EWMA boundary between COLD-stub + /// (idle) and COLD-ann (queried) demotion in the frequency classifier. #[arg(long = "segment-cold-min-qps", default_value_t = 0.1)] pub segment_cold_min_qps: f64, @@ -463,14 +475,14 @@ pub struct ServerConfig { pub memory_arenas_cap: u32, /// DiskANN beam width for disk-resident vector search. - /// Not yet consumed — reserved for the experimental cold-tier search path - /// (gated by MOON_VEC_COLD_TIER). + /// [reserved: M5] — consumed when the COLD-ann search path is + /// productionized (gated by MOON_VEC_COLD_TIER until then). #[arg(long = "vec-diskann-beam-width", default_value_t = 8)] pub vec_diskann_beam_width: u32, /// Number of HNSW upper levels cached in memory for DiskANN hybrid search. - /// Not yet consumed — reserved for the experimental cold-tier cache layer - /// (gated by MOON_VEC_COLD_TIER). + /// [reserved: M5] — consumed when the COLD-ann cache layer is + /// productionized (gated by MOON_VEC_COLD_TIER until then). #[arg(long = "vec-diskann-cache-levels", default_value_t = 3)] pub vec_diskann_cache_levels: u32, @@ -1030,6 +1042,22 @@ impl ServerConfig { Self::parse_size(&self.vec_warm_mmap_budget).unwrap_or(2 * 1024 * 1024 * 1024) } + /// Per-shard share of `--vec-warm-mmap-budget` (accounting-spine A5, + /// tiering-v2 D3): the flag is an INSTANCE-TOTAL cap divided across + /// shards, matching `maxmemory_per_shard` semantics. Previously each + /// shard applied the full value — an N-shard instance silently allowed + /// N× the configured WARM memory. `0` still disables enforcement; a + /// nonzero total floors at 1 byte per shard (0 would flip semantics to + /// "unlimited", the unsafe direction). Division floor is fine otherwise: + /// under-allocating a soft budget is the safe direction. + pub fn vec_warm_mmap_budget_bytes_per_shard(&self) -> u64 { + let total = self.vec_warm_mmap_budget_bytes(); + if total == 0 { + return 0; + } + (total / self.shards.max(1) as u64).max(1) + } + /// Returns the effective disk offload directory, falling back to --dir. pub fn effective_disk_offload_dir(&self) -> PathBuf { self.disk_offload_dir @@ -2435,6 +2463,41 @@ mod tests { assert_eq!(config.vec_diskann_cache_levels, 3); } + /// Accounting-spine A5 (tiering-v2 D3): `--vec-warm-mmap-budget` is an + /// INSTANCE-TOTAL cap divided across shards, matching + /// `maxmemory_per_shard` semantics. Previously each shard applied the + /// full value — an N-shard instance silently allowed N× the configured + /// WARM memory. RED until the per-shard accessor exists and divides. + #[test] + fn test_vec_warm_budget_divided_per_shard() { + let mut config = ServerConfig::parse_from::<[&str; 0], &str>([]); + + // 4 shards × default "2gb" ⇒ 512 MiB per shard. + config.shards = 4; + assert_eq!( + config.vec_warm_mmap_budget_bytes_per_shard(), + 512 * 1024 * 1024 + ); + + // Single shard: unchanged (full 2 GiB). + config.shards = 1; + assert_eq!( + config.vec_warm_mmap_budget_bytes_per_shard(), + 2 * 1024 * 1024 * 1024 + ); + + // "0" disables enforcement regardless of shard count. + config.vec_warm_mmap_budget = "0".to_string(); + config.shards = 8; + assert_eq!(config.vec_warm_mmap_budget_bytes_per_shard(), 0); + + // A tiny budget over many shards must floor at 1 byte, NOT 0 — + // 0 flips semantics to "unlimited", the unsafe direction. + config.vec_warm_mmap_budget = "100".to_string(); + config.shards = 128; + assert_eq!(config.vec_warm_mmap_budget_bytes_per_shard(), 1); + } + #[test] fn test_cold_tier_custom() { let config = ServerConfig::parse_from([ diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index 502210b8f..bf2fd7e50 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -659,8 +659,9 @@ impl super::Shard { // Per-shard warm-segment mmap budget enforcer. // Owned exclusively by this event-loop task; no locking needed. + // A5: the flag is an instance-total cap; each shard enforces its share. let mut warm_mmap_budget = crate::vector::persistence::mmap_budget::MmapBudget::new( - server_config.vec_warm_mmap_budget_bytes(), + server_config.vec_warm_mmap_budget_bytes_per_shard(), ); // Tokio path doesn't take these into the spawn signatures; suppress warnings. let (_, _, _) = (&spill_sender, &spill_file_id, &disk_offload_dir); diff --git a/src/shard/persistence_tick.rs b/src/shard/persistence_tick.rs index 8494a5259..ed7cce144 100644 --- a/src/shard/persistence_tick.rs +++ b/src/shard/persistence_tick.rs @@ -336,30 +336,13 @@ pub(crate) fn run_eviction_tick( // per 100ms tick. Siblings read the published snapshot on their own // ticks, so every budget is at most one tick stale — the same slack the // static scheme already has between eviction passes. - { - let rt = runtime_config.read(); - // C5 / Phase 3: compute per-shard KV memory via ShardSlice without - // lock acquisitions (avoids per-DB read locks; estimated_memory() is - // an O(1) accumulator read). Published unconditionally: MEMORY DOCTOR - // and the Prometheus KV gauge read this atomic even when maxmemory is - // unlimited — gating it on maxmemory > 0 left them at a permanent 0. - let used = crate::shard::slice::with_shard(|s| { - s.databases - .iter() - .map(|db| db.estimated_memory()) - .sum::() - }); - shard_databases.publish_memory(shard_id, used); - // Elastic budgets only exist under a finite maxmemory cap. - if rt.maxmemory > 0 { - shard_databases.recompute_elastic_budget(shard_id, &rt); - } - } - // C5 / M4: publish vector/text/graph store memory for lock-free observers. // 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. + // A4 review (LOW): published BEFORE the KV publish + elastic recompute + // below so the recompute's vector-aware donor/hot classification reads + // THIS tick's vector figure, not last tick's (siblings' figures remain + // ≤ 1 tick stale by design). let vector_resident_bytes = crate::shard::slice::with_shard(|s| { use std::sync::atomic::Ordering; let (mutable, immutable) = s.vector_store.resident_bytes(); @@ -385,6 +368,26 @@ pub(crate) fn run_eviction_tick( mutable + immutable }); + { + let rt = runtime_config.read(); + // C5 / Phase 3: compute per-shard KV memory via ShardSlice without + // lock acquisitions (avoids per-DB read locks; estimated_memory() is + // an O(1) accumulator read). Published unconditionally: MEMORY DOCTOR + // and the Prometheus KV gauge read this atomic even when maxmemory is + // unlimited — gating it on maxmemory > 0 left them at a permanent 0. + let used = crate::shard::slice::with_shard(|s| { + s.databases + .iter() + .map(|db| db.estimated_memory()) + .sum::() + }); + shard_databases.publish_memory(shard_id, used); + // Elastic budgets only exist under a finite maxmemory cap. + if rt.maxmemory > 0 { + shard_databases.recompute_elastic_budget(shard_id, &rt); + } + } + if server_config.disk_offload_enabled() && should_run_pressure_cascade( runtime_config, @@ -631,6 +634,16 @@ pub(crate) fn handle_memory_pressure( // budget (maxmemory/num_shards) so the summed eviction across shards bounds // aggregate RSS at the whole-instance maxmemory. // + // A3 review (MEDIUM): this used-term stays KV-only DELIBERATELY, unlike + // `timers::run_eviction` (the disk-offload-off path), which adds the + // shard's vector bytes. Inside the cascade the vector term has already + // pulled its weight: it fired the trigger (`should_run_pressure_cascade` + // is vector-inclusive) and step 2 sheds vector memory directly via + // offload-to-COLD. Adding it here too would evict KV to pay for memory + // that step 2 is already reclaiming more cheaply; the trigger refires + // every 100ms tick, so the cascade converges with vectors shedding + // first and KV eviction as the residual step. + // // When a SpillThread is available, use the async path: entries are removed // from DashTable immediately (freeing RAM) and pwrite is deferred to the // background thread. Otherwise, fall back to synchronous spill. diff --git a/src/shard/shared_databases.rs b/src/shard/shared_databases.rs index a1092af39..8be31064d 100644 --- a/src/shard/shared_databases.rs +++ b/src/shard/shared_databases.rs @@ -234,10 +234,21 @@ impl ShardDatabases { // SmallVec: most deployments run <=16 shards, so this 100ms-tick // snapshot stays fully on the stack; only larger shard counts spill // to a single heap allocation (still one per call, same as before). + // + // A4 (accounting spine, tiering-v2 D3): each shard's used-term is + // KV + published vector resident bytes. A vector-heavy/KV-light + // shard was misclassified as an idle donor — it lent headroom to + // siblings while its true footprint was already over base, and the + // pressure cascade then compared a vector-INCLUSIVE used against a + // budget inflated by that donation. Two Relaxed loads per shard. let used: SmallVec<[usize; 16]> = self .memory_per_shard .iter() - .map(|a| a.load(Ordering::Relaxed)) + .zip(self.store_memory_per_shard.iter()) + .map(|(kv, store)| { + kv.load(Ordering::Relaxed) + .saturating_add(store.vector.load(Ordering::Relaxed)) + }) .collect(); let budget = crate::storage::eviction::compute_elastic_budget(shard_id, base, &used); self.elastic_budgets[shard_id].store(budget, Ordering::Relaxed); @@ -892,6 +903,45 @@ mod tests { assert_eq!(shared.recompute_elastic_budget(1, &rt), 100); } + /// Accounting-spine A4 (tiering-v2 D3): the donor/hot classification must + /// see vector resident bytes. A vector-heavy/KV-light shard was + /// misclassified as an idle donor — it lent per-shard headroom to + /// siblings while its true resident footprint (KV + vector) was already + /// over base, inflating the budget the pressure cascade later compares a + /// vector-INCLUSIVE used-term against. RED until `recompute_elastic_budget` + /// adds the published vector bytes to its `used` snapshot. + #[test] + fn recompute_elastic_budget_vector_heavy_shard_not_donor() { + let shared = new_shared(4, 1); + let rt = rt_config(400, 4); // base = 100 per shard + + shared.publish_memory(0, 120); // hot + shared.publish_memory(1, 10); // KV-light... + shared.store_memory_per_shard[1] + .vector + .store(200, Ordering::Relaxed); // ...but 200 of vector RAM + shared.publish_memory(2, 10); + shared.publish_memory(3, 10); + + // KV-blind math: shard 1 looks idle (10 < 100) and donates 90 — + // surplus 270, one hot shard ⇒ shard 0 budget 370. + // Vector-aware: shard 1's true used is 210 > base — it is HOT, not + // a donor. Surplus = 90 + 90 (shards 2,3), split across the two hot + // shards ⇒ 100 + 180/2 = 190 each. + assert_eq!( + shared.recompute_elastic_budget(0, &rt), + 190, + "vector-heavy shard must not be classified as an idle donor" + ); + assert_eq!( + shared.recompute_elastic_budget(1, &rt), + 190, + "the vector-heavy shard itself is hot and shares the pool" + ); + // True idle shards keep base. + assert_eq!(shared.recompute_elastic_budget(2, &rt), 100); + } + #[test] fn recompute_elastic_budget_correct_beyond_smallvec_inline_capacity() { // The per-call `used` snapshot is a `SmallVec<[usize; 16]>` — pin diff --git a/src/shard/timers.rs b/src/shard/timers.rs index ff8afe575..0d4f49d15 100644 --- a/src/shard/timers.rs +++ b/src/shard/timers.rs @@ -46,10 +46,51 @@ pub(crate) fn run_eviction( if rt.maxmemory > 0 { // GAP-1: enforce the elastic budget (0 = static fallback inside). let budget = shard_databases.elastic_budget(shard_id); + // A3 (accounting spine, tiering-v2 D3): vector segment memory counts + // toward maxmemory — it is data-plane memory. Published by the same + // 100ms tick (`run_eviction_tick` stores it before calling us; lag + // ≤ 1 tick, documented acceptable). One Relaxed load per tick, no + // recompute on this path. + // + // REVIEW FIX (CRITICAL, perf+security reviews concurring): the + // used-term is an AGGREGATE — Σ all dbs' KV + the shard's vector + // bytes, computed ONCE — never the shard-wide vector term re-added + // to each logical db's independent check. Per-db duplication both + // under-detected (KV spread across dbs, no single db + vector over + // budget ⇒ nothing evicted while aggregate RSS overran) and + // over-evicted (vector term over budget ⇒ every db drained instead + // of stopping at the aggregate target). + // + // Semantics when the un-evictable vector term alone exceeds the + // budget: KV drains then errors OOM — the vector tier legitimately + // owns that memory (shared-budget, same as one tenant's KV growth + // evicting siblings under allkeys-lru); per-db quotas (WS5b) remain + // the tenant-isolation mechanism. The pressure cascade (which CAN + // shrink vectors via offload to COLD) fires earlier at + // `--disk-offload-threshold` (0.85×budget), so with disk-offload + // enabled vectors shed first. + let vector_bytes = shard_databases.store_memory_per_shard[shard_id] + .vector + .load(std::sync::atomic::Ordering::Relaxed); let db_count = shard_databases.db_count(); + let mut kv_total: usize = 0; + for i in 0..db_count { + kv_total = kv_total.saturating_add(crate::shard::slice::with_shard_db(i, |db| { + db.estimated_memory() + })); + } + // Running aggregate: each db's eviction reduces it by what that db + // actually freed; once it drops to the budget, remaining dbs see an + // under-budget total and return immediately (no eviction). + let mut remaining = kv_total.saturating_add(vector_bytes); for i in 0..db_count { crate::shard::slice::with_shard_db(i, |db| { - let _ = crate::storage::eviction::try_evict_if_needed_budget(db, &rt, budget); + let before = db.estimated_memory(); + let _ = crate::storage::eviction::try_evict_if_needed_with_spill_and_total_budget( + db, &rt, None, remaining, budget, + ); + let freed = before.saturating_sub(db.estimated_memory()); + remaining = remaining.saturating_sub(freed); }); } } @@ -370,3 +411,125 @@ pub(crate) fn run_mvcc_sweep( ); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::storage::db::Database; + use bytes::Bytes; + use std::sync::atomic::Ordering; + + /// A3 review fix (CRITICAL, both adversarial reviews): the shard-wide + /// vector term must gate an AGGREGATE check (Σ all dbs' KV + vector), + /// not be re-added to every logical db's independent check. The per-db + /// duplication had two failure modes: + /// 1. under-detection — KV spread across dbs where no single + /// `db_i + vector` crosses budget ⇒ nothing evicts while true + /// aggregate RSS overruns (the very bug A3 claimed to fix); + /// 2. over-eviction — a vector term over budget independently drained + /// EVERY db on the shard (multi-tenant blast radius), instead of + /// evicting only until the aggregate is back under budget. + /// This test pins both: two dbs each individually under budget with the + /// vector term, aggregate over — eviction must fire (kills mode 1) and + /// must stop once the aggregate is satisfied, leaving db 1 untouched + /// (kills mode 2). + #[test] + fn test_run_eviction_gates_on_aggregate_not_per_db() { + let dbs = vec![vec![Database::new(), Database::new()]]; + let (shared, mut inits) = ShardDatabases::new(dbs); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); + + // ~256 KiB of KV per db (64 keys × 4 KiB values). + for db_idx in 0..2 { + crate::shard::slice::with_shard_db(db_idx, |db| { + for i in 0..64u32 { + db.set_string( + Bytes::from(format!("db{db_idx}:key:{i}")), + Bytes::from(vec![b'v'; 4096]), + ); + } + }); + } + + let mut rt = RuntimeConfig::default(); + rt.maxmemory = 1024 * 1024; // 1 MiB budget, 1 shard + rt.num_shards = 1; + rt.maxmemory_policy = "allkeys-lru".to_string(); + let runtime_config = Arc::new(parking_lot::RwLock::new(rt)); + + // 700 KiB vector term: each db alone is ~256K + 700K < 1 MiB (the + // per-db check would pass ⇒ old code evicts NOTHING), but the + // aggregate ~512K + 700K > 1 MiB must evict. + shared.store_memory_per_shard[0] + .vector + .store(700 * 1024, Ordering::Relaxed); + run_eviction(&shared, 0, &runtime_config); + + let len0 = crate::shard::slice::with_shard_db(0, |db| db.len()); + let len1 = crate::shard::slice::with_shard_db(1, |db| db.len()); + assert!( + len0 < 64, + "aggregate over budget must evict even when no single db crosses it (db0 {len0})" + ); + // The overage (~190 KiB) is smaller than db0's ~256 KiB, so eviction + // must satisfy the aggregate within db0 and never touch db1. + assert_eq!( + len1, 64, + "eviction must stop at the aggregate target, not drain sibling dbs" + ); + } + + /// A3 (accounting spine): published vector resident bytes must count + /// toward the background maxmemory eviction check. A vector-heavy shard + /// whose KV alone is under budget previously never evicted — vector RAM + /// was invisible to `--maxmemory`, so a pure-vector workload could drive + /// RSS to OOM while eviction reported "under budget". RED until + /// `run_eviction` adds the published vector term to the used total. + /// + /// Note on semantics: when the un-evictable vector term alone exceeds + /// the budget, KV on the shard fully drains then errors OOM — the vector + /// tier legitimately owns that memory (shared-budget semantics, same as + /// one tenant's KV growth evicting siblings under allkeys-lru), and the + /// pressure cascade (which CAN shrink vectors via offload) fires earlier + /// at 0.85×budget. Per-db tenant isolation remains the per-db quota + /// mechanism (WS5b), not --maxmemory. + #[test] + fn test_run_eviction_counts_vector_memory() { + let dbs = vec![vec![Database::new()]]; + let (shared, mut inits) = ShardDatabases::new(dbs); + crate::shard::slice::reset_test_shard(crate::shard::slice::ShardSlice::new( + inits.remove(0), + )); + + crate::shard::slice::with_shard_db(0, |db| { + for i in 0..100u32 { + db.set_string(Bytes::from(format!("key:{i}")), Bytes::from(vec![b'v'; 64])); + } + }); + + let mut rt = RuntimeConfig::default(); + rt.maxmemory = 1024 * 1024; + rt.num_shards = 1; + rt.maxmemory_policy = "allkeys-lru".to_string(); + let runtime_config = Arc::new(parking_lot::RwLock::new(rt)); + + // KV alone (~tens of KB) is far under the 1 MiB budget: no eviction. + run_eviction(&shared, 0, &runtime_config); + let before = crate::shard::slice::with_shard_db(0, |db| db.len()); + assert_eq!(before, 100, "KV under budget must not evict"); + + // 2 MiB of published vector-segment memory pushes the shard over + // budget: KV must now be evicted (vector memory is data-plane memory). + shared.store_memory_per_shard[0] + .vector + .store(2 * 1024 * 1024, Ordering::Relaxed); + run_eviction(&shared, 0, &runtime_config); + let after = crate::shard::slice::with_shard_db(0, |db| db.len()); + assert!( + after < before, + "vector bytes over budget must trigger KV eviction ({after} vs {before})" + ); + } +} diff --git a/src/vector/diskann/pq.rs b/src/vector/diskann/pq.rs index b549e8105..d7ad955ac 100644 --- a/src/vector/diskann/pq.rs +++ b/src/vector/diskann/pq.rs @@ -192,6 +192,12 @@ impl ProductQuantizer { pub fn dim(&self) -> usize { self.dim } + + /// Estimated resident heap bytes of the trained quantizer: the flat + /// codebook (`m * ksub * dsub` floats) plus fixed struct overhead. + pub fn resident_bytes(&self) -> usize { + std::mem::size_of::() + self.centroids.len() * std::mem::size_of::() + } } /// Scalar squared-L2 for sub-vectors. diff --git a/src/vector/diskann/segment.rs b/src/vector/diskann/segment.rs index 427412a2f..731ae4148 100644 --- a/src/vector/diskann/segment.rs +++ b/src/vector/diskann/segment.rs @@ -489,6 +489,16 @@ impl DiskAnnSegment { self.num_vectors } + /// Estimated resident heap bytes of this cold segment (accounting-spine + /// A1): PQ codes (`num_vectors * m` bytes, kept in RAM) + the trained + /// codebook, plus fixed struct overhead. The Vamana graph lives on NVMe + /// and is deliberately excluded; the Linux io_uring read buffers are + /// small, bounded, and also excluded. Previously this tier reported 0 — + /// untracked resident memory (D9 quarantine). + pub fn resident_bytes(&self) -> usize { + std::mem::size_of::() + self.pq_codes.len() + self.pq.resident_bytes() + } + /// Maximum graph degree (R parameter). #[inline] pub fn max_degree(&self) -> u32 { @@ -591,6 +601,25 @@ mod tests { (seg, vectors, tmp) } + /// Accounting-spine A1: DiskANN cold segments keep PQ codes + codebook in + /// RAM but contributed a hardcoded 0 to `SegmentHolder::resident_bytes()` + /// — untracked resident memory. RED until `resident_bytes()` exists and + /// counts at least the PQ-code floor plus a non-empty codebook. + #[test] + fn test_diskann_resident_bytes_accounts_pq() { + let n = 64usize; + let m = 4usize; + let (seg, _, _tmp) = build_test_segment(n, 16, m, 8); + let rb = seg.resident_bytes(); + // pq_codes alone are n*m bytes; the trained codebook (m*ksub*dsub + // floats) must push the total strictly above that floor. + assert!( + rb > n * m, + "DiskANN resident_bytes ({rb}) must count pq_codes + codebook (> {})", + n * m + ); + } + #[test] fn test_diskann_segment_search_recall() { let n = 50; diff --git a/src/vector/segment/holder.rs b/src/vector/segment/holder.rs index c037e89c9..16a956e61 100644 --- a/src/vector/segment/holder.rs +++ b/src/vector/segment/holder.rs @@ -354,8 +354,9 @@ impl SegmentHolder { /// so once segments age past `--segment-warm-after` it dominates resident /// vector memory — it MUST be counted here or the memory-pressure trigger /// (see `should_run_pressure_cascade`) and INFO/Prometheus go blind to it. - /// IVF and DiskANN-cold segments have no resident accessor yet (IVF is - /// in-memory but small; cold lives on disk) and still contribute 0. + /// IVF (centroids + posting lists) and DiskANN-cold (PQ codes + codebook; + /// graph on NVMe excluded) are counted too (accounting-spine A1) — every + /// tier that pins heap reports it here. pub fn resident_bytes(&self) -> (usize, usize) { let snapshot = self.load(); let mutable = snapshot.mutable.resident_bytes(); @@ -372,6 +373,15 @@ impl SegmentHolder { for stub in &snapshot.unloaded { immutable += stub.resident_bytes(); } + // IVF: in-memory centroids + interleaved posting lists. + for ivf_seg in &snapshot.ivf { + immutable += ivf_seg.resident_bytes(); + } + // DiskANN cold (COLD-ann, distinct from the COLD stubs above): PQ + // codes + codebook stay resident; the Vamana graph lives on disk. + for cold_seg in &snapshot.cold { + immutable += cold_seg.resident_bytes(); + } (mutable, immutable) } @@ -1336,6 +1346,83 @@ mod tests { ); } + /// Accounting-spine A1: the IVF and DiskANN-cold tiers hold real heap + /// memory (IVF: centroids + posting lists; DiskANN: PQ codes + codebook) + /// but contributed a hardcoded 0 to the roll-up — blinding the pressure + /// trigger and observability for those tiers. RED until the roll-up sums + /// both tiers' `resident_bytes()`. + #[test] + fn test_resident_bytes_includes_ivf_and_cold_tiers() { + use crate::vector::diskann::pq::ProductQuantizer; + use crate::vector::diskann::segment::DiskAnnSegment; + use crate::vector::segment::ivf; + + distance::init(); + let dim = 8usize; + let pdim = padded_dimension(dim as u32) as usize; + let dim_half = pdim / 2; + + // Small IVF segment (20 vectors, 2 clusters). + let n = 20; + let mut sign_flips = vec![1.0f32; pdim]; + for (i, s) in sign_flips.iter_mut().enumerate() { + if i % 3 == 0 { + *s = -1.0; + } + } + let mut vectors = Vec::with_capacity(n * dim); + let mut tq_codes = Vec::with_capacity(n); + let mut norms = Vec::with_capacity(n); + let ids: Vec = (1000..1000 + n as u32).collect(); + for i in 0..n { + let v: Vec = (0..dim).map(|d| (i * dim + d) as f32 * 0.01).collect(); + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + norms.push(if norm > 0.0 { norm } else { 1.0 }); + vectors.extend_from_slice(&v); + tq_codes.push(vec![(i & 0xF) as u8; dim_half]); + } + let ivf_seg = + ivf::build_ivf_segment(&vectors, &tq_codes, &norms, &ids, dim, 2, &sign_flips); + let ivf_bytes = ivf_seg.resident_bytes(); + assert!(ivf_bytes > 0); + + // Small DiskANN cold segment: PQ trained on the same vectors; the + // vamana graph lives on disk (constructor only opens the file). + let tmp = tempfile::tempdir().unwrap(); + let vamana_path = tmp.path().join("vamana.mpf"); + std::fs::write(&vamana_path, vec![0u8; 4096]).unwrap(); + let m = 4usize; + let pq = ProductQuantizer::train(&vectors, dim, m, 8); + let mut pq_codes = Vec::with_capacity(n * m); + for i in 0..n { + pq_codes.extend_from_slice(&pq.encode(&vectors[i * dim..(i + 1) * dim])); + } + let cold_seg = + DiskAnnSegment::new(pq_codes, pq, vamana_path, dim, n as u32, 0, 8, 1).unwrap(); + let cold_bytes = cold_seg.resident_bytes(); + assert!(cold_bytes > 0); + + let collection = make_test_collection(dim as u32); + let holder = SegmentHolder::new(dim as u32, collection.clone()); + let (_, base_imm) = holder.resident_bytes(); + + holder.swap(SegmentList { + mutable: Arc::new(MutableSegment::new(dim as u32, collection)), + immutable: Vec::new(), + ivf: vec![Arc::new(ivf_seg)], + warm: Vec::new(), + cold: vec![Arc::new(cold_seg)], + unloaded: Vec::new(), + }); + + let (_, imm) = holder.resident_bytes(); + assert!( + imm >= base_imm + ivf_bytes + cold_bytes, + "IVF + cold tiers must contribute to resident_bytes: {imm} \ + (base {base_imm}, ivf {ivf_bytes}, cold {cold_bytes})" + ); + } + /// #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`. diff --git a/src/vector/segment/ivf.rs b/src/vector/segment/ivf.rs index b133a3ea3..fa9dc43cf 100644 --- a/src/vector/segment/ivf.rs +++ b/src/vector/segment/ivf.rs @@ -226,6 +226,25 @@ impl IvfSegment { self.posting_lists.iter().map(|pl| pl.count as u64).sum() } + /// Estimated resident heap bytes of this segment (accounting-spine A1). + /// + /// Counts the centroid table, FWHT sign flips, and every posting list's + /// interleaved codes + ids + norms, plus fixed struct overhead. Feeds + /// `SegmentHolder::resident_bytes()` → memory-pressure trigger, + /// MEMORY DOCTOR, and Prometheus — previously this tier reported 0. + pub fn resident_bytes(&self) -> usize { + let mut bytes = std::mem::size_of::(); + bytes += self.centroids.len() * std::mem::size_of::(); + bytes += self.sign_flips.len() * std::mem::size_of::(); + for pl in &self.posting_lists { + bytes += std::mem::size_of::(); + bytes += pl.codes.len(); // u8 codes + bytes += pl.ids.len() * std::mem::size_of::(); + bytes += pl.norms.len() * std::mem::size_of::(); + } + bytes + } + /// Reference to the FWHT sign flips for query rotation. #[inline] pub fn sign_flips(&self) -> &[f32] { @@ -1047,6 +1066,45 @@ mod tests { assert_eq!(seg.dimension(), dim as u32); } + /// Accounting-spine A1: IVF segments hold real heap memory (centroids, + /// interleaved posting-list codes, ids, norms) but contributed a hardcoded + /// 0 to `SegmentHolder::resident_bytes()` — invisible to the memory-pressure + /// trigger, MEMORY DOCTOR, and Prometheus. RED until `resident_bytes()` + /// exists and counts at least the known heap floor. + #[test] + fn test_ivf_resident_bytes_accounts_heap() { + crate::vector::distance::init(); + let dim = 8; + let pdim = padded_dimension(dim as u32) as usize; + let dim_half = pdim / 2; + let n = 100; + let n_clusters = 4; + let signs = test_sign_flips(pdim, 42); + + let mut vectors = Vec::with_capacity(n * dim); + let mut tq_codes = Vec::with_capacity(n); + let mut norms = Vec::with_capacity(n); + let ids: Vec = (0..n as u32).collect(); + for i in 0..n { + let v = det_f32(dim, i as u64 + 1); + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + norms.push(norm); + vectors.extend_from_slice(&v); + tq_codes.push(vec![(i & 0xFF) as u8; dim_half]); + } + let seg = build_ivf_segment(&vectors, &tq_codes, &norms, &ids, dim, n_clusters, &signs); + + // Provable heap floor: interleaved codes are >= n * dim_half bytes + // (32-block round-up only adds), ids + norms are 8 bytes/vector, and + // centroids are n_clusters * pdim floats. + let floor = n * dim_half + n * 8 + n_clusters * pdim * 4; + let rb = seg.resident_bytes(); + assert!( + rb >= floor, + "IVF resident_bytes ({rb}) must count codes+ids+norms+centroids (>= {floor})" + ); + } + #[test] fn test_recall_at_10_nprobe_32() { // Recall test: 10K vectors from 256 synthetic Gaussian clusters.