src/cache_budget.rs
//! **One byte-budget policy for every in-process cache that holds entries of
//! unbounded size.**
//!
//! ## The class this exists to kill
//!
//! A cache capped by ENTRY COUNT bounds nothing when one entry can be a whole
//! decoded table, a whole snapshot's vectors, or a whole repo's call graph. The
//! symptom is a monotone RSS climb with no error, no log line and no ceiling to
//! hit — indistinguishable from a leak until you read the eviction policy.
//! The warehouse working-set cache (removed with Mímir v2) learned it the
//! expensive way and its note recorded the price: *13 entries at 5M rows each blew 44 GB under a count
//! cap of 64*. Three sibling caches copied only the front half of that lesson
//! (`vector::store`'s `INDEX_CACHE`, `knowledge::query`'s `VIEW_CACHE`, and —
//! in another crate — skade's table-handle cache).
//!
//! So the budget lives HERE, once, and the caches are callers. A count cap is
//! still allowed, but it is a *convenience* bound on lookup cost; the memory
//! bound is always [`Budget::max_bytes`].
//!
//! ## Charge on USE, not only on insert
//!
//! Several entries in this codebase GROW after admission, inside an already-
//! shared `Arc`:
//!
//! * `Reconstructed` builds its HNSW graph lazily in a `OnceLock` — and the
//! graph keeps its own copy of the `n * dim` float matrix, so the entry can
//! more than double.
//! * `VectorIndex` builds its int8 matrix lazily in a second `OnceLock`.
//! * `KnowledgeView` builds its `CallIndex` — and then two `stree` inverted
//! indexes inside THAT — lazily on first query.
//!
//! A budget charged only at insert under-counts all three. [`BudgetedCache::recharge`]
//! is the answer: after a read path may have grown an entry, re-weigh it and let
//! the budget evict on the true figure. A cache whose entries can grow and which
//! never calls `recharge` is back to being count-capped in disguise.
//!
//! ## Contract
//!
//! * Eviction is LRU (`get` promotes to the back, eviction takes the front).
//! * Every eviction and every refused-too-big insert is LOGGED with the entry
//! label, its size and the knob that would have kept it. A cache that silently
//! drops the caller's data is the house disease, not a fix.
//! * An entry larger than the whole budget is SERVED to its caller and never
//! retained (`refused_too_big`), so a single huge decode can never pin the
//! cache.
//! * The budget may evict down to ZERO entries. A budget that cannot evict its
//! own smallest entry is the blind-guard version of this fix.
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
/// A byte budget shared by one or more [`BudgetedCache`]s, plus the counters the
/// tests and the warm summary read.
///
/// `env_var` names the operator knob (parsed by
/// [`parse_bytes`]: a plain byte count, or
/// `<n>k` / `<n>m` / `<n>g`). A test override wins over the environment.
pub struct Budget {
label: &'static str,
env_var: &'static str,
default_bytes: usize,
override_bytes: AtomicUsize,
/// Bytes resident across EVERY cache on this budget.
resident: AtomicUsize,
evictions: AtomicUsize,
refused_too_big: AtomicUsize,
}
impl Budget {
pub const fn new(label: &'static str, env_var: &'static str, default_bytes: usize) -> Self {
Self {
label,
env_var,
default_bytes,
override_bytes: AtomicUsize::new(0),
resident: AtomicUsize::new(0),
evictions: AtomicUsize::new(0),
refused_too_big: AtomicUsize::new(0),
}
}
/// The operator knob's name — quoted in every cap-hit log line so the reader
/// is told what to raise.
pub fn env_var(&self) -> &'static str {
self.env_var
}
/// The budget in force. A test override wins; otherwise the environment;
/// otherwise the default. A mis-typed knob falls back LOUDLY (never to zero,
/// which would be the silent opposite of the intent).
pub fn max_bytes(&self) -> usize {
match self.override_bytes.load(Ordering::Relaxed) {
0 => env_bytes(self.env_var, self.default_bytes),
n => n,
}
}
/// Test/bench hook: pin the budget (`0` restores env/default).
pub fn set_for_test(&self, n: usize) {
self.override_bytes.store(n, Ordering::Relaxed);
}
/// Bytes currently resident across every cache on this budget.
pub fn resident_bytes(&self) -> usize {
self.resident.load(Ordering::Relaxed)
}
/// Entries evicted for being over budget (or over a count cap) since start.
pub fn evictions(&self) -> usize {
self.evictions.load(Ordering::Relaxed)
}
/// Inserts served but NOT retained because one entry exceeded the whole budget.
pub fn refused_too_big(&self) -> usize {
self.refused_too_big.load(Ordering::Relaxed)
}
fn add(&self, n: usize) {
self.resident.fetch_add(n, Ordering::Relaxed);
}
fn sub(&self, n: usize) {
// Saturating: two caches on one budget update independently, and an
// underflow here would wrap to usize::MAX and evict everything forever.
let _ = self
.resident
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
Some(v.saturating_sub(n))
});
}
}
/// Parse a byte budget: a plain integer, or an integer with a `k`/`m`/`g` suffix
/// (binary units, case-insensitive, optional trailing `b`/`ib`). `None` on junk —
/// a mis-typed knob must fall back to the default LOUDLY at the read site, never
/// to zero (zero would mean "cache nothing", the silent opposite of the intent).
pub fn parse_bytes(s: &str) -> Option<usize> {
let t = s.trim().to_ascii_lowercase();
let t = t
.strip_suffix("ib")
.or_else(|| t.strip_suffix('b'))
.unwrap_or(&t);
let (num, mult): (&str, usize) = match t.chars().last()? {
'k' => (&t[..t.len() - 1], 1 << 10),
'm' => (&t[..t.len() - 1], 1 << 20),
'g' => (&t[..t.len() - 1], 1 << 30),
c if c.is_ascii_digit() => (t, 1),
_ => return None,
};
let n: usize = num.trim().parse().ok()?;
n.checked_mul(mult)
}
fn env_bytes(var: &str, default: usize) -> usize {
match std::env::var(var) {
Ok(v) => match parse_bytes(&v) {
Some(n) if n > 0 => n,
_ => {
eprintln!(
"nornir: {var}={v:?} is not a byte size (use e.g. 4g / 512m / 1073741824) — \
using the default {default}"
);
default
}
},
Err(_) => default,
}
}
/// Human byte figure for the log lines: `NNN B` / `N.N KiB` / `N.N MiB` / `N.NN GiB`.
pub fn human(b: usize) -> String {
const K: f64 = 1024.0;
let f = b as f64;
if f < K {
format!("{b} B")
} else if f < K * K {
format!("{:.1} KiB", f / K)
} else if f < K * K * K {
format!("{:.1} MiB", f / (K * K))
} else {
format!("{:.2} GiB", f / (K * K * K))
}
}
struct Entry<K, V> {
key: K,
/// What this entry is currently charged at. Updated by [`BudgetedCache::recharge`].
bytes: usize,
value: Arc<V>,
}
/// A point-in-time reading of one cache.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct CacheStats {
pub entries: usize,
pub bytes: usize,
}
/// An LRU cache bounded by BYTES (and, secondarily, by an entry count).
///
/// `K` is matched by closure rather than `Eq` so compound keys and
/// generation-supersede rules (drop this table's other snapshots) express
/// naturally; `V` is always handed out behind an `Arc`.
pub struct BudgetedCache<K, V> {
/// Named in every log line ("vector index", "knowledge view", …).
label: &'static str,
budget: &'static Budget,
/// Secondary bound on lookup cost. NOT the memory bound.
count_cap: usize,
entries: Mutex<Vec<Entry<K, V>>>,
/// This cache's own contribution to `budget.resident`.
own_bytes: AtomicUsize,
}
impl<K, V> BudgetedCache<K, V> {
pub const fn new(label: &'static str, budget: &'static Budget, count_cap: usize) -> Self {
Self {
label,
budget,
count_cap,
entries: Mutex::new(Vec::new()),
own_bytes: AtomicUsize::new(0),
}
}
/// The budget this cache charges against.
pub fn budget(&self) -> &'static Budget {
self.budget
}
/// Cached value whose key satisfies `hit`, promoted to most-recently-used.
pub fn get(&self, hit: impl Fn(&K) -> bool) -> Option<Arc<V>> {
let mut c = self.entries.lock().unwrap();
let pos = c.iter().position(|e| hit(&e.key))?;
let entry = c.remove(pos);
let value = entry.value.clone();
c.push(entry); // MRU → back
Some(value)
}
/// Insert `value` at `bytes`, after dropping every entry whose key satisfies
/// `supersedes` (the generation rule — e.g. "this repo's other snapshots").
///
/// An entry bigger than the whole budget is NOT retained; the caller already
/// holds it, and pinning it would put the cache permanently over budget.
pub fn put(&self, key: K, bytes: usize, value: Arc<V>, supersedes: impl Fn(&K) -> bool) {
let budget = self.budget.max_bytes();
let mut c = self.entries.lock().unwrap();
let mut freed = 0usize;
c.retain(|e| {
if supersedes(&e.key) {
freed += e.bytes;
false
} else {
true
}
});
if freed > 0 {
self.charge_sub(freed);
}
if bytes > budget {
self.budget.refused_too_big.fetch_add(1, Ordering::Relaxed);
eprintln!(
"nornir: {} cache: an entry of {} is larger than the whole {} budget \
({}); served but NOT retained — the next reader rebuilds it",
self.label,
human(bytes),
human(budget),
self.budget.env_var
);
return;
}
c.push(Entry { key, bytes, value });
self.charge_add(bytes);
self.enforce(&mut c, budget);
}
/// **Charge on use.** Re-weigh the entry whose key satisfies `hit` at
/// `bytes` and evict if the cache is now over budget. Call this after any
/// read path that may have grown an entry in place (a lazily built graph,
/// index or quantized matrix inside an already-admitted `Arc`).
///
/// Does nothing when the key is absent — an entry that was already evicted
/// is not resurrected.
pub fn recharge(&self, hit: impl Fn(&K) -> bool, bytes: usize) {
let budget = self.budget.max_bytes();
let mut c = self.entries.lock().unwrap();
let Some(pos) = c.iter().position(|e| hit(&e.key)) else {
return;
};
let was = c[pos].bytes;
if was == bytes {
return;
}
c[pos].bytes = bytes;
if bytes > was {
self.charge_add(bytes - was);
} else {
self.charge_sub(was - bytes);
}
self.enforce(&mut c, budget);
}
/// Drop every entry. Test/bench hook — forces the cold rebuild path.
pub fn clear(&self) {
let mut c = self.entries.lock().unwrap();
c.clear();
let held = self.own_bytes.swap(0, Ordering::Relaxed);
self.budget.sub(held);
}
/// The bytes ONE entry is currently charged, if it is resident.
///
/// The per-entry counterpart to [`Self::stats`]. A guard that asserts a
/// whole-cache total is at the mercy of anything else sharing the process;
/// this asks about the entry the caller actually put in.
pub fn charge_of(&self, hit: impl Fn(&K) -> bool) -> Option<usize> {
let c = self.entries.lock().unwrap();
c.iter().find(|e| hit(&e.key)).map(|e| e.bytes)
}
/// Entries and bytes currently resident in THIS cache.
pub fn stats(&self) -> CacheStats {
let c = self.entries.lock().unwrap();
CacheStats {
entries: c.len(),
bytes: c.iter().map(|e| e.bytes).sum(),
}
}
fn charge_add(&self, n: usize) {
self.own_bytes.fetch_add(n, Ordering::Relaxed);
self.budget.add(n);
}
fn charge_sub(&self, n: usize) {
let _ = self
.own_bytes
.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |v| {
Some(v.saturating_sub(n))
});
self.budget.sub(n);
}
/// Evict LRU-first until the SHARED budget and this cache's count cap both
/// hold. May empty the cache — a budget below one entry must still be a
/// budget, not a floor of one.
fn enforce(&self, c: &mut Vec<Entry<K, V>>, budget: usize) {
while !c.is_empty()
&& (c.len() > self.count_cap || self.budget.resident_bytes() > budget)
{
let gone = c.remove(0);
self.charge_sub(gone.bytes);
self.budget.evictions.fetch_add(1, Ordering::Relaxed);
eprintln!(
"nornir: {} cache: evicted an entry of {} — over the {} budget ({})",
self.label,
human(gone.bytes),
human(budget),
self.budget.env_var
);
}
}
}
/// Serialises tests that pin a process-global budget or clear a process-global
/// cache ([`crate::vector::store`]'s index cache, [`crate::knowledge::query`]'s
/// view cache). `cargo test` runs in
/// parallel; without this two tests fight over the same static and the failure
/// looks exactly like a broken budget, which is the worst possible false signal
/// for a guard whose whole job is to be believable.
#[cfg(test)]
pub(crate) static CACHE_TEST_LOCK: Mutex<()> = Mutex::new(());
/// Take [`CACHE_TEST_LOCK`], surviving a poisoning from an earlier failure.
#[cfg(test)]
pub(crate) fn exclusive() -> std::sync::MutexGuard<'static, ()> {
CACHE_TEST_LOCK.lock().unwrap_or_else(|p| p.into_inner())
}
#[cfg(test)]
mod tests {
use super::*;
/// RED WITHOUT THE BUDGET: entries whose sizes differ by four orders of
/// magnitude must be bounded in BYTES, not in count. With a count cap alone
/// (the defect this module exists to kill) all four fat entries stay resident
/// and `bytes` lands at 4 MiB — 32× over the budget — while `entries` looks
/// perfectly healthy at 4 of a cap of 64.
#[test]
fn a_count_cap_would_not_have_bounded_this() {
// Each test owns its own budget + cache, so they cannot interfere under
// `cargo test`'s default parallelism.
static B: Budget = Budget::new("t1", "NORNIR_CACHE_BUDGET_T1", 1 << 30);
let cache: BudgetedCache<u32, Vec<u8>> = BudgetedCache::new("test", &B, 64);
B.set_for_test(128 * 1024); // 128 KiB
// 100 tiny entries: 100 B each. Count 100 ≫ nothing; bytes 10 KB ≪ budget.
for i in 0..100u32 {
cache.put(i, 100, Arc::new(vec![0u8; 100]), |k| *k == i);
}
let tiny = cache.stats();
assert!(
tiny.bytes <= 128 * 1024,
"100 tiny entries must fit the budget, got {}",
tiny.bytes
);
// 4 fat entries at 64 KiB each — 655x a tiny one, and each comfortably
// UNDER the whole budget, so they are admitted and must EVICT rather
// than be refused at the door. (The first version of this guard used
// 1 MiB entries, which were bigger than the budget itself: they took the
// refusal path, nothing was ever evicted, and the test passed with
// eviction disabled entirely. Refusal and eviction are different paths.)
for i in 1000..1004u32 {
cache.put(i, 64 * 1024, Arc::new(vec![0u8; 64 * 1024]), |k| *k == i);
}
let s = cache.stats();
assert!(
s.bytes <= 128 * 1024,
"BYTES must be bounded by the budget, got {} resident in {} entries \
(a count cap would have left 4 MiB here)",
human(s.bytes),
s.entries
);
assert!(
s.entries < 64,
"the byte budget, not the count cap, must be what evicted: {} entries \
(a 64-entry cap would have held 64 of these, i.e. up to 4 MiB)",
s.entries
);
cache.clear();
}
/// A budget BELOW one entry must still evict — down to zero. A cache that
/// keeps its last entry no matter what has a floor, not a budget, and that
/// floor is exactly the unbounded-entry hole.
#[test]
fn a_budget_below_one_entry_evicts_to_empty() {
static B: Budget = Budget::new("t2", "NORNIR_CACHE_BUDGET_T2", 1 << 30);
let cache: BudgetedCache<u32, Vec<u8>> = BudgetedCache::new("test", &B, 64);
B.set_for_test(4096);
cache.put(1, 1024, Arc::new(vec![0u8; 1024]), |k| *k == 1);
assert_eq!(cache.stats().entries, 1);
// Now the entry GROWS in place past the budget — the charge-on-use path.
cache.recharge(|k| *k == 1, 1 << 20);
assert_eq!(
cache.stats().entries,
0,
"an entry that grew past the whole budget must be evicted, not kept as a floor"
);
assert_eq!(cache.stats().bytes, 0);
assert_eq!(B.resident_bytes(), 0);
cache.clear();
}
/// Charge-on-use: an entry admitted small and grown large must be charged at
/// its GROWN size. Without `recharge` the cache still reports the admission
/// figure and the budget is blind to the growth.
#[test]
fn recharge_sees_growth_after_admission() {
static B: Budget = Budget::new("t3", "NORNIR_CACHE_BUDGET_T3", 1 << 30);
let cache: BudgetedCache<u32, Vec<u8>> = BudgetedCache::new("test", &B, 64);
B.set_for_test(64 * 1024 * 1024);
cache.put(7, 1024, Arc::new(Vec::new()), |k| *k == 7);
assert_eq!(cache.stats().bytes, 1024);
cache.recharge(|k| *k == 7, 8 * 1024 * 1024);
assert_eq!(
cache.stats().bytes,
8 * 1024 * 1024,
"the entry grew 8192× after admission and the budget must see it"
);
assert!(cache.get(|k| *k == 7).is_some(), "growth must not lose the entry");
cache.clear();
}
/// The eviction order is LRU and `get` promotes. Proven by ORDER, not by
/// count: the entry read most recently is the one that survives.
#[test]
fn get_promotes_and_eviction_takes_the_least_recently_used() {
static B: Budget = Budget::new("t4", "NORNIR_CACHE_BUDGET_T4", 1 << 30);
let cache: BudgetedCache<u32, Vec<u8>> = BudgetedCache::new("test", &B, 64);
B.set_for_test(3 * 1024);
for i in 0..3u32 {
cache.put(i, 1024, Arc::new(vec![0u8; 1024]), |k| *k == i);
}
assert_eq!(cache.stats().entries, 3);
let _ = cache.get(|k| *k == 0); // 0 becomes MRU; 1 is now LRU
cache.put(9, 1024, Arc::new(vec![0u8; 1024]), |k| *k == 9);
assert!(cache.get(|k| *k == 0).is_some(), "the promoted entry survived");
assert!(cache.get(|k| *k == 1).is_none(), "the LRU entry was evicted");
cache.clear();
}
/// An entry bigger than the entire budget is served but never retained — it
/// must not sit in the cache holding it permanently over budget.
#[test]
fn an_entry_bigger_than_the_budget_is_never_retained() {
static B: Budget = Budget::new("t5", "NORNIR_CACHE_BUDGET_T5", 1 << 30);
let cache: BudgetedCache<u32, Vec<u8>> = BudgetedCache::new("test", &B, 64);
B.set_for_test(1024);
let before = B.refused_too_big();
cache.put(1, 1 << 20, Arc::new(vec![0u8; 16]), |k| *k == 1);
assert_eq!(cache.stats().entries, 0);
assert_eq!(B.refused_too_big(), before + 1);
cache.clear();
}
/// The knob is honoured and a junk value falls back to the default rather
/// than to zero (zero would mean "cache nothing", silently).
#[test]
fn the_knob_parses_and_junk_falls_back_loudly() {
static B: Budget = Budget::new("knob", "NORNIR_CACHE_BUDGET_KNOB_TEST", 7 * 1024);
// SAFETY: single-threaded test process env mutation.
unsafe { std::env::set_var("NORNIR_CACHE_BUDGET_KNOB_TEST", "512m") };
assert_eq!(B.max_bytes(), 512 * 1024 * 1024);
unsafe { std::env::set_var("NORNIR_CACHE_BUDGET_KNOB_TEST", "banana") };
assert_eq!(B.max_bytes(), 7 * 1024);
unsafe { std::env::remove_var("NORNIR_CACHE_BUDGET_KNOB_TEST") };
assert_eq!(B.max_bytes(), 7 * 1024);
}
}