lean_ctx/core/
bm25_cache.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::time::{Instant, SystemTime};
4
5use super::bm25_index::BM25Index;
6
7const DEFAULT_TTL_SECS: u64 = 60;
8
9/// Cheap content fingerprint of the persisted index file: `(mtime, size)`.
10///
11/// mtime alone is not enough — many filesystems only resolve mtime to 1–2 s, so
12/// a background rebuild that lands in the same tick as the load would be missed.
13/// Pairing it with the file size catches those same-second rewrites without the
14/// cost of hashing a multi-MB index file on every per-query freshness check.
15#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)]
16pub struct IndexFingerprint {
17    mtime: Option<SystemTime>,
18    size: u64,
19}
20
21pub struct Bm25CacheEntry {
22    pub root: PathBuf,
23    pub index: Arc<BM25Index>,
24    pub loaded_at: Instant,
25    /// Fingerprint of the persisted index file when this entry was loaded.
26    pub fingerprint: IndexFingerprint,
27}
28
29impl Bm25CacheEntry {
30    pub fn is_fresh(&self) -> bool {
31        if self.loaded_at.elapsed().as_secs() >= ttl_secs() {
32            return false;
33        }
34        // Precise invalidation: if a background rebuild changed the index file
35        // on disk, the resident copy is stale even within the TTL window.
36        index_fingerprint(&self.root) == self.fingerprint
37    }
38}
39
40/// `(mtime, size)` fingerprint of the persisted BM25 index file for `root`.
41pub(crate) fn index_fingerprint(root: &Path) -> IndexFingerprint {
42    match std::fs::metadata(BM25Index::index_file_path(root)) {
43        Ok(m) => IndexFingerprint {
44            mtime: m.modified().ok(),
45            size: m.len(),
46        },
47        Err(_) => IndexFingerprint::default(),
48    }
49}
50
51fn ttl_secs() -> u64 {
52    std::env::var("LEAN_CTX_BM25_CACHE_TTL")
53        .ok()
54        .and_then(|v| v.parse().ok())
55        .unwrap_or(DEFAULT_TTL_SECS)
56}
57
58pub type SharedBm25Cache = std::sync::Arc<std::sync::Mutex<Option<Bm25CacheEntry>>>;
59
60/// Get the BM25 index from cache if available and fresh, otherwise load/build,
61/// cache it, and return. Uses Arc to avoid cloning the entire index.
62pub fn get_or_load(cache: &SharedBm25Cache, root: &Path) -> Arc<BM25Index> {
63    {
64        let guard = cache
65            .lock()
66            .unwrap_or_else(std::sync::PoisonError::into_inner);
67        if let Some(ref entry) = *guard
68            && entry.root == root
69            && entry.is_fresh()
70        {
71            return Arc::clone(&entry.index);
72        }
73    }
74
75    let index = Arc::new(BM25Index::load_or_build_fast(root));
76
77    let mut guard = cache
78        .lock()
79        .unwrap_or_else(std::sync::PoisonError::into_inner);
80    *guard = Some(Bm25CacheEntry {
81        root: root.to_path_buf(),
82        index: Arc::clone(&index),
83        loaded_at: Instant::now(),
84        fingerprint: index_fingerprint(root),
85    });
86
87    index
88}
89
90/// Get index from cache (fresh or stale), triggering background rebuild if stale.
91/// Returns None only if no cache entry exists at all.
92pub fn get_or_background(cache: &SharedBm25Cache, root: &Path) -> Option<Arc<BM25Index>> {
93    let guard = cache
94        .lock()
95        .unwrap_or_else(std::sync::PoisonError::into_inner);
96    let entry = guard.as_ref()?;
97    if entry.root != root {
98        return None;
99    }
100
101    let idx = Arc::clone(&entry.index);
102
103    if !entry.is_fresh() {
104        let root_str = root.to_string_lossy().to_string();
105        let cache_clone = cache.clone();
106        let root_clone = root.to_path_buf();
107        std::thread::spawn(move || {
108            // Isolate panics (corrupt index file, FS race): a panic here must not
109            // kill the worker silently — the stale index keeps serving and the
110            // next call retries the refresh.
111            let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
112                let rebuilt = BM25Index::load_or_build(&root_clone);
113                let rebuilt_fp = index_fingerprint(&root_clone);
114                let mut g = cache_clone
115                    .lock()
116                    .unwrap_or_else(std::sync::PoisonError::into_inner);
117                *g = Some(Bm25CacheEntry {
118                    root: root_clone,
119                    index: Arc::new(rebuilt),
120                    loaded_at: Instant::now(),
121                    fingerprint: rebuilt_fp,
122                });
123            }));
124            if result.is_ok() {
125                tracing::debug!("[bm25_cache: background refresh done for {root_str}]");
126            } else {
127                tracing::warn!(
128                    "[bm25_cache: background refresh panicked for {root_str}; serving stale index]"
129                );
130            }
131        });
132    }
133
134    Some(idx)
135}
136
137/// Drops the cached BM25 index, freeing its heap memory.
138/// The index will be rebuilt from disk on the next search.
139pub fn unload(cache: &SharedBm25Cache) {
140    let mut guard = cache
141        .lock()
142        .unwrap_or_else(std::sync::PoisonError::into_inner);
143    if guard.is_some() {
144        *guard = None;
145        tracing::info!("[bm25_cache] unloaded index to free memory");
146    }
147}
148
149/// Returns the approximate heap memory used by the cached BM25 index, or 0.
150pub fn memory_usage(cache: &SharedBm25Cache) -> usize {
151    let guard = cache
152        .lock()
153        .unwrap_or_else(std::sync::PoisonError::into_inner);
154    guard.as_ref().map_or(0, |e| e.index.memory_usage_bytes())
155}
156
157/// Trims the RESIDENT cached index for `root` so each chunk keeps only its first
158/// `keep_lines` lines of `content`, reclaiming the RAM held by full source
159/// bodies once the embedding pass has consumed them.
160///
161/// Call this ONLY after embeddings for the current fingerprint are built and
162/// persisted (see `ctx_semantic_search::ensure_embeddings`). The on-disk index
163/// is untouched, so a reload restores full bodies; the resident `content_truncated`
164/// flag guards a later embedding pass against re-embedding the trimmed bodies.
165///
166/// Truncation happens in place via `Arc::get_mut`, so it is a no-op (and costs
167/// nothing) whenever another owner still holds the Arc — e.g. the search handler
168/// has not yet dropped its clone, or a background refresh is in flight. The next
169/// search call retries against the then-sole-owner cache entry. Returns the bytes
170/// reclaimed (0 if skipped).
171pub fn shrink_resident_to_snippet(
172    cache: &SharedBm25Cache,
173    root: &Path,
174    keep_lines: usize,
175) -> usize {
176    let mut guard = cache
177        .lock()
178        .unwrap_or_else(std::sync::PoisonError::into_inner);
179    let Some(entry) = guard.as_mut() else {
180        return 0;
181    };
182    if entry.root != root || entry.index.content_truncated {
183        return 0;
184    }
185    // Only mutate when the cache is the sole owner — cloning a multi-MB index
186    // just to trim it would defeat the purpose.
187    let Some(index) = Arc::get_mut(&mut entry.index) else {
188        tracing::debug!(
189            "[bm25_cache] resident index still shared; skipping content shrink for now"
190        );
191        return 0;
192    };
193    let before = index.memory_usage_bytes();
194    index.shrink_resident_content_to_snippet(keep_lines);
195    before.saturating_sub(index.memory_usage_bytes())
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    use std::sync::Arc;
202
203    #[test]
204    fn fresh_cache_returns_same_instance() {
205        let cache: SharedBm25Cache = Arc::new(std::sync::Mutex::new(None));
206        let tmp = tempfile::tempdir().unwrap();
207        let root = tmp.path();
208        std::fs::write(root.join("main.rs"), "fn main() {}\n").unwrap();
209
210        let idx1 = get_or_load(&cache, root);
211        assert!(idx1.doc_count > 0);
212
213        let idx2 = get_or_load(&cache, root);
214        assert_eq!(idx1.doc_count, idx2.doc_count);
215    }
216
217    #[test]
218    fn different_root_invalidates() {
219        let cache: SharedBm25Cache = Arc::new(std::sync::Mutex::new(None));
220        let tmp1 = tempfile::tempdir().unwrap();
221        let tmp2 = tempfile::tempdir().unwrap();
222        std::fs::write(tmp1.path().join("a.rs"), "fn a() {}\n").unwrap();
223        std::fs::write(tmp2.path().join("b.rs"), "fn b() {}\n").unwrap();
224
225        let _ = get_or_load(&cache, tmp1.path());
226        let idx2 = get_or_load(&cache, tmp2.path());
227
228        let guard = cache.lock().unwrap();
229        let entry = guard.as_ref().unwrap();
230        assert_eq!(entry.root, tmp2.path());
231        assert_eq!(entry.index.doc_count, idx2.doc_count);
232    }
233
234    #[test]
235    fn get_or_background_returns_none_on_empty() {
236        let cache: SharedBm25Cache = Arc::new(std::sync::Mutex::new(None));
237        let tmp = tempfile::tempdir().unwrap();
238        assert!(get_or_background(&cache, tmp.path()).is_none());
239    }
240
241    #[test]
242    fn fingerprint_default_when_index_file_absent() {
243        let tmp = tempfile::tempdir().unwrap();
244        // No persisted index file → default (None, 0) fingerprint.
245        assert_eq!(index_fingerprint(tmp.path()), IndexFingerprint::default());
246    }
247
248    #[test]
249    fn fingerprint_detects_size_change_under_equal_mtime() {
250        // Two fingerprints with the same mtime but different size must differ,
251        // proving size catches same-second rewrites that mtime alone misses.
252        let mtime = Some(SystemTime::UNIX_EPOCH);
253        let a = IndexFingerprint { mtime, size: 100 };
254        let b = IndexFingerprint { mtime, size: 200 };
255        assert_ne!(a, b);
256        assert_eq!(a, IndexFingerprint { mtime, size: 100 });
257    }
258}
lean_ctx/core/bm25_cache.rs

lean_ctx/core/
bm25_cache.rs
