| 1 | |
|
| 2 | |
|
| 3 | |
|
| 4 | |
|
| 5 | |
|
| 6 | |
|
| 7 | | use std::collections::HashMap; |
| 8 | | use std::sync::{Arc, Mutex}; |
| 9 | |
|
| 10 | | pub struct Lru<V> { |
| 11 | | inner: Mutex<Inner<V>>, |
| 12 | | } |
| 13 | |
|
| 14 | | struct Inner<V> { |
| 15 | | budget: usize, |
| 16 | | used: usize, |
| 17 | | tick: u64, |
| 18 | | hits: u64, |
| 19 | | misses: u64, |
| 20 | | map: HashMap<[u8; 32], (u64, Arc<V>, usize)>, |
| 21 | | } |
| 22 | |
|
| 23 | |
|
| 24 | | pub trait Weigh { |
| 25 | | fn weight(&self) -> usize; |
| 26 | | } |
| 27 | |
|
| 28 | | impl Weigh for String { |
| 29 | 65× | fn weight(&self) -> usize { |
| 30 | 65× | self.len() + 32 |
| 31 | 65× | } |
| 32 | | } |
| 33 | |
|
| 34 | | impl Weigh for Vec<String> { |
| 35 | 16× | fn weight(&self) -> usize { |
| 36 | 49× | self.iter().map(|s| s.len() + 24).sum::<usize>() + 32 |
| 37 | 16× | } |
| 38 | | } |
| 39 | |
|
| 40 | |
|
| 41 | 342× | pub fn key(parts: &[&[u8]]) -> [u8; 32] { |
| 42 | 342× | let mut h = blake3::Hasher::new(); |
| 43 | 1261× | for p in parts { |
| 44 | 1261× | h.update(&(p.len() as u64).to_le_bytes()); |
| 45 | 1261× | h.update(p); |
| 46 | 1261× | } |
| 47 | 342× | *h.finalize().as_bytes() |
| 48 | 342× | } |
| 49 | |
|
| 50 | | impl<V: Weigh> Lru<V> { |
| 51 | 48× | pub fn new(budget: usize) -> Lru<V> { |
| 52 | 48× | Lru { inner: Mutex::new(Inner { budget, used: 0, tick: 0, hits: 0, misses: 0, map: HashMap::new() }) } |
| 53 | 48× | } |
| 54 | |
|
| 55 | | |
| 56 | | |
| 57 | 340× | pub fn get_or(&self, key: [u8; 32], make: impl FnOnce() -> V) -> Arc<V> { |
| 58 | | { |
| 59 | 340× | let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); |
| 60 | 340× | g.tick += 1; |
| 61 | 340× | let tick = g.tick; |
| 62 | 340× | if let Some(entry) = g.map.get_mut(&key) { |
| 63 | 244× | entry.0 = tick; |
| 64 | 244× | let v = entry.1.clone(); |
| 65 | 244× | g.hits += 1; |
| 66 | 244× | return v; |
| 67 | 96× | } |
| 68 | 96× | g.misses += 1; |
| 69 | | } |
| 70 | 96× | let value = Arc::new(make()); |
| 71 | 96× | let weight = value.weight(); |
| 72 | 96× | let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); |
| 73 | 96× | if weight > g.budget / 4 { |
| 74 | 4× | return value; |
| 75 | 92× | } |
| 76 | 92× | g.tick += 1; |
| 77 | 92× | let tick = g.tick; |
| 78 | 92× | if let Some(old) = g.map.insert(key, (tick, value.clone(), weight)) { |
| 79 | 0× | g.used -= old.2; |
| 80 | 92× | } |
| 81 | 92× | g.used += weight; |
| 82 | 92× | while g.used > g.budget { |
| 83 | 0× | let Some((&oldest, _)) = g.map.iter().min_by_key(|(_, e)| e.0) else { break }; |
| 84 | 0× | if let Some(e) = g.map.remove(&oldest) { |
| 85 | 0× | g.used -= e.2; |
| 86 | 0× | } |
| 87 | | } |
| 88 | 92× | value |
| 89 | 340× | } |
| 90 | |
|
| 91 | | |
| 92 | 4× | pub fn stats(&self) -> (u64, u64, usize) { |
| 93 | 4× | let g = self.inner.lock().unwrap_or_else(|p| p.into_inner()); |
| 94 | 4× | (g.hits, g.misses, g.used) |
| 95 | 4× | } |
| 96 | | } |
| 97 | |
|
| 98 | | #[cfg(test)] |
| 99 | | mod tests { |
| 100 | | use super::*; |
| 101 | |
|
| 102 | | #[test] |
| 103 | 1× | fn hits_misses_and_eviction() { |
| 104 | 1× | let c: Lru<String> = Lru::new(400); |
| 105 | 1× | let mut runs = 0; |
| 106 | 1× | for _ in 0..3 { |
| 107 | 3× | c.get_or(key(&[b"a"]), || { |
| 108 | 1× | runs += 1; |
| 109 | 1× | "x".repeat(50) |
| 110 | 1× | }); |
| 111 | | } |
| 112 | 1× | assert_eq!(runs, 1, "rendered once"); |
| 113 | 1× | assert_eq!(c.stats().0, 2); |
| 114 | | |
| 115 | 1× | c.get_or(key(&[b"b"]), || "y".repeat(90)); |
| 116 | 1× | c.get_or(key(&[b"a"]), || unreachable!()); |
| 117 | 1× | c.get_or(key(&[b"c"]), || "z".repeat(90)); |
| 118 | 1× | c.get_or(key(&[b"d"]), || "w".repeat(90)); |
| 119 | 1× | assert!(c.stats().2 <= 400); |
| 120 | 1× | let mut again = false; |
| 121 | 1× | c.get_or(key(&[b"b"]), || { |
| 122 | 1× | again = true; |
| 123 | 1× | String::new() |
| 124 | 1× | }); |
| 125 | 1× | assert!(again, "b was the least recently used and went first"); |
| 126 | | |
| 127 | 1× | c.get_or(key(&[b"huge"]), || "h".repeat(200)); |
| 128 | 1× | let mut rerun = false; |
| 129 | 1× | c.get_or(key(&[b"huge"]), || { |
| 130 | 1× | rerun = true; |
| 131 | 1× | String::new() |
| 132 | 1× | }); |
| 133 | 1× | assert!(rerun); |
| 134 | | |
| 135 | 1× | assert_ne!(key(&[b"ab", b"c"]), key(&[b"a", b"bc"])); |
| 136 | 1× | } |
| 137 | | } |