OXYGENHUB
# 8f0efa0729b0 oxygenhub / crates / hub / src / render / cache.rs
4.1 KiB · Rust · last changed Cache rendered Markdown and highlighted files by content, in memory-bounded LRUs 1 month ago source blame history raw
f5c44337 Andy McCall ↶ 1 
//! Caches of rendered output keyed by what the output depends on (hub spec
2 
//! §16.3): the BLAKE3 of the content and of anything else that changes the
3 
//! result, such as a Markdown file's link bases or a file's name, which
4 
//! picks its syntax. Bounded by the bytes they hold; the least recently
5 
//! used entries go first.
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 
/// Something whose size counts against a cache's budget.
24 
pub trait Weigh {
25 
    fn weight(&self) -> usize;
26 
}
27
28 
impl Weigh for String {
29 
    fn weight(&self) -> usize {
30 
        self.len() + 32
31 
    }
32 
}
33
34 
impl Weigh for Vec<String> {
35 
    fn weight(&self) -> usize {
36 
        self.iter().map(|s| s.len() + 24).sum::<usize>() + 32
37 
    }
38 
}
39
40 
/// The key for a set of parts: BLAKE3 of each part's length and bytes.
41 
pub fn key(parts: &[&[u8]]) -> [u8; 32] {
42 
    let mut h = blake3::Hasher::new();
43 
    for p in parts {
44 
        h.update(&(p.len() as u64).to_le_bytes());
45 
        h.update(p);
46 
    }
47 
    *h.finalize().as_bytes()
48 
}
49
50 
impl<V: Weigh> Lru<V> {
51 
    pub fn new(budget: usize) -> Lru<V> {
52 
        Lru { inner: Mutex::new(Inner { budget, used: 0, tick: 0, hits: 0, misses: 0, map: HashMap::new() }) }
53 
    }
54
55 
    /// The cached value for `key`, or `make()` stored and returned. `make`
56 
    /// runs without the lock held, so a slow render blocks nobody.
57 
    pub fn get_or(&self, key: [u8; 32], make: impl FnOnce() -> V) -> Arc<V> {
58 
        {
59 
            let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner());
60 
            g.tick += 1;
61 
            let tick = g.tick;
62 
            if let Some(entry) = g.map.get_mut(&key) {
63 
                entry.0 = tick;
64 
                let v = entry.1.clone();
65 
                g.hits += 1;
66 
                return v;
67 
            }
68 
            g.misses += 1;
69 
        }
70 
        let value = Arc::new(make());
71 
        let weight = value.weight();
72 
        let mut g = self.inner.lock().unwrap_or_else(|p| p.into_inner());
73 
        if weight > g.budget / 4 {
74 
            return value; // too big to be worth keeping
75 
        }
76 
        g.tick += 1;
77 
        let tick = g.tick;
78 
        if let Some(old) = g.map.insert(key, (tick, value.clone(), weight)) {
79 
            g.used -= old.2;
80 
        }
81 
        g.used += weight;
82 
        while g.used > g.budget {
83 
            let Some((&oldest, _)) = g.map.iter().min_by_key(|(_, e)| e.0) else { break };
84 
            if let Some(e) = g.map.remove(&oldest) {
85 
                g.used -= e.2;
86 
            }
87 
        }
88 
        value
89 
    }
90
91 
    /// (hits, misses, bytes held).
92 
    pub fn stats(&self) -> (u64, u64, usize) {
93 
        let g = self.inner.lock().unwrap_or_else(|p| p.into_inner());
94 
        (g.hits, g.misses, g.used)
95 
    }
96 
}
97
98 
#[cfg(test)]
99 
mod tests {
100 
    use super::*;
101
102 
    #[test]
103 
    fn hits_misses_and_eviction() {
104 
        let c: Lru<String> = Lru::new(400);
105 
        let mut runs = 0;
106 
        for _ in 0..3 {
107 
            c.get_or(key(&[b"a"]), || {
108 
                runs += 1;
109 
                "x".repeat(50)
110 
            });
111 
        }
112 
        assert_eq!(runs, 1, "rendered once");
113 
        assert_eq!(c.stats().0, 2);
114 
        // Filling past the budget drops the least recently used.
115 
        c.get_or(key(&[b"b"]), || "y".repeat(90));
116 
        c.get_or(key(&[b"a"]), || unreachable!());
117 
        c.get_or(key(&[b"c"]), || "z".repeat(90));
118 
        c.get_or(key(&[b"d"]), || "w".repeat(90));
119 
        assert!(c.stats().2 <= 400);
120 
        let mut again = false;
121 
        c.get_or(key(&[b"b"]), || {
122 
            again = true;
123 
            String::new()
124 
        });
125 
        assert!(again, "b was the least recently used and went first");
126 
        // Something larger than a quarter of the budget is not kept.
127 
        c.get_or(key(&[b"huge"]), || "h".repeat(200));
128 
        let mut rerun = false;
129 
        c.get_or(key(&[b"huge"]), || {
130 
            rerun = true;
131 
            String::new()
132 
        });
133 
        assert!(rerun);
134 
        // Parts are length-framed: ("ab","c") and ("a","bc") differ.
135 
        assert_ne!(key(&[b"ab", b"c"]), key(&[b"a", b"bc"]));
136 
    }
137 
}