Skip to main content

tokenburn_core/utils/
cache.rs

1//! A tiny per-file result cache so a refresh only re-parses what changed.
2//!
3//! Every collector parses one log file into rows. Re-reading a year of
4//! transcripts every few seconds would burn CPU for no reason, so each file's
5//! rows are cached against its *signature* (length + modification time): an
6//! unchanged file is skipped entirely, a growing one is re-parsed. Idle
7//! refreshes therefore cost a directory walk and a few `stat` calls.
8//!
9//! The cache lives for the life of the process (a TUI/GUI/web server keeps
10//! refreshing); a short-lived CLI call never benefits and never pays.
11
12use std::collections::{HashMap, HashSet};
13use std::path::{Path, PathBuf};
14use std::sync::atomic::{AtomicUsize, Ordering};
15use std::sync::{Arc, Mutex};
16use std::time::SystemTime;
17
18/// What identifies "this file, unchanged": size, mtime and a caller-chosen tag
19/// (e.g. the window start for sqlite sources whose query depends on it).
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct Sig {
22    len: u64,
23    modified: Option<SystemTime>,
24    tag: i64,
25}
26
27impl Sig {
28    /// Signature of `path`, or `None` if it cannot be `stat`ed.
29    pub fn of(path: &Path, tag: i64) -> Option<Sig> {
30        let m = std::fs::metadata(path).ok()?;
31        Some(Sig {
32            len: m.len(),
33            modified: m.modified().ok(),
34            tag,
35        })
36    }
37
38    /// Signature of several related files (a sqlite database and its `-wal`):
39    /// sizes add up, the newest mtime wins.
40    pub fn of_all(paths: &[PathBuf], tag: i64) -> Option<Sig> {
41        let mut sig: Option<Sig> = None;
42        for p in paths {
43            let Some(s) = Sig::of(p, tag) else { continue };
44            sig = Some(match sig {
45                None => s,
46                Some(a) => Sig {
47                    len: a.len + s.len,
48                    modified: a.modified.max(s.modified),
49                    tag,
50                },
51            });
52        }
53        sig
54    }
55}
56
57/// Parsed results per file, valid while the file's [`Sig`] is unchanged.
58pub struct FileCache<T> {
59    inner: Mutex<HashMap<PathBuf, (Sig, Arc<T>)>>,
60    hits: AtomicUsize,
61    misses: AtomicUsize,
62}
63
64impl<T> Default for FileCache<T> {
65    fn default() -> Self {
66        Self {
67            inner: Mutex::new(HashMap::new()),
68            hits: AtomicUsize::new(0),
69            misses: AtomicUsize::new(0),
70        }
71    }
72}
73
74impl<T> FileCache<T> {
75    /// The cached result for `path` if its signature still matches `sig`,
76    /// otherwise run `parse` and remember the result.
77    pub fn get_or_parse(&self, path: &Path, sig: Sig, parse: impl FnOnce() -> T) -> Arc<T> {
78        if let Some((s, v)) = self.lock().get(path) {
79            if *s == sig {
80                self.hits.fetch_add(1, Ordering::Relaxed);
81                return Arc::clone(v);
82            }
83        }
84        self.misses.fetch_add(1, Ordering::Relaxed);
85        // Parse outside the lock: other collectors/threads keep running.
86        let value = Arc::new(parse());
87        self.lock()
88            .insert(path.to_path_buf(), (sig, Arc::clone(&value)));
89        value
90    }
91
92    /// Forget files that no longer exist / are no longer in scope.
93    pub fn prune(&self, keep: &HashSet<PathBuf>) {
94        self.lock().retain(|p, _| keep.contains(p));
95    }
96
97    /// Forget cached files below `root` that are not in `keep` (deleted or out of
98    /// scope), leaving every other root's entries alone.
99    pub fn prune_under(&self, root: &Path, keep: &HashSet<PathBuf>) {
100        self.lock()
101            .retain(|p, _| !p.starts_with(root) || keep.contains(p));
102    }
103
104    /// `(hits, misses)` since creation — what tests (and `--bench`) look at.
105    pub fn stats(&self) -> (usize, usize) {
106        (
107            self.hits.load(Ordering::Relaxed),
108            self.misses.load(Ordering::Relaxed),
109        )
110    }
111
112    /// Number of cached files.
113    pub fn len(&self) -> usize {
114        self.lock().len()
115    }
116
117    pub fn is_empty(&self) -> bool {
118        self.len() == 0
119    }
120
121    fn lock(&self) -> std::sync::MutexGuard<'_, HashMap<PathBuf, (Sig, Arc<T>)>> {
122        // A panic while holding the lock only ever leaves a consistent map behind.
123        self.inner.lock().unwrap_or_else(|e| e.into_inner())
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use std::fs;
131
132    #[test]
133    fn unchanged_files_are_served_from_the_cache() {
134        let d = tempfile::tempdir().unwrap();
135        let p = d.path().join("a.jsonl");
136        fs::write(&p, "one").unwrap();
137        let cache: FileCache<String> = FileCache::default();
138        let sig = Sig::of(&p, 0).unwrap();
139
140        let mut parses = 0;
141        for _ in 0..3 {
142            let v = cache.get_or_parse(&p, sig, || {
143                parses += 1;
144                "parsed".to_string()
145            });
146            assert_eq!(*v, "parsed");
147        }
148        assert_eq!(parses, 1, "parsed once, then reused");
149        assert_eq!(cache.stats(), (2, 1));
150    }
151
152    #[test]
153    fn a_changed_file_is_parsed_again() {
154        let d = tempfile::tempdir().unwrap();
155        let p = d.path().join("a.jsonl");
156        fs::write(&p, "one").unwrap();
157        let cache: FileCache<usize> = FileCache::default();
158        let a = cache.get_or_parse(&p, Sig::of(&p, 0).unwrap(), || 1);
159        fs::write(&p, "one two three").unwrap(); // size changes
160        let b = cache.get_or_parse(&p, Sig::of(&p, 0).unwrap(), || 2);
161        assert_eq!((*a, *b), (1, 2));
162        assert_eq!(cache.stats(), (0, 2));
163    }
164
165    #[test]
166    fn the_tag_is_part_of_the_identity() {
167        let d = tempfile::tempdir().unwrap();
168        let p = d.path().join("db");
169        fs::write(&p, "x").unwrap();
170        let cache: FileCache<i64> = FileCache::default();
171        let a = cache.get_or_parse(&p, Sig::of(&p, 100).unwrap(), || 100);
172        let b = cache.get_or_parse(&p, Sig::of(&p, 200).unwrap(), || 200);
173        assert_eq!(
174            (*a, *b),
175            (100, 200),
176            "a different window start must not reuse rows"
177        );
178    }
179
180    #[test]
181    fn prune_forgets_files_that_left_scope() {
182        let d = tempfile::tempdir().unwrap();
183        let (a, b) = (d.path().join("a"), d.path().join("b"));
184        fs::write(&a, "x").unwrap();
185        fs::write(&b, "y").unwrap();
186        let cache: FileCache<u8> = FileCache::default();
187        cache.get_or_parse(&a, Sig::of(&a, 0).unwrap(), || 1);
188        cache.get_or_parse(&b, Sig::of(&b, 0).unwrap(), || 2);
189        assert_eq!(cache.len(), 2);
190        cache.prune(&HashSet::from([a]));
191        assert_eq!(cache.len(), 1);
192        assert!(!cache.is_empty());
193    }
194
195    #[test]
196    fn prune_under_only_touches_its_own_root() {
197        let d = tempfile::tempdir().unwrap();
198        let (r1, r2) = (d.path().join("r1"), d.path().join("r2"));
199        fs::create_dir_all(&r1).unwrap();
200        fs::create_dir_all(&r2).unwrap();
201        let (a, b) = (r1.join("a"), r2.join("b"));
202        fs::write(&a, "x").unwrap();
203        fs::write(&b, "y").unwrap();
204        let cache: FileCache<u8> = FileCache::default();
205        cache.get_or_parse(&a, Sig::of(&a, 0).unwrap(), || 1);
206        cache.get_or_parse(&b, Sig::of(&b, 0).unwrap(), || 2);
207        cache.prune_under(&r1, &HashSet::new()); // r1 has no live files any more
208        assert_eq!(cache.len(), 1, "r2's entry survives");
209    }
210
211    #[test]
212    fn signatures_of_missing_files_are_none_and_groups_combine() {
213        let d = tempfile::tempdir().unwrap();
214        assert!(Sig::of(&d.path().join("nope"), 0).is_none());
215        let (a, b) = (d.path().join("db"), d.path().join("db-wal"));
216        fs::write(&a, "12345").unwrap();
217        assert!(
218            Sig::of_all(&[a.clone(), b.clone()], 0).is_some(),
219            "the -wal may not exist"
220        );
221        let one = Sig::of_all(std::slice::from_ref(&a), 0).unwrap();
222        fs::write(&b, "67").unwrap();
223        let both = Sig::of_all(&[a, b], 0).unwrap();
224        assert_ne!(one, both, "a growing -wal changes the group signature");
225        assert!(Sig::of_all(&[], 0).is_none());
226    }
227
228    #[test]
229    fn a_panic_in_parse_does_not_poison_the_cache() {
230        let d = tempfile::tempdir().unwrap();
231        let p = d.path().join("a");
232        fs::write(&p, "x").unwrap();
233        let cache: Arc<FileCache<u8>> = Arc::new(FileCache::default());
234        let c2 = Arc::clone(&cache);
235        let p2 = p.clone();
236        let _ = std::thread::spawn(move || {
237            c2.get_or_parse(&p2, Sig::of(&p2, 0).unwrap(), || panic!("boom"));
238        })
239        .join();
240        assert_eq!(*cache.get_or_parse(&p, Sig::of(&p, 0).unwrap(), || 7), 7);
241    }
242}