veil_engine/
scan_cache.rs

1//! Remembers where the matches are in text the engine has scanned before.
2//!
3//! Chat clients resend the whole conversation on every turn, so almost all of
4//! a request was scanned on the previous turn. Scanning depends only on the
5//! rules and the text, not on the scope or the vault, so the result can be
6//! reused for any conversation; placeholder numbers still come from the vault
7//! on every request. Entries are keyed by a SHA-256 of the rules' identity and
8//! the text, and hold only offsets and kinds: no text is kept.
9
10use std::mem::size_of;
11use std::sync::{Arc, Mutex};
12
13use lru::LruCache;
14use sha2::{Digest, Sha256};
15
16use crate::{EntityKind, Match, Source};
17
18/// Shorter text is scanned again: hashing it costs about as much.
19pub(crate) const MIN_LEN: usize = 256;
20
21/// Memory for cached scans, by default.
22pub const DEFAULT_BUDGET: usize = 32 * 1024 * 1024;
23
24/// What scanning one text found.
25#[derive(Debug, PartialEq, Eq)]
26pub(crate) struct Scanned {
27    pub matches: Vec<Match>,
28    /// Placeholder-shaped text the user wrote: (kind, number).
29    pub tokens: Vec<(EntityKind, u64)>,
30}
31
32impl Scanned {
33    /// Roughly what the entry costs in memory, for the budget.
34    fn cost(&self) -> usize {
35        let names: usize = self
36            .matches
37            .iter()
38            .map(|m| match &m.source {
39                Source::Regex(name) => name.len(),
40                _ => 0,
41            })
42            .sum();
43        96 + self.matches.len() * size_of::<Match>() + names + self.tokens.len() * size_of::<(EntityKind, u64)>()
44    }
45}
46
47/// How the cache is doing, for tests and diagnostics.
48#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
49#[non_exhaustive]
50pub struct ScanCacheStats {
51    /// Scans held.
52    pub entries: usize,
53    /// Roughly how much memory they use, in bytes.
54    pub bytes: usize,
55    /// Scans answered from the cache.
56    pub hits: u64,
57    /// Scans that had to run in full.
58    pub misses: u64,
59}
60
61pub(crate) struct ScanCache {
62    budget: usize,
63    inner: Mutex<Inner>,
64}
65
66struct Inner {
67    map: LruCache<[u8; 32], (Arc<Scanned>, usize)>,
68    bytes: usize,
69    hits: u64,
70    misses: u64,
71}
72
73impl ScanCache {
74    /// A cache using up to `budget` bytes; 0 turns it off.
75    pub fn new(budget: usize) -> Self {
76        ScanCache { budget, inner: Mutex::new(Inner { map: LruCache::unbounded(), bytes: 0, hits: 0, misses: 0 }) }
77    }
78
79    pub fn enabled(&self) -> bool {
80        self.budget > 0
81    }
82
83    pub fn key(rules_id: u64, text: &str) -> [u8; 32] {
84        let mut h = Sha256::new();
85        h.update(rules_id.to_le_bytes());
86        h.update(text.as_bytes());
87        h.finalize().into()
88    }
89
90    pub fn get(&self, key: &[u8; 32]) -> Option<Arc<Scanned>> {
91        let mut g = self.inner.lock().unwrap_or_else(|e| e.into_inner());
92        let found = g.map.get(key).map(|(s, _)| s.clone());
93        match found {
94            Some(_) => g.hits += 1,
95            None => g.misses += 1,
96        }
97        found
98    }
99
100    pub fn put(&self, key: [u8; 32], scanned: Arc<Scanned>) {
101        let cost = scanned.cost();
102        if cost > self.budget {
103            return;
104        }
105        let mut g = self.inner.lock().unwrap_or_else(|e| e.into_inner());
106        if let Some((_, old)) = g.map.put(key, (scanned, cost)) {
107            g.bytes -= old;
108        }
109        g.bytes += cost;
110        while g.bytes > self.budget {
111            match g.map.pop_lru() {
112                Some((_, (_, c))) => g.bytes -= c,
113                None => break,
114            }
115        }
116    }
117
118    pub fn clear(&self) {
119        let mut g = self.inner.lock().unwrap_or_else(|e| e.into_inner());
120        g.map.clear();
121        g.bytes = 0;
122    }
123
124    pub fn stats(&self) -> ScanCacheStats {
125        let g = self.inner.lock().unwrap_or_else(|e| e.into_inner());
126        ScanCacheStats { entries: g.map.len(), bytes: g.bytes, hits: g.hits, misses: g.misses }
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use crate::kinds;
134
135    fn scanned(n: usize) -> Arc<Scanned> {
136        let kind = EntityKind::new(kinds::EMAIL).unwrap();
137        let matches =
138            (0..n).map(|i| Match { start: i, end: i + 1, kind: kind.clone(), source: Source::BuiltIn }).collect();
139        Arc::new(Scanned { matches, tokens: Vec::new() })
140    }
141
142    #[test]
143    fn keys_depend_on_the_rules_and_the_text() {
144        assert_eq!(ScanCache::key(1, "abc"), ScanCache::key(1, "abc"));
145        assert_ne!(ScanCache::key(1, "abc"), ScanCache::key(2, "abc"));
146        assert_ne!(ScanCache::key(1, "abc"), ScanCache::key(1, "abd"));
147    }
148
149    #[test]
150    fn counts_hits_and_misses() {
151        let c = ScanCache::new(DEFAULT_BUDGET);
152        let k = ScanCache::key(1, "x");
153        assert!(c.get(&k).is_none());
154        c.put(k, scanned(2));
155        assert_eq!(c.get(&k).unwrap().matches.len(), 2);
156        let s = c.stats();
157        assert_eq!((s.entries, s.hits, s.misses), (1, 1, 1));
158        assert_eq!(s.bytes, scanned(2).cost());
159    }
160
161    #[test]
162    fn stays_within_the_budget_and_drops_the_least_recently_used() {
163        let one = scanned(10).cost();
164        let c = ScanCache::new(one * 3);
165        let keys: Vec<_> = (0..3).map(|i| ScanCache::key(1, &i.to_string())).collect();
166        for k in &keys {
167            c.put(*k, scanned(10));
168        }
169        // Use the first, so the second is the least recently used.
170        assert!(c.get(&keys[0]).is_some());
171        c.put(ScanCache::key(1, "new"), scanned(10));
172        assert!(c.stats().bytes <= one * 3);
173        assert!(c.get(&keys[1]).is_none(), "least recently used goes first");
174        assert!(c.get(&keys[0]).is_some());
175        assert!(c.get(&keys[2]).is_some());
176    }
177
178    #[test]
179    fn an_entry_larger_than_the_budget_is_not_kept() {
180        let c = ScanCache::new(scanned(1).cost());
181        c.put(ScanCache::key(1, "big"), scanned(1000));
182        assert_eq!(c.stats().entries, 0);
183    }
184
185    #[test]
186    fn replacing_an_entry_keeps_the_byte_count_right() {
187        let c = ScanCache::new(DEFAULT_BUDGET);
188        let k = ScanCache::key(1, "x");
189        c.put(k, scanned(5));
190        c.put(k, scanned(1));
191        assert_eq!(c.stats().bytes, scanned(1).cost());
192        c.clear();
193        assert_eq!(c.stats(), ScanCacheStats { entries: 0, bytes: 0, hits: 0, misses: 0 });
194    }
195}
veil_engine/scan_cache.rs

veil_engine/
scan_cache.rs
