rullst_core/cache/
memory.rs

1//! When an insertion would exceed a bound, expired entries are dropped first
2//! and then the oldest insertions are evicted until the store is back at three
3//! quarters of both limits, which amortizes eviction work. A memoized call
4//! whose entry is not cached simply runs uncached.
5
6use std::collections::HashMap;
7use std::sync::{Arc, Mutex, OnceLock, PoisonError};
8use std::time::{Duration, Instant};
9
10/// Maximum number of memoized entries kept by the process.
11const MAX_ENTRIES: usize = 4_096;
12/// Maximum key plus value bytes of one memoized entry.
13const MAX_ENTRY_BYTES: usize = 256 * 1024;
14/// Maximum key plus value bytes of all memoized entries.
15const MAX_TOTAL_BYTES: usize = 32 * 1024 * 1024;
16/// Lifetime of a memoized entry.
17const TTL: Duration = Duration::from_secs(3600);
18
19static GLOBAL_MEMO_CACHE: OnceLock<MemoStore> = OnceLock::new();
20
21fn get_cache() -> &'static MemoStore {
22    GLOBAL_MEMO_CACHE.get_or_init(|| {
23        MemoStore::new(Limits {
24            max_entries: MAX_ENTRIES,
25            max_entry_bytes: MAX_ENTRY_BYTES,
26            max_total_bytes: MAX_TOTAL_BYTES,
27            ttl: TTL,
28        })
29    })
30}
31
32/// Retrieve a value from the global memoize cache.
33pub fn get(key: &str) -> Option<String> {
34    get_cache().get(key).map(|value| value.to_string())
35}
36
37/// Store a value in the global memoize cache for one hour.
38///
39/// The entry is skipped when its key plus value exceeds 256 KiB; the oldest
40/// entries are evicted when the store would exceed 4,096 entries or 32 MiB.
41pub fn set(key: &str, value: &str) {
42    get_cache().set(key, value);
43}
44
45#[derive(Clone, Copy)]
46struct Limits {
47    max_entries: usize,
48    max_entry_bytes: usize,
49    max_total_bytes: usize,
50    ttl: Duration,
51}
52
53struct Entry {
54    value: Arc<str>,
55    expires_at: Instant,
56    inserted: u64,
57    bytes: usize,
58}
59
60#[derive(Default)]
61struct State {
62    entries: HashMap<Box<str>, Entry>,
63    total_bytes: usize,
64    next_insertion: u64,
65}
66
67impl State {
68    fn remove(&mut self, key: &str) {
69        if let Some(entry) = self.entries.remove(key) {
70            self.total_bytes = self.total_bytes.saturating_sub(entry.bytes);
71        }
72    }
73
74    fn retain(&mut self, keep: impl Fn(&Entry) -> bool) {
75        self.entries.retain(|_, entry| keep(entry));
76        self.total_bytes = self.entries.values().map(|entry| entry.bytes).sum();
77    }
78
79    fn fits(&self, limits: &Limits, bytes: usize) -> bool {
80        self.entries.len() < limits.max_entries
81            && self.total_bytes.saturating_add(bytes) <= limits.max_total_bytes
82    }
83
84    /// Evicts the oldest insertions until the store and the new entry fit
85    /// within three quarters of both limits.
86    fn evict_oldest(&mut self, limits: &Limits, bytes: usize) {
87        let entry_target = limits.max_entries - limits.max_entries / 4;
88        let byte_target = limits.max_total_bytes - limits.max_total_bytes / 4;
89        let mut order: Vec<(u64, usize)> = self
90            .entries
91            .values()
92            .map(|entry| (entry.inserted, entry.bytes))
93            .collect();
94        order.sort_unstable();
95
96        let mut remaining_entries = self.entries.len();
97        let mut remaining_bytes = self.total_bytes;
98        let mut newest_evicted = None;
99        for (inserted, entry_bytes) in order {
100            if remaining_entries < entry_target
101                && remaining_bytes.saturating_add(bytes) <= byte_target
102            {
103                break;
104            }
105            remaining_entries -= 1;
106            remaining_bytes = remaining_bytes.saturating_sub(entry_bytes);
107            newest_evicted = Some(inserted);
108        }
109        if let Some(newest_evicted) = newest_evicted {
110            self.retain(|entry| entry.inserted > newest_evicted);
111        }
112    }
113}
114
115struct MemoStore {
116    state: Mutex<State>,
117    limits: Limits,
118}
119
120impl MemoStore {
121    fn new(limits: Limits) -> Self {
122        Self {
123            state: Mutex::new(State::default()),
124            limits,
125        }
126    }
127
128    fn state(&self) -> std::sync::MutexGuard<'_, State> {
129        // Every update leaves the map and byte count consistent, so a panic
130        // in another thread cannot leave a state that is unsafe to reuse.
131        self.state.lock().unwrap_or_else(PoisonError::into_inner)
132    }
133
134    fn get(&self, key: &str) -> Option<Arc<str>> {
135        let mut state = self.state();
136        let entry = state.entries.get(key)?;
137        if Instant::now() < entry.expires_at {
138            return Some(Arc::clone(&entry.value));
139        }
140        state.remove(key);
141        None
142    }
143
144    fn set(&self, key: &str, value: &str) {
145        let bytes = key.len().saturating_add(value.len());
146        let expires_at = Instant::now().checked_add(self.limits.ttl);
147        let mut state = self.state();
148        // A replaced or skipped key never keeps its previous value.
149        state.remove(key);
150        let Some(expires_at) = expires_at else {
151            return;
152        };
153        if bytes > self.limits.max_entry_bytes || bytes > self.limits.max_total_bytes {
154            return;
155        }
156        if !state.fits(&self.limits, bytes) {
157            let now = Instant::now();
158            state.retain(|entry| now < entry.expires_at);
159        }
160        if !state.fits(&self.limits, bytes) {
161            state.evict_oldest(&self.limits, bytes);
162        }
163        if !state.fits(&self.limits, bytes) {
164            return;
165        }
166        let inserted = state.next_insertion;
167        state.next_insertion = inserted.wrapping_add(1);
168        state.total_bytes = state.total_bytes.saturating_add(bytes);
169        state.entries.insert(
170            key.into(),
171            Entry {
172                value: value.into(),
173                expires_at,
174                inserted,
175                bytes,
176            },
177        );
178    }
179
180    #[cfg(test)]
181    fn usage(&self) -> (usize, usize) {
182        let state = self.state();
183        (state.entries.len(), state.total_bytes)
184    }
185}
186
187#[cfg(test)]
188#[allow(clippy::unwrap_used, clippy::expect_used)]
189mod tests {
190    use super::*;
191
192    fn store(max_entries: usize, max_total_bytes: usize) -> MemoStore {
193        MemoStore::new(Limits {
194            max_entries,
195            max_entry_bytes: 64,
196            max_total_bytes,
197            ttl: TTL,
198        })
199    }
200
201    #[test]
202    fn entry_count_is_bounded_and_oldest_entries_are_evicted_first() {
203        let store = store(8, 1_024);
204        for index in 0..100 {
205            store.set(&format!("key-{index}"), "value");
206            let (entries, _) = store.usage();
207            assert!(entries <= 8, "{entries} entries after insert {index}");
208        }
209        // The newest insertion is always kept; the oldest were evicted.
210        assert_eq!(store.get("key-99").as_deref(), Some("value"));
211        assert_eq!(store.get("key-0"), None);
212    }
213
214    #[test]
215    fn total_bytes_are_bounded() {
216        let store = store(1_000, 100);
217        for index in 0..50 {
218            store.set(&format!("k{index:02}"), "0123456789");
219            let (_, bytes) = store.usage();
220            assert!(bytes <= 100, "{bytes} bytes after insert {index}");
221        }
222        assert_eq!(store.get("k49").as_deref(), Some("0123456789"));
223    }
224
225    #[test]
226    fn oversized_entries_are_skipped_and_replace_stale_values() {
227        let store = store(8, 1_024);
228        store.set("key", "small");
229        store.set("key", &"x".repeat(100));
230        assert_eq!(store.get("key"), None);
231        store.set(&"k".repeat(100), "v");
232        assert_eq!(store.usage(), (0, 0));
233    }
234
235    #[test]
236    fn replacing_a_key_keeps_the_byte_count_exact() {
237        let store = store(8, 1_024);
238        store.set("key", "first");
239        store.set("key", "second value");
240        assert_eq!(store.usage(), (1, "key".len() + "second value".len()));
241        assert_eq!(store.get("key").as_deref(), Some("second value"));
242    }
243
244    #[test]
245    fn expired_entries_are_not_returned_and_are_reclaimed_before_eviction() {
246        let store = MemoStore::new(Limits {
247            max_entries: 2,
248            max_entry_bytes: 64,
249            max_total_bytes: 1_024,
250            ttl: Duration::ZERO,
251        });
252        store.set("a", "1");
253        store.set("b", "2");
254        assert_eq!(store.usage(), (2, 4));
255        // Both expired entries are reclaimed instead of evicting by age.
256        store.set("c", "3");
257        assert_eq!(store.usage(), (1, 2));
258        assert_eq!(store.get("c"), None);
259        assert_eq!(store.usage(), (0, 0));
260    }
261}
rullst_core/cache/memory.rs

rullst_core/cache/
memory.rs
