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::collections::HashMap;
26use std::fs::Metadata;
27use std::path::{Path, PathBuf};
28use std::sync::{Arc, Mutex, OnceLock};
29use std::time::UNIX_EPOCH;
30
31/// Default resident byte budget when `LEAN_CTX_CONTENT_CACHE_MB` is unset.
32const DEFAULT_BUDGET_MB: usize = 128;
33
34/// Identity of one file *version*. A changed mtime or size ⇒ stale ⇒ cache miss.
35/// Mirrors the `(mtime, size)` pair the BM25 index already trusts for staleness.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub struct FileState {
38    pub mtime_ms: u64,
39    pub size_bytes: u64,
40}
41
42impl FileState {
43    /// Build from an already-`stat`ed [`Metadata`] (no extra syscall) — callers
44    /// in the hot path typically have this in hand from their size/regular-file
45    /// checks. Returns `None` only when the platform cannot report mtime.
46    pub fn from_metadata(meta: &Metadata) -> Option<Self> {
47        let mtime_ms = meta
48            .modified()
49            .ok()
50            .and_then(|t| t.duration_since(UNIX_EPOCH).ok())
51            .map(|d| d.as_millis() as u64)?;
52        Some(Self {
53            mtime_ms,
54            size_bytes: meta.len(),
55        })
56    }
57
58    /// Convenience: `stat` the path then build the state. Costs one syscall.
59    pub fn from_path(path: &Path) -> Option<Self> {
60        Self::from_metadata(&path.metadata().ok()?)
61    }
62}
63
64struct Entry {
65    state: FileState,
66    content: Arc<str>,
67    /// Logical clock tick of the last hit/insert — drives approximate LRU.
68    last_used: u64,
69}
70
71struct Cache {
72    map: HashMap<PathBuf, Entry>,
73    total_bytes: usize,
74    budget_bytes: usize,
75    clock: u64,
76    hits: u64,
77    misses: u64,
78    inserts: u64,
79    evictions: u64,
80}
81
82impl Cache {
83    fn new(budget_bytes: usize) -> Self {
84        Self {
85            map: HashMap::new(),
86            total_bytes: 0,
87            budget_bytes,
88            clock: 0,
89            hits: 0,
90            misses: 0,
91            inserts: 0,
92            evictions: 0,
93        }
94    }
95
96    fn tick(&mut self) -> u64 {
97        self.clock += 1;
98        self.clock
99    }
100
101    fn remove_entry(&mut self, path: &Path) {
102        if let Some(old) = self.map.remove(path) {
103            self.total_bytes = self.total_bytes.saturating_sub(old.content.len());
104        }
105    }
106
107    /// Evict approximate-LRU entries until the budget is satisfied. Eviction
108    /// only runs after an over-budget insert, so the `O(n)` min-scan is rare and
109    /// dwarfed by the disk reads it prevents.
110    fn evict_to_budget(&mut self) {
111        while self.total_bytes > self.budget_bytes && !self.map.is_empty() {
112            let Some(victim) = self
113                .map
114                .iter()
115                .min_by_key(|(_, e)| e.last_used)
116                .map(|(p, _)| p.clone())
117            else {
118                break;
119            };
120            self.remove_entry(&victim);
121            self.evictions += 1;
122        }
123    }
124}
125
126static CACHE: OnceLock<Mutex<Cache>> = OnceLock::new();
127
128fn budget_bytes() -> usize {
129    let mb = std::env::var("LEAN_CTX_CONTENT_CACHE_MB")
130        .ok()
131        .and_then(|v| v.trim().parse::<usize>().ok())
132        .unwrap_or(DEFAULT_BUDGET_MB);
133    mb.saturating_mul(1024 * 1024)
134}
135
136fn disabled() -> bool {
137    // A zero byte budget (or the explicit disable flag) turns the cache into a
138    // no-op pass-through — every `get` misses and `insert` is dropped.
139    std::env::var("LEAN_CTX_DISABLE_CONTENT_CACHE")
140        .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
141        || budget_bytes() == 0
142}
143
144fn cache() -> &'static Mutex<Cache> {
145    CACHE.get_or_init(|| Mutex::new(Cache::new(budget_bytes())))
146}
147
148fn lock() -> std::sync::MutexGuard<'static, Cache> {
149    cache()
150        .lock()
151        .unwrap_or_else(std::sync::PoisonError::into_inner)
152}
153
154/// Look up `path`; returns the cached content only when the supplied current
155/// `(mtime, size)` matches the stored identity. A mismatch evicts the stale
156/// entry and reports a miss. `state` is passed in (not re-`stat`ed) because hot
157/// callers already hold the metadata.
158pub fn get(path: &Path, current: FileState) -> Option<Arc<str>> {
159    if disabled() {
160        return None;
161    }
162    let mut c = lock();
163    let Some(entry) = c.map.get(path) else {
164        c.misses += 1;
165        return None;
166    };
167    let matches = entry.state == current;
168    if !matches {
169        // Stale version cached — drop it so we don't keep paying for it.
170        c.remove_entry(path);
171        c.misses += 1;
172        return None;
173    }
174    let tick = c.tick();
175    c.hits += 1;
176    // The entry is present under the lock we still hold, but degrade gracefully
177    // instead of panicking on the read hot path if that invariant ever changes.
178    let entry = c.map.get_mut(path)?;
179    entry.last_used = tick;
180    Some(Arc::clone(&entry.content))
181}
182
183/// Insert (or replace) the content for `path` at version `state`. Skipped while
184/// the process is under memory pressure or when the cache is disabled, so the
185/// cache never *adds* to a memory problem.
186pub fn insert(path: &Path, state: FileState, content: Arc<str>) {
187    if disabled() || crate::core::memory_guard::is_under_pressure() {
188        return;
189    }
190    let len = content.len();
191    let mut c = lock();
192    // A single file larger than the whole budget would thrash eviction — skip it.
193    if len > c.budget_bytes {
194        return;
195    }
196    c.remove_entry(path);
197    let tick = c.tick();
198    c.map.insert(
199        path.to_path_buf(),
200        Entry {
201            state,
202            content,
203            last_used: tick,
204        },
205    );
206    c.total_bytes += len;
207    c.inserts += 1;
208    if c.total_bytes > c.budget_bytes {
209        c.evict_to_budget();
210    }
211}
212
213/// Read a file through the cache: returns cached content on a fresh hit, else
214/// reads from disk (UTF-8), populates the cache, and returns it. `None` on a
215/// non-UTF-8/unreadable/unstatable file. Convenience for callers without their
216/// own size/special-file gating (the search-index build and `ctx_search` use
217/// the explicit [`get`]/[`insert`] pair so they keep their own skip rules).
218pub fn get_or_read(path: &Path) -> Option<Arc<str>> {
219    let state = FileState::from_path(path)?;
220    if let Some(hit) = get(path, state) {
221        return Some(hit);
222    }
223    let content = std::fs::read_to_string(path).ok()?;
224    let arc: Arc<str> = Arc::from(content);
225    insert(path, state, Arc::clone(&arc));
226    Some(arc)
227}
228
229/// Drop all entries, freeing the heap. Called by the eviction orchestrator under
230/// memory pressure; the cache simply re-warms on subsequent reads.
231pub fn clear() {
232    if CACHE.get().is_none() {
233        return;
234    }
235    let mut c = lock();
236    c.map.clear();
237    c.total_bytes = 0;
238}
239
240/// Approximate resident heap used by cached contents, in bytes.
241pub fn memory_usage_bytes() -> usize {
242    if CACHE.get().is_none() {
243        return 0;
244    }
245    lock().total_bytes
246}
247
248/// Observability snapshot: `(hits, misses, entries, bytes, evictions)`.
249#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
250pub struct CacheStats {
251    pub hits: u64,
252    pub misses: u64,
253    pub entries: usize,
254    pub bytes: usize,
255    pub inserts: u64,
256    pub evictions: u64,
257}
258
259pub fn stats() -> CacheStats {
260    if CACHE.get().is_none() {
261        return CacheStats::default();
262    }
263    let c = lock();
264    CacheStats {
265        hits: c.hits,
266        misses: c.misses,
267        entries: c.map.len(),
268        bytes: c.total_bytes,
269        inserts: c.inserts,
270        evictions: c.evictions,
271    }
272}
273
274#[cfg(test)]
275mod tests {
276    use super::*;
277
278    /// The cache is a process-wide global and tests mutate it (and the budget
279    /// env var). Serialize them so they cannot observe each other's state.
280    static TEST_LOCK: Mutex<()> = Mutex::new(());
281
282    fn fresh_cache(budget_bytes: usize) {
283        std::env::remove_var("LEAN_CTX_CONTENT_CACHE_MB");
284        std::env::remove_var("LEAN_CTX_DISABLE_CONTENT_CACHE");
285        let mut c = lock();
286        *c = Cache::new(budget_bytes);
287    }
288
289    fn write(dir: &Path, name: &str, body: &str) -> PathBuf {
290        let p = dir.join(name);
291        std::fs::write(&p, body).unwrap();
292        p
293    }
294
295    #[test]
296    fn hit_after_insert_with_matching_state() {
297        let _g = TEST_LOCK
298            .lock()
299            .unwrap_or_else(std::sync::PoisonError::into_inner);
300        fresh_cache(1024 * 1024);
301        let dir = tempfile::tempdir().unwrap();
302        let p = write(dir.path(), "a.rs", "fn main() {}\n");
303        let state = FileState::from_path(&p).unwrap();
304        assert!(get(&p, state).is_none(), "cold cache must miss");
305        insert(&p, state, Arc::from("fn main() {}\n"));
306        let got = get(&p, state).expect("warm cache must hit");
307        assert_eq!(&*got, "fn main() {}\n");
308    }
309
310    #[test]
311    fn mtime_or_size_change_invalidates() {
312        let _g = TEST_LOCK
313            .lock()
314            .unwrap_or_else(std::sync::PoisonError::into_inner);
315        fresh_cache(1024 * 1024);
316        let dir = tempfile::tempdir().unwrap();
317        let p = write(dir.path(), "a.rs", "v1\n");
318        let s1 = FileState::from_path(&p).unwrap();
319        insert(&p, s1, Arc::from("v1\n"));
320        assert!(get(&p, s1).is_some());
321
322        // Different size ⇒ different state ⇒ miss, and the stale entry is dropped.
323        let s_bigger = FileState {
324            size_bytes: s1.size_bytes + 10,
325            ..s1
326        };
327        assert!(get(&p, s_bigger).is_none(), "size change must miss");
328        assert!(
329            get(&p, s1).is_none(),
330            "stale entry must be evicted on mismatch"
331        );
332
333        // Different mtime ⇒ miss as well.
334        insert(&p, s1, Arc::from("v1\n"));
335        let s_newer = FileState {
336            mtime_ms: s1.mtime_ms + 1,
337            ..s1
338        };
339        assert!(get(&p, s_newer).is_none(), "mtime change must miss");
340    }
341
342    #[test]
343    fn get_or_read_populates_then_serves_from_cache() {
344        let _g = TEST_LOCK
345            .lock()
346            .unwrap_or_else(std::sync::PoisonError::into_inner);
347        fresh_cache(1024 * 1024);
348        let dir = tempfile::tempdir().unwrap();
349        let p = write(dir.path(), "a.rs", "hello world\n");
350
351        let before = stats();
352        let first = get_or_read(&p).unwrap();
353        assert_eq!(&*first, "hello world\n");
354        let after_first = stats();
355        assert_eq!(
356            after_first.inserts,
357            before.inserts + 1,
358            "first read inserts"
359        );
360
361        let second = get_or_read(&p).unwrap();
362        assert_eq!(&*second, "hello world\n");
363        let after_second = stats();
364        assert_eq!(
365            after_second.inserts, after_first.inserts,
366            "second read must NOT re-insert (served from cache)"
367        );
368        assert!(after_second.hits > after_first.hits, "second read is a hit");
369    }
370
371    #[test]
372    fn eviction_keeps_cache_within_budget() {
373        let _g = TEST_LOCK
374            .lock()
375            .unwrap_or_else(std::sync::PoisonError::into_inner);
376        // Budget fits ~2 small files; a third insert must evict the LRU one.
377        fresh_cache(64);
378        let dir = tempfile::tempdir().unwrap();
379        let pa = write(dir.path(), "a", "aaaaaaaaaaaaaaaaaaaaaaaaaaaa"); // 28 bytes
380        let pb = write(dir.path(), "b", "bbbbbbbbbbbbbbbbbbbbbbbbbbbb");
381        let pc = write(dir.path(), "c", "cccccccccccccccccccccccccccc");
382        let sa = FileState::from_path(&pa).unwrap();
383        let sb = FileState::from_path(&pb).unwrap();
384        let sc = FileState::from_path(&pc).unwrap();
385
386        insert(&pa, sa, Arc::from("aaaaaaaaaaaaaaaaaaaaaaaaaaaa"));
387        // Touch a so b becomes the LRU victim.
388        let _ = get(&pa, sa);
389        insert(&pb, sb, Arc::from("bbbbbbbbbbbbbbbbbbbbbbbbbbbb"));
390        let _ = get(&pa, sa);
391        insert(&pc, sc, Arc::from("cccccccccccccccccccccccccccc"));
392
393        let st = stats();
394        assert!(st.bytes <= 64, "cache must respect byte budget: {st:?}");
395        assert!(st.evictions >= 1, "an eviction must have occurred: {st:?}");
396        assert!(get(&pa, sa).is_some(), "recently-used entry must survive");
397    }
398
399    #[test]
400    fn disabled_via_zero_budget_is_passthrough() {
401        let _g = TEST_LOCK
402            .lock()
403            .unwrap_or_else(std::sync::PoisonError::into_inner);
404        fresh_cache(1024 * 1024);
405        std::env::set_var("LEAN_CTX_CONTENT_CACHE_MB", "0");
406        let dir = tempfile::tempdir().unwrap();
407        let p = write(dir.path(), "a.rs", "x\n");
408        let state = FileState::from_path(&p).unwrap();
409        insert(&p, state, Arc::from("x\n"));
410        assert!(get(&p, state).is_none(), "zero-budget cache is a no-op");
411        std::env::remove_var("LEAN_CTX_CONTENT_CACHE_MB");
412    }
413}