lean_ctx/core/
bm25_cache.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3use std::time::Instant;
4
5use super::bm25_index::BM25Index;
6
7const DEFAULT_TTL_SECS: u64 = 60;
8
9pub struct Bm25CacheEntry {
10    pub root: PathBuf,
11    pub index: Arc<BM25Index>,
12    pub loaded_at: Instant,
13}
14
15impl Bm25CacheEntry {
16    pub fn is_fresh(&self) -> bool {
17        self.loaded_at.elapsed().as_secs() < ttl_secs()
18    }
19}
20
21fn ttl_secs() -> u64 {
22    std::env::var("LEAN_CTX_BM25_CACHE_TTL")
23        .ok()
24        .and_then(|v| v.parse().ok())
25        .unwrap_or(DEFAULT_TTL_SECS)
26}
27
28pub type SharedBm25Cache = std::sync::Arc<std::sync::Mutex<Option<Bm25CacheEntry>>>;
29
30/// Get the BM25 index from cache if available and fresh, otherwise load/build,
31/// cache it, and return. Uses Arc to avoid cloning the entire index.
32pub fn get_or_load(cache: &SharedBm25Cache, root: &Path) -> Arc<BM25Index> {
33    {
34        let guard = cache
35            .lock()
36            .unwrap_or_else(std::sync::PoisonError::into_inner);
37        if let Some(ref entry) = *guard {
38            if entry.root == root && entry.is_fresh() {
39                return Arc::clone(&entry.index);
40            }
41        }
42    }
43
44    let index = Arc::new(BM25Index::load_or_build_fast(root));
45
46    let mut guard = cache
47        .lock()
48        .unwrap_or_else(std::sync::PoisonError::into_inner);
49    *guard = Some(Bm25CacheEntry {
50        root: root.to_path_buf(),
51        index: Arc::clone(&index),
52        loaded_at: Instant::now(),
53    });
54
55    index
56}
57
58/// Get index from cache (fresh or stale), triggering background rebuild if stale.
59/// Returns None only if no cache entry exists at all.
60pub fn get_or_background(cache: &SharedBm25Cache, root: &Path) -> Option<Arc<BM25Index>> {
61    let guard = cache
62        .lock()
63        .unwrap_or_else(std::sync::PoisonError::into_inner);
64    let entry = guard.as_ref()?;
65    if entry.root != root {
66        return None;
67    }
68
69    let idx = Arc::clone(&entry.index);
70
71    if !entry.is_fresh() {
72        let root_str = root.to_string_lossy().to_string();
73        let cache_clone = cache.clone();
74        let root_clone = root.to_path_buf();
75        std::thread::spawn(move || {
76            let rebuilt = BM25Index::load_or_build(&root_clone);
77            let mut g = cache_clone
78                .lock()
79                .unwrap_or_else(std::sync::PoisonError::into_inner);
80            *g = Some(Bm25CacheEntry {
81                root: root_clone,
82                index: Arc::new(rebuilt),
83                loaded_at: Instant::now(),
84            });
85            tracing::debug!("[bm25_cache: background refresh done for {root_str}]");
86        });
87    }
88
89    Some(idx)
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95    use std::sync::Arc;
96
97    #[test]
98    fn fresh_cache_returns_same_instance() {
99        let cache: SharedBm25Cache = Arc::new(std::sync::Mutex::new(None));
100        let tmp = tempfile::tempdir().unwrap();
101        let root = tmp.path();
102        std::fs::write(root.join("main.rs"), "fn main() {}\n").unwrap();
103
104        let idx1 = get_or_load(&cache, root);
105        assert!(idx1.doc_count > 0);
106
107        let idx2 = get_or_load(&cache, root);
108        assert_eq!(idx1.doc_count, idx2.doc_count);
109    }
110
111    #[test]
112    fn different_root_invalidates() {
113        let cache: SharedBm25Cache = Arc::new(std::sync::Mutex::new(None));
114        let tmp1 = tempfile::tempdir().unwrap();
115        let tmp2 = tempfile::tempdir().unwrap();
116        std::fs::write(tmp1.path().join("a.rs"), "fn a() {}\n").unwrap();
117        std::fs::write(tmp2.path().join("b.rs"), "fn b() {}\n").unwrap();
118
119        let _ = get_or_load(&cache, tmp1.path());
120        let idx2 = get_or_load(&cache, tmp2.path());
121
122        let guard = cache.lock().unwrap();
123        let entry = guard.as_ref().unwrap();
124        assert_eq!(entry.root, tmp2.path());
125        assert_eq!(entry.index.doc_count, idx2.doc_count);
126    }
127
128    #[test]
129    fn get_or_background_returns_none_on_empty() {
130        let cache: SharedBm25Cache = Arc::new(std::sync::Mutex::new(None));
131        let tmp = tempfile::tempdir().unwrap();
132        assert!(get_or_background(&cache, tmp.path()).is_none());
133    }
134}
lean_ctx/core/bm25_cache.rs

lean_ctx/core/
bm25_cache.rs
