Line data Source code
1 : // SPDX-License-Identifier: EUPL-1.2
2 : //! wasm32 only: a bounded map with moka's call surface. moka's clock
3 : //! calls `std::time::Instant::now()` unconditionally, which panics on
4 : //! wasm32-unknown-unknown — so the browser build swaps in this FIFO-evicting
5 : //! cache behind the same methods and builder the loader uses. Bounded is the
6 : //! requirement, by the same budget as moka's: an entry count, or with a
7 : //! weigher the sum of the entries' weights (the loader's byte and term
8 : //! budgets). Recency-optimal eviction is not a requirement.
9 : // FIFO eviction is deliberate; port moka's TinyLFU here only if a browser
10 : // workload ever shows cache thrash.
11 :
12 : use std::collections::{HashMap, VecDeque};
13 : use std::hash::Hash;
14 : use std::sync::Mutex;
15 :
16 : type Weigher<K, V> = Box<dyn Fn(&K, &V) -> u32 + Send + Sync>;
17 :
18 : /// Eviction policy, in full: entries leave in the order they were FIRST
19 : /// inserted, until the new entry fits the budget. `get` is a pure read — it
20 : /// does not refresh recency — and re-inserting a resident key overwrites its
21 : /// value without moving it back in the queue. The held weight (one per entry
22 : /// without a weigher) never exceeds the capacity the cache was built with,
23 : /// including a capacity of zero; an entry heavier than the whole budget is
24 : /// not held.
25 : pub struct Cache<K, V> {
26 : cap: u64,
27 : weigher: Option<Weigher<K, V>>,
28 : inner: Mutex<Inner<K, V>>,
29 : }
30 :
31 : /// The entries with their weights, the insertion queue and the weight held.
32 : type Inner<K, V> = (HashMap<K, (V, u64)>, VecDeque<K>, u64);
33 :
34 : /// moka's `Cache::builder()` surface: a capacity and an optional weigher.
35 : pub struct CacheBuilder<K, V> {
36 : cap: u64,
37 : weigher: Option<Weigher<K, V>>,
38 : }
39 :
40 : impl<K: Eq + Hash + Clone, V: Clone> CacheBuilder<K, V> {
41 34 : pub fn max_capacity(mut self, cap: u64) -> Self {
42 34 : self.cap = cap;
43 34 : self
44 34 : }
45 :
46 6 : pub fn weigher(mut self, weigh: impl Fn(&K, &V) -> u32 + Send + Sync + 'static) -> Self {
47 6 : self.weigher = Some(Box::new(weigh));
48 6 : self
49 6 : }
50 :
51 34 : pub fn build(self) -> Cache<K, V> {
52 34 : Cache {
53 34 : cap: self.cap,
54 34 : weigher: self.weigher,
55 34 : inner: Mutex::new((HashMap::new(), VecDeque::new(), 0)),
56 34 : }
57 34 : }
58 : }
59 :
60 : impl<K: Eq + Hash + Clone, V: Clone> Cache<K, V> {
61 28 : pub fn new(cap: u64) -> Self {
62 28 : Self::builder().max_capacity(cap).build()
63 28 : }
64 :
65 34 : pub fn builder() -> CacheBuilder<K, V> {
66 34 : CacheBuilder {
67 34 : cap: 0,
68 34 : weigher: None,
69 34 : }
70 34 : }
71 :
72 : /// A cache must not turn a panic elsewhere into a permanently poisoned
73 : /// broker: nothing under this lock can leave the map inconsistent, so a
74 : /// poisoned guard is taken as-is.
75 36482 : fn lock(&self) -> std::sync::MutexGuard<'_, Inner<K, V>> {
76 36482 : self.inner.lock().unwrap_or_else(|e| e.into_inner())
77 36482 : }
78 :
79 8080 : pub fn get<Q>(&self, key: &Q) -> Option<V>
80 8080 : where
81 8080 : K: std::borrow::Borrow<Q>,
82 8080 : Q: Eq + Hash + ?Sized,
83 : {
84 8080 : self.lock().0.get(key).map(|(v, _)| v.clone())
85 8080 : }
86 :
87 18274 : pub fn insert(&self, key: K, value: V) {
88 18274 : let weight = self
89 18274 : .weigher
90 18274 : .as_ref()
91 18274 : .map_or(1, |w| u64::from(w(&key, &value)));
92 18274 : let mut g = self.lock();
93 18274 : let (map, order, held) = &mut *g;
94 18274 : if let Some((_, old)) = map.remove(&key) {
95 6 : *held = held.saturating_sub(old);
96 18268 : }
97 18274 : if weight > self.cap {
98 204 : order.retain(|k| k != &key);
99 204 : return;
100 18070 : }
101 35817 : while held.saturating_add(weight) > self.cap {
102 17747 : let Some(at) = order.iter().position(|k| k != &key) else {
103 0 : break;
104 : };
105 17747 : if let Some(old) = order.remove(at) {
106 17747 : if let Some((_, w)) = map.remove(&old) {
107 17747 : *held = held.saturating_sub(w);
108 17747 : }
109 0 : }
110 : }
111 18070 : if !order.contains(&key) {
112 18066 : order.push_back(key.clone());
113 18066 : }
114 18070 : *held = held.saturating_add(weight);
115 18070 : map.insert(key, (value, weight));
116 18274 : }
117 :
118 94 : pub fn invalidate<Q>(&self, key: &Q)
119 94 : where
120 94 : K: std::borrow::Borrow<Q>,
121 94 : Q: Eq + Hash + ?Sized,
122 : {
123 94 : let mut g = self.lock();
124 94 : let (map, order, held) = &mut *g;
125 94 : if let Some((_, w)) = map.remove(key) {
126 87 : *held = held.saturating_sub(w);
127 87 : }
128 1360 : order.retain(|k| map.contains_key::<K>(k));
129 94 : }
130 :
131 : /// Snapshot of the resident entries, in moka's `iter()` item shape — the
132 : /// lock is released before the caller walks them, so a callback may touch
133 : /// this cache again.
134 4 : pub fn iter(&self) -> std::vec::IntoIter<(std::sync::Arc<K>, V)> {
135 4 : let g = self.lock();
136 4 : g.0.iter()
137 14 : .map(|(k, (v, _))| (std::sync::Arc::new(k.clone()), v.clone()))
138 4 : .collect::<Vec<_>>()
139 4 : .into_iter()
140 4 : }
141 :
142 2 : pub fn run_pending_tasks(&self) {}
143 :
144 8036 : pub fn entry_count(&self) -> u64 {
145 8036 : self.lock().0.len() as u64
146 8036 : }
147 :
148 : /// The weight held, moka's `weighted_size()`.
149 : #[cfg(test)]
150 1994 : pub fn weighted_size(&self) -> u64 {
151 1994 : self.lock().2
152 1994 : }
153 : }
154 :
155 : #[cfg(test)]
156 : mod tests {
157 : use super::*;
158 :
159 24 : fn filled(cap: u64, n: u32) -> Cache<u32, u32> {
160 24 : let c = Cache::new(cap);
161 8258 : for i in 0..n {
162 8258 : c.insert(i, i);
163 8258 : }
164 24 : c
165 24 : }
166 :
167 : /// The entry count is a hard ceiling: inserting far past the capacity
168 : /// leaves exactly `cap` entries.
169 : #[test]
170 2 : fn capacity_is_a_hard_bound() {
171 8 : for cap in [1u64, 2, 7, 64] {
172 8 : let c = filled(cap, 1000);
173 8 : assert_eq!(c.entry_count(), cap, "cap {cap}");
174 : }
175 2 : }
176 :
177 : /// A zero capacity means "hold nothing" — it must not degrade into an
178 : /// unbounded (or one-entry) map.
179 : #[test]
180 2 : fn zero_capacity_holds_nothing() {
181 2 : let c = filled(0, 100);
182 2 : assert_eq!(c.entry_count(), 0);
183 2 : assert_eq!(c.get(&99), None);
184 2 : }
185 :
186 : /// Eviction is FIFO by first insertion and therefore deterministic: after
187 : /// filling past the bound exactly the last `cap` keys survive.
188 : #[test]
189 2 : fn eviction_is_deterministic_fifo() {
190 2 : let c = filled(4, 6);
191 2 : assert_eq!(c.entry_count(), 4);
192 4 : for gone in [0u32, 1] {
193 4 : assert_eq!(c.get(&gone), None, "key {gone} should have been evicted");
194 : }
195 8 : for kept in [2u32, 3, 4, 5] {
196 8 : assert_eq!(c.get(&kept), Some(kept), "key {kept} should be resident");
197 : }
198 2 : }
199 :
200 : /// Documented policy: `get` is a pure read — it does NOT refresh recency.
201 : /// A key read on every lookup is still evicted in insertion order.
202 : #[test]
203 2 : fn get_does_not_refresh_recency() {
204 2 : let c = filled(4, 4);
205 2 : for _ in 0..10 {
206 20 : assert_eq!(c.get(&0), Some(0));
207 : }
208 2 : c.insert(4, 4);
209 2 : assert_eq!(
210 2 : c.get(&0),
211 : None,
212 : "get must not protect an entry from FIFO eviction"
213 : );
214 2 : assert_eq!(c.get(&1), Some(1));
215 2 : }
216 :
217 : /// Re-inserting an existing key overwrites the value, keeps the count flat
218 : /// and does NOT move the key to the back of the eviction queue.
219 : #[test]
220 2 : fn reinsert_updates_value_without_reordering() {
221 2 : let c = filled(3, 3);
222 2 : c.insert(0, 99);
223 2 : assert_eq!(c.entry_count(), 3);
224 2 : assert_eq!(c.get(&0), Some(99));
225 2 : c.insert(3, 3);
226 2 : assert_eq!(c.get(&0), None, "re-insert must not refresh eviction order");
227 2 : assert_eq!(c.entry_count(), 3);
228 2 : }
229 :
230 : /// `invalidate` drops the entry and its queue slot, so the freed room is
231 : /// reused before anything else is evicted.
232 : #[test]
233 2 : fn invalidate_frees_a_slot() {
234 2 : let c = filled(2, 2);
235 2 : c.invalidate(&0);
236 2 : assert_eq!(c.entry_count(), 1);
237 2 : c.insert(2, 2);
238 2 : assert_eq!(c.entry_count(), 2);
239 2 : assert_eq!(
240 2 : c.get(&1),
241 : Some(1),
242 : "the surviving entry must not be evicted early"
243 : );
244 2 : c.insert(3, 3);
245 2 : assert_eq!(c.entry_count(), 2);
246 2 : assert_eq!(c.get(&1), None);
247 2 : assert_eq!(c.get(&3), Some(3));
248 2 : }
249 :
250 : /// Invalidating an absent key is a no-op, and the bound still holds after.
251 : #[test]
252 2 : fn invalidate_absent_key_is_a_noop() {
253 2 : let c = filled(3, 3);
254 2 : c.invalidate(&42);
255 2 : assert_eq!(c.entry_count(), 3);
256 6 : for i in 0..3 {
257 6 : assert_eq!(c.get(&i), Some(i));
258 : }
259 2 : }
260 :
261 : /// `iter` yields every resident entry exactly once — it is what the
262 : /// caller filters to invalidate a subset, so a missed entry would leave a
263 : /// stale one behind. Evicted keys must not appear.
264 : #[test]
265 2 : fn iter_yields_every_resident_entry() {
266 2 : let c = filled(4, 6);
267 2 : let mut seen: Vec<u32> = c
268 2 : .iter()
269 8 : .map(|(k, v)| {
270 8 : assert_eq!(*k, v, "key and value must belong to the same entry");
271 8 : *k
272 8 : })
273 2 : .collect();
274 2 : seen.sort_unstable();
275 2 : assert_eq!(
276 : seen,
277 2 : vec![2, 3, 4, 5],
278 : "the evicted keys must not be listed"
279 : );
280 2 : c.invalidate(&3);
281 6 : let mut seen: Vec<u32> = c.iter().map(|(k, _)| *k).collect();
282 2 : seen.sort_unstable();
283 2 : assert_eq!(seen, vec![2, 4, 5]);
284 2 : }
285 :
286 6 : fn weighed(cap: u64) -> Cache<u32, u32> {
287 6 : Cache::builder()
288 6 : .max_capacity(cap)
289 6 : .weigher(|_k: &u32, v: &u32| *v)
290 6 : .build()
291 6 : }
292 :
293 : /// With a weigher the budget is the sum of the weights, as moka's is:
294 : /// a heavy entry evicts as many of the oldest as it needs, never more,
295 : /// and the held weight never passes the capacity.
296 : #[test]
297 2 : fn a_weigher_bounds_the_held_weight() {
298 2 : let c = weighed(10);
299 6 : for (k, w) in [(0, 3), (1, 3), (2, 3)] {
300 6 : c.insert(k, w);
301 6 : }
302 2 : assert_eq!((c.entry_count(), c.weighted_size()), (3, 9));
303 2 : c.insert(3, 5);
304 2 : assert_eq!(c.get(&0), None, "the oldest made room");
305 2 : assert_eq!(c.get(&1), None, "and the next, since 3 + 3 + 5 > 10");
306 2 : assert_eq!((c.get(&2), c.get(&3)), (Some(3), Some(5)));
307 2 : assert_eq!(c.weighted_size(), 8);
308 1980 : for k in 10..1000 {
309 1980 : c.insert(k, k % 7 + 1);
310 1980 : assert!(c.weighted_size() <= 10, "bound breached at {k}");
311 : }
312 2 : }
313 :
314 : /// An entry heavier than the whole budget is not held, and it evicts
315 : /// nothing: a hostile 5 MiB document cannot flush the warm ones.
316 : #[test]
317 2 : fn an_entry_over_the_budget_is_not_held() {
318 2 : let c = weighed(10);
319 2 : c.insert(0, 4);
320 2 : c.insert(1, 11);
321 2 : assert_eq!(c.get(&1), None);
322 2 : assert_eq!((c.get(&0), c.weighted_size()), (Some(4), 4));
323 : // a resident key re-inserted over the budget leaves too
324 2 : c.insert(0, 11);
325 2 : assert_eq!(
326 2 : (c.get(&0), c.entry_count(), c.weighted_size()),
327 : (None, 0, 0)
328 : );
329 2 : c.insert(2, 10);
330 2 : assert_eq!(c.get(&2), Some(10), "the queue holds no stale key");
331 2 : }
332 :
333 : /// Re-inserting a key with a new weight charges the new weight, and
334 : /// invalidating returns it to the budget.
335 : #[test]
336 2 : fn reinsert_and_invalidate_keep_the_weight_exact() {
337 2 : let c = weighed(10);
338 2 : c.insert(0, 2);
339 2 : c.insert(1, 2);
340 2 : c.insert(0, 7);
341 2 : assert_eq!((c.weighted_size(), c.entry_count()), (9, 2));
342 2 : c.insert(2, 3);
343 2 : assert_eq!(
344 2 : (c.get(&0), c.get(&1)),
345 : (None, Some(2)),
346 : "FIFO: the re-inserted 0 kept its place at the head of the queue"
347 : );
348 2 : assert_eq!(c.weighted_size(), 5);
349 2 : c.invalidate(&1);
350 2 : assert_eq!(c.weighted_size(), 3);
351 2 : }
352 :
353 : /// `run_pending_tasks` exists only to match moka's call surface; it must
354 : /// be a no-op that leaves the occupancy unchanged.
355 : #[test]
356 2 : fn run_pending_tasks_is_a_noop() {
357 2 : let c = filled(3, 5);
358 2 : c.run_pending_tasks();
359 2 : assert_eq!(c.entry_count(), 3);
360 2 : assert_eq!(c.get(&4), Some(4));
361 2 : }
362 :
363 : /// Borrowed lookups: a `String`-keyed cache is queried with `&str`, which
364 : /// is how the loader uses it.
365 : #[test]
366 2 : fn borrowed_key_lookup() {
367 2 : let c: Cache<String, u32> = Cache::new(2);
368 2 : c.insert("a".to_owned(), 1);
369 2 : assert_eq!(c.get("a"), Some(1));
370 2 : c.invalidate("a");
371 2 : assert_eq!(c.get("a"), None);
372 2 : assert_eq!(c.entry_count(), 0);
373 2 : }
374 :
375 : /// The bound is enforced under the lock, so it holds no matter how the
376 : /// inserts interleave. (Native-only: wasm32 has no threads — the browser
377 : /// build reaches the cache from a single thread.)
378 : #[cfg(not(target_arch = "wasm32"))]
379 : #[test]
380 2 : fn bound_holds_under_concurrent_inserts() {
381 : use std::sync::Arc;
382 2 : let c: Arc<Cache<u32, u32>> = Arc::new(Cache::new(16));
383 2 : let mut handles = Vec::new();
384 16 : for t in 0..8u32 {
385 16 : let c = Arc::clone(&c);
386 16 : handles.push(std::thread::spawn(move || {
387 8000 : for i in 0..500u32 {
388 8000 : let k = t * 500 + i;
389 8000 : c.insert(k, k);
390 8000 : let _ = c.get(&k);
391 8000 : if k % 97 == 0 {
392 84 : c.invalidate(&k);
393 7916 : }
394 8000 : assert!(c.entry_count() <= 16, "bound breached: {}", c.entry_count());
395 : }
396 16 : }));
397 : }
398 16 : for h in handles {
399 16 : h.join().expect("worker thread");
400 16 : }
401 2 : assert!(c.entry_count() <= 16);
402 2 : }
403 : }
|
