Skip to main content

lean_ctx/core/
content_cache.rs

1//! Resident, bounded file-content cache shared across the search-index build and
2//! `ctx_search` (issue #148).
3//!
4//! Before this module the trigram [`search_index`](crate::core::search_index)
5//! build read *every* file in the corpus to extract trigrams and then threw the
6//! content away, after which `ctx_search` read the narrowed candidate files
7//! **again** to run the regex line-by-line — the corpus was read from disk
8//! twice. This cache lets the first reader (whichever it is) populate file
9//! contents once, keyed by absolute path and validated by `(mtime, size)`, and
10//! every subsequent reader reuse them as an in-memory hit.
11//!
12//! Correctness: an entry is only ever served when the file's *current*
13//! `(mtime, size)` exactly matches the stored identity, so any edit (which
14//! changes mtime, and usually size) is a guaranteed miss — results can never go
15//! stale. A miss simply falls back to a disk read.
16//!
17//! Bounds & safety:
18//! - Total resident bytes are capped (`LEAN_CTX_CONTENT_CACHE_MB`, default
19//!   128 MB) with approximate-LRU eviction, so a large corpus cannot grow the
20//!   cache without limit.
21//! - Inserts are skipped while the process is under memory pressure, and the
22//!   eviction orchestrator can [`clear`] the cache on `UnloadIndices` /
23//!   `EmergencyDrop`.
24
25use std::fs::Metadata;
26use std::path::{Path, PathBuf};
27use std::sync::{Arc, Mutex, OnceLock};
28use std::time::UNIX_EPOCH;
29
30use lru::LruCache;
31
32/// Default resident byte budget when `LEAN_CTX_CONTENT_CACHE_MB` is unset.
33const DEFAULT_BUDGET_MB: usize = 128;
34
35/// Identity of one file *version*. A changed mtime or size ⇒ stale ⇒ cache miss.
36/// Mirrors the `(mtime, size)` pair the BM25 index already trusts for staleness.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub struct FileState {
39    pub mtime_ms: u64,
40    pub size_bytes: u64,
41}
42
43impl FileState {
44    /// Build from an already-`stat`ed [`Metadata`] (no extra syscall) — callers
45    /// in the hot path typically have this in hand from their size/regular-file
46    /// checks. Returns `None` only when the platform cannot report mtime.
47    pub fn from_metadata(meta: &Metadata) -> Option<Self> {
48        let mtime_ms = meta
49            .modified()
50            .ok()
51            .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
52            .map(|d| d.as_millis() as u64)?;
53        Some(Self {
54            mtime_ms,
55            size_bytes: meta.len(),
56        })
57    }
58
59    /// Convenience: `stat` the path then build the state. Costs one syscall.
60    pub fn from_path(path: &Path) -> Option<Self> {
61        Self::from_metadata(&path.metadata().ok()?)
62    }
63}
64
65struct Entry {
66    state: FileState,
67    content: Arc<str>,
68}
69
70struct Cache {
71    /// Unbounded by entry count — eviction is driven by `total_bytes` vs
72    /// `budget_bytes` below, not by a capacity `lru::LruCache` would enforce on
73    /// its own. `LruCache` gives O(1) recency tracking and O(1) LRU-victim
74    /// lookup (`pop_lru`), replacing a hand-rolled `last_used` clock plus an
75    /// O(n) `min_by_key` scan per eviction — the scan was rare per insert as
76    /// originally noted, but `trim_oldest_percent` called it in a loop, making
77    /// a big trim O(n²) over the resident set.
78    map: LruCache<PathBuf, Entry>,
79    total_bytes: usize,
80    budget_bytes: usize,
81    hits: u64,
82    misses: u64,
83    inserts: u64,
84    evictions: u64,
85}
86
87impl Cache {
88    fn new(budget_bytes: usize) -> Self {
89        Self {
90            map: LruCache::unbounded(),
91            total_bytes: 0,
92            budget_bytes,
93            hits: 0,
94            misses: 0,
95            inserts: 0,
96            evictions: 0,
97        }
98    }
99
100    fn remove_entry(&mut self, path: &Path) {
101        if let Some(old) = self.map.pop(path) {
102            self.total_bytes = self.total_bytes.saturating_sub(old.content.len());
103        }
104    }
105
106    /// Evict LRU entries until the budget is satisfied. O(1) per eviction.
107    fn evict_to_budget(&mut self) {
108        while self.total_bytes > self.budget_bytes {
109            let Some((_, victim)) = self.map.pop_lru() else {
110                break;
111            };
112            self.total_bytes = self.total_bytes.saturating_sub(victim.content.len());
113            self.evictions += 1;
114        }
115    }
116}
117
118static CACHE: OnceLock<Mutex<Cache>> = OnceLock::new();
119
120fn budget_bytes() -> usize {
121    let mb = std::env::var("LEAN_CTX_CONTENT_CACHE_MB")
122        .ok()
123        .and_then(|v| v.trim().parse::<usize>().ok())
124        .unwrap_or(DEFAULT_BUDGET_MB);
125    mb.saturating_mul(1024 * 1024)
126}
127
128fn disabled() -> bool {
129    // A zero byte budget (or the explicit disable flag) turns the cache into a
130    // no-op pass-through — every `get` misses and `insert` is dropped.
131    std::env::var("LEAN_CTX_DISABLE_CONTENT_CACHE")
132        .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
133        || budget_bytes() == 0
134}
135
136fn cache() -> &'static Mutex<Cache> {
137    CACHE.get_or_init(|| Mutex::new(Cache::new(budget_bytes())))
138}
139
140fn lock() -> std::sync::MutexGuard<'static, Cache> {
141    cache()
142        .lock()
143        .unwrap_or_else(std::sync::PoisonError::into_inner)
144}
145
146/// Look up `path`; returns the cached content only when the supplied current
147/// `(mtime, size)` matches the stored identity. A mismatch evicts the stale
148/// entry and reports a miss. `state` is passed in (not re-`stat`ed) because hot
149/// callers already hold the metadata.
150pub fn get(path: &Path, current: FileState) -> Option<Arc<str>> {
151    if disabled() {
152        return None;
153    }
154    let mut c = lock();
155    // `peek` (not `get`) first: a stale hit must not promote to MRU on its way
156    // to being evicted.
157    let Some(entry) = c.map.peek(path) else {
158        c.misses += 1;
159        crate::core::telemetry::global_metrics().record_cache(false);
160        return None;
161    };
162    if entry.state != current {
163        // Stale version cached — drop it so we don't keep paying for it.
164        c.remove_entry(path);
165        c.misses += 1;
166        crate::core::telemetry::global_metrics().record_cache(false);
167        return None;
168    }
169    c.hits += 1;
170    crate::core::telemetry::global_metrics().record_cache(true);
171    // `get` promotes to MRU; present under the lock we still hold, but degrade
172    // gracefully instead of panicking on the read hot path if that invariant
173    // ever changes.
174    let entry = c.map.get(path)?;
175    Some(Arc::clone(&entry.content))
176}
177
178/// Insert (or replace) the content for `path` at version `state`. Skipped while
179/// the process is under memory pressure or when the cache is disabled, so the
180/// cache never *adds* to a memory problem.
181pub fn insert(path: &Path, state: FileState, content: Arc<str>) {
182    if disabled() || crate::core::memory_guard::is_under_pressure() {
183        return;
184    }
185    let len = content.len();
186    let mut c = lock();
187    // A single file larger than the whole budget would thrash eviction — skip it.
188    if len > c.budget_bytes {
189        return;
190    }
191    c.remove_entry(path);
192    c.map.put(path.to_path_buf(), Entry { state, content });
193    c.total_bytes += len;
194    c.inserts += 1;
195    if c.total_bytes > c.budget_bytes {
196        c.evict_to_budget();
197    }
198}
199
200/// Read a file through the cache: returns cached content on a fresh hit, else
201/// reads from disk (UTF-8), populates the cache, and returns it. `None` on a
202/// non-UTF-8/unreadable/unstatable file. Convenience for callers without their
203/// own size/special-file gating (the search-index build and `ctx_search` use
204/// the explicit [`get`]/[`insert`] pair so they keep their own skip rules).
205pub fn get_or_read(path: &Path) -> Option<Arc<str>> {
206    let state = FileState::from_path(path)?;
207    if let Some(hit) = get(path, state) {
208        return Some(hit);
209    }
210    let content = std::fs::read_to_string(path).ok()?;
211    let arc: Arc<str> = Arc::from(content);
212    insert(path, state, Arc::clone(&arc));
213    Some(arc)
214}
215
216/// Drop all entries, freeing the heap. Called by the eviction orchestrator under
217/// memory pressure; the cache simply re-warms on subsequent reads.
218pub fn clear() {
219    if CACHE.get().is_none() {
220        return;
221    }
222    let mut c = lock();
223    c.map.clear();
224    c.total_bytes = 0;
225}
226
227/// Evict the oldest `percent` of entries by LRU clock. Used after index builds
228/// to release file contents that were only needed for chunking/tokenization.
229/// `percent` is clamped to `[0, 100]`.
230pub fn trim_oldest_percent(percent: u8) {
231    if CACHE.get().is_none() {
232        return;
233    }
234    let mut c = lock();
235    if c.map.is_empty() {
236        return;
237    }
238    let pct = (percent.min(100)) as usize;
239    let target_evictions = c.map.len() * pct / 100;
240    for _ in 0..target_evictions {
241        let Some((_, victim)) = c.map.pop_lru() else {
242            break;
243        };
244        c.total_bytes = c.total_bytes.saturating_sub(victim.content.len());
245        c.evictions += 1;
246    }
247}
248
249/// Approximate resident heap used by cached contents, in bytes.
250pub fn memory_usage_bytes() -> usize {
251    if CACHE.get().is_none() {
252        return 0;
253    }
254    lock().total_bytes
255}
256
257/// Observability snapshot: `(hits, misses, entries, bytes, evictions)`.
258#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
259pub struct CacheStats {
260    pub hits: u64,
261    pub misses: u64,
262    pub entries: usize,
263    pub bytes: usize,
264    pub inserts: u64,
265    pub evictions: u64,
266}
267
268pub fn stats() -> CacheStats {
269    if CACHE.get().is_none() {
270        return CacheStats::default();
271    }
272    let c = lock();
273    CacheStats {
274        hits: c.hits,
275        misses: c.misses,
276        entries: c.map.len(),
277        bytes: c.total_bytes,
278        inserts: c.inserts,
279        evictions: c.evictions,
280    }
281}
282
283#[cfg(test)]
284mod tests {
285    use super::*;
286
287    /// The cache is a process-wide global and tests mutate it (and the budget
288    /// env var). Serialize them so they cannot observe each other's state.
289    static TEST_LOCK: Mutex<()> = Mutex::new(());
290
291    fn fresh_cache(budget_bytes: usize) {
292        crate::test_env::remove_var("LEAN_CTX_CONTENT_CACHE_MB");
293        crate::test_env::remove_var("LEAN_CTX_DISABLE_CONTENT_CACHE");
294        let mut c = lock();
295        *c = Cache::new(budget_bytes);
296    }
297
298    fn write(dir: &Path, name: &str, body: &str) -> PathBuf {
299        let p = dir.join(name);
300        std::fs::write(&p, body).unwrap();
301        p
302    }
303
304    #[test]
305    fn hit_after_insert_with_matching_state() {
306        let _g = TEST_LOCK
307            .lock()
308            .unwrap_or_else(std::sync::PoisonError::into_inner);
309        fresh_cache(1024 * 1024);
310        let dir = tempfile::tempdir().unwrap();
311        let p = write(dir.path(), "a.rs", "fn main() {}\n");
312        let state = FileState::from_path(&p).unwrap();
313        assert!(get(&p, state).is_none(), "cold cache must miss");
314        insert(&p, state, Arc::from("fn main() {}\n"));
315        let got = get(&p, state).expect("warm cache must hit");
316        assert_eq!(&*got, "fn main() {}\n");
317    }
318
319    #[test]
320    fn miss_paths_update_local_and_central_stats() {
321        let _g = TEST_LOCK
322            .lock()
323            .unwrap_or_else(std::sync::PoisonError::into_inner);
324        fresh_cache(1024 * 1024);
325        let dir = tempfile::tempdir().unwrap();
326        let p = write(dir.path(), "cold.rs", "cold\n");
327        let state = FileState::from_path(&p).unwrap();
328        let stale_state = FileState {
329            size_bytes: state.size_bytes + 1,
330            ..state
331        };
332        let local_before = stats();
333        let central = crate::core::telemetry::global_metrics();
334        let misses_before = central
335            .cache_misses
336            .load(std::sync::atomic::Ordering::Relaxed);
337
338        assert!(get(&p, state).is_none());
339        insert(&p, state, Arc::from("cold\n"));
340        assert!(get(&p, stale_state).is_none());
341
342        let local_after = stats();
343        assert!(local_after.misses >= local_before.misses + 2);
344        assert!(
345            central
346                .cache_misses
347                .load(std::sync::atomic::Ordering::Relaxed)
348                >= misses_before + 2
349        );
350    }
351
352    #[test]
353    fn warm_hit_updates_local_and_central_stats() {
354        let _g = TEST_LOCK
355            .lock()
356            .unwrap_or_else(std::sync::PoisonError::into_inner);
357        fresh_cache(1024 * 1024);
358        let dir = tempfile::tempdir().unwrap();
359        let p = write(dir.path(), "warm.rs", "warm\n");
360        let state = FileState::from_path(&p).unwrap();
361        insert(&p, state, Arc::from("warm\n"));
362        let local_before = stats();
363        let central = crate::core::telemetry::global_metrics();
364        let hits_before = central
365            .cache_hits
366            .load(std::sync::atomic::Ordering::Relaxed);
367
368        assert!(get(&p, state).is_some());
369
370        let local_after = stats();
371        assert!(local_after.hits > local_before.hits);
372        assert!(
373            central
374                .cache_hits
375                .load(std::sync::atomic::Ordering::Relaxed)
376                > hits_before
377        );
378    }
379
380    #[test]
381    fn mtime_or_size_change_invalidates() {
382        let _g = TEST_LOCK
383            .lock()
384            .unwrap_or_else(std::sync::PoisonError::into_inner);
385        fresh_cache(1024 * 1024);
386        let dir = tempfile::tempdir().unwrap();
387        let p = write(dir.path(), "a.rs", "v1\n");
388        let s1 = FileState::from_path(&p).unwrap();
389        insert(&p, s1, Arc::from("v1\n"));
390        assert!(get(&p, s1).is_some());
391
392        // Different size ⇒ different state ⇒ miss, and the stale entry is dropped.
393        let s_bigger = FileState {
394            size_bytes: s1.size_bytes + 10,
395            ..s1
396        };
397        assert!(get(&p, s_bigger).is_none(), "size change must miss");
398        assert!(
399            get(&p, s1).is_none(),
400            "stale entry must be evicted on mismatch"
401        );
402
403        // Different mtime ⇒ miss as well.
404        insert(&p, s1, Arc::from("v1\n"));
405        let s_newer = FileState {
406            mtime_ms: s1.mtime_ms + 1,
407            ..s1
408        };
409        assert!(get(&p, s_newer).is_none(), "mtime change must miss");
410    }
411
412    #[test]
413    fn get_or_read_populates_then_serves_from_cache() {
414        let _g = TEST_LOCK
415            .lock()
416            .unwrap_or_else(std::sync::PoisonError::into_inner);
417        fresh_cache(1024 * 1024);
418        let dir = tempfile::tempdir().unwrap();
419        let p = write(dir.path(), "a.rs", "hello world\n");
420
421        let before = stats();
422        let first = get_or_read(&p).unwrap();
423        assert_eq!(&*first, "hello world\n");
424        let after_first = stats();
425        assert_eq!(
426            after_first.inserts,
427            before.inserts + 1,
428            "first read inserts"
429        );
430
431        let second = get_or_read(&p).unwrap();
432        assert_eq!(&*second, "hello world\n");
433        let after_second = stats();
434        assert_eq!(
435            after_second.inserts, after_first.inserts,
436            "second read must NOT re-insert (served from cache)"
437        );
438        assert!(after_second.hits > after_first.hits, "second read is a hit");
439    }
440
441    #[test]
442    fn eviction_keeps_cache_within_budget() {
443        let _g = TEST_LOCK
444            .lock()
445            .unwrap_or_else(std::sync::PoisonError::into_inner);
446        // Budget fits ~2 small files; a third insert must evict the LRU one.
447        fresh_cache(64);
448        let dir = tempfile::tempdir().unwrap();
449        let pa = write(dir.path(), "a", "aaaaaaaaaaaaaaaaaaaaaaaaaaaa"); // 28 bytes
450        let pb = write(dir.path(), "b", "bbbbbbbbbbbbbbbbbbbbbbbbbbbb");
451        let pc = write(dir.path(), "c", "cccccccccccccccccccccccccccc");
452        let sa = FileState::from_path(&pa).unwrap();
453        let sb = FileState::from_path(&pb).unwrap();
454        let sc = FileState::from_path(&pc).unwrap();
455
456        insert(&pa, sa, Arc::from("aaaaaaaaaaaaaaaaaaaaaaaaaaaa"));
457        // Touch a so b becomes the LRU victim.
458        let _ = get(&pa, sa);
459        insert(&pb, sb, Arc::from("bbbbbbbbbbbbbbbbbbbbbbbbbbbb"));
460        let _ = get(&pa, sa);
461        insert(&pc, sc, Arc::from("cccccccccccccccccccccccccccc"));
462
463        let st = stats();
464        assert!(st.bytes <= 64, "cache must respect byte budget: {st:?}");
465        assert!(st.evictions >= 1, "an eviction must have occurred: {st:?}");
466        assert!(get(&pa, sa).is_some(), "recently-used entry must survive");
467    }
468
469    #[test]
470    fn disabled_via_zero_budget_is_passthrough() {
471        let _env_lock = crate::core::data_dir::test_env_lock();
472        let _g = TEST_LOCK
473            .lock()
474            .unwrap_or_else(std::sync::PoisonError::into_inner);
475        fresh_cache(1024 * 1024);
476        crate::test_env::set_var("LEAN_CTX_CONTENT_CACHE_MB", "0");
477        let dir = tempfile::tempdir().unwrap();
478        let p = write(dir.path(), "a.rs", "x\n");
479        let state = FileState::from_path(&p).unwrap();
480        insert(&p, state, Arc::from("x\n"));
481        assert!(get(&p, state).is_none(), "zero-budget cache is a no-op");
482        crate::test_env::remove_var("LEAN_CTX_CONTENT_CACHE_MB");
483    }
484}