rust_zero_core/
cache.rs

1use std::{
2    collections::{HashMap, VecDeque},
3    future::Future,
4    hash::Hash,
5    sync::{
6        atomic::{AtomicU64, Ordering},
7        Mutex,
8    },
9    time::{Duration, Instant},
10};
11
12/// Adds a bounded, deterministic-per-call spread to an expiry duration.
13///
14/// Cache adapters use this to avoid a large set of records expiring at exactly the same instant.
15/// The caller owns the sequence so independent cache instances do not contend on a global RNG.
16#[cfg(any(
17    feature = "stores-sql",
18    feature = "stores-mongo",
19    feature = "stores-redis",
20    test
21))]
22pub(crate) fn jittered_ttl(base: Duration, jitter: Duration, sequence: u64) -> Duration {
23    if jitter.is_zero() {
24        return base;
25    }
26
27    // SplitMix64 gives a well-distributed value without pulling an RNG into the cache hot path.
28    let mut value = sequence.wrapping_add(0x9e37_79b9_7f4a_7c15);
29    value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
30    value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
31    value ^= value >> 31;
32
33    let jitter_nanos = jitter.as_nanos();
34    let added_nanos = u128::from(value) % (jitter_nanos.saturating_add(1));
35    let added = Duration::new(
36        (added_nanos / 1_000_000_000).min(u128::from(u64::MAX)) as u64,
37        (added_nanos % 1_000_000_000) as u32,
38    );
39    base.saturating_add(added)
40}
41
42use crate::{SingleFlight, SingleFlightError};
43
44struct Entry<V> {
45    value: V,
46    expires_at: Instant,
47}
48
49/// A thread-safe in-memory cache that expires entries on access.
50pub struct TtlCache<K, V> {
51    entries: Mutex<HashMap<K, Entry<V>>>,
52}
53
54impl<K, V> Default for TtlCache<K, V>
55where
56    K: Eq + Hash,
57{
58    fn default() -> Self {
59        Self {
60            entries: Mutex::new(HashMap::new()),
61        }
62    }
63}
64
65impl<K, V> TtlCache<K, V>
66where
67    K: Eq + Hash,
68{
69    pub fn new() -> Self {
70        Self::default()
71    }
72
73    pub fn insert(&self, key: K, value: V, ttl: Duration) {
74        assert!(!ttl.is_zero(), "cache TTL must be greater than zero");
75        self.entries.lock().expect("cache lock poisoned").insert(
76            key,
77            Entry {
78                value,
79                expires_at: Instant::now() + ttl,
80            },
81        );
82    }
83
84    pub fn get(&self, key: &K) -> Option<V>
85    where
86        V: Clone,
87    {
88        let mut entries = self.entries.lock().expect("cache lock poisoned");
89        let expired = entries
90            .get(key)
91            .is_some_and(|entry| entry.expires_at <= Instant::now());
92        if expired {
93            entries.remove(key);
94            None
95        } else {
96            entries.get(key).map(|entry| entry.value.clone())
97        }
98    }
99
100    pub fn remove(&self, key: &K) -> Option<V> {
101        self.entries
102            .lock()
103            .expect("cache lock poisoned")
104            .remove(key)
105            .map(|entry| entry.value)
106    }
107
108    pub fn len(&self) -> usize {
109        let now = Instant::now();
110        let mut entries = self.entries.lock().expect("cache lock poisoned");
111        entries.retain(|_, entry| entry.expires_at > now);
112        entries.len()
113    }
114
115    pub fn is_empty(&self) -> bool {
116        self.len() == 0
117    }
118}
119
120/// A snapshot of cache hit, miss, insertion, and eviction counters.
121#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
122pub struct CacheStats {
123    pub hits: u64,
124    pub misses: u64,
125    pub insertions: u64,
126    pub evictions: u64,
127}
128
129#[derive(Default)]
130struct CacheCounters {
131    hits: AtomicU64,
132    misses: AtomicU64,
133    insertions: AtomicU64,
134    evictions: AtomicU64,
135}
136
137struct MemoryState<K, V> {
138    entries: HashMap<K, Entry<V>>,
139    least_to_most_recent: VecDeque<K>,
140}
141
142/// A bounded in-process LRU cache with per-entry expiry and cache statistics.
143///
144/// Unlike [`TtlCache`], this cache has a hard item limit. Reads refresh recency, and expired
145/// entries are removed lazily. This mirrors go-zero's bounded memory-cache behavior without
146/// requiring a background timing-wheel task.
147pub struct MemoryCache<K, V> {
148    capacity: usize,
149    state: Mutex<MemoryState<K, V>>,
150    counters: CacheCounters,
151}
152
153impl<K, V> MemoryCache<K, V>
154where
155    K: Clone + Eq + Hash,
156{
157    pub fn new(capacity: usize) -> Self {
158        assert!(capacity > 0, "cache capacity must be greater than zero");
159        Self {
160            capacity,
161            state: Mutex::new(MemoryState {
162                entries: HashMap::new(),
163                least_to_most_recent: VecDeque::new(),
164            }),
165            counters: CacheCounters::default(),
166        }
167    }
168
169    pub fn insert(&self, key: K, value: V, ttl: Duration) {
170        assert!(!ttl.is_zero(), "cache TTL must be greater than zero");
171        let mut state = self.state.lock().expect("cache lock poisoned");
172        remove_from_recency(&mut state.least_to_most_recent, &key);
173        state.least_to_most_recent.push_back(key.clone());
174        state.entries.insert(
175            key,
176            Entry {
177                value,
178                expires_at: Instant::now() + ttl,
179            },
180        );
181        self.counters.insertions.fetch_add(1, Ordering::Relaxed);
182
183        while state.entries.len() > self.capacity {
184            let Some(oldest) = state.least_to_most_recent.pop_front() else {
185                break;
186            };
187            if state.entries.remove(&oldest).is_some() {
188                self.counters.evictions.fetch_add(1, Ordering::Relaxed);
189            }
190        }
191    }
192
193    pub fn get(&self, key: &K) -> Option<V>
194    where
195        V: Clone,
196    {
197        let mut state = self.state.lock().expect("cache lock poisoned");
198        let now = Instant::now();
199        let expired = state
200            .entries
201            .get(key)
202            .is_some_and(|entry| entry.expires_at <= now);
203        if expired {
204            state.entries.remove(key);
205            remove_from_recency(&mut state.least_to_most_recent, key);
206            self.counters.evictions.fetch_add(1, Ordering::Relaxed);
207            self.counters.misses.fetch_add(1, Ordering::Relaxed);
208            return None;
209        }
210
211        let value = state.entries.get(key).map(|entry| entry.value.clone());
212        if value.is_some() {
213            remove_from_recency(&mut state.least_to_most_recent, key);
214            state.least_to_most_recent.push_back(key.clone());
215            self.counters.hits.fetch_add(1, Ordering::Relaxed);
216        } else {
217            self.counters.misses.fetch_add(1, Ordering::Relaxed);
218        }
219        value
220    }
221
222    pub fn remove(&self, key: &K) -> Option<V> {
223        let mut state = self.state.lock().expect("cache lock poisoned");
224        remove_from_recency(&mut state.least_to_most_recent, key);
225        state.entries.remove(key).map(|entry| entry.value)
226    }
227
228    /// Removes every entry accepted by `predicate` and returns the number removed.
229    ///
230    /// The predicate and removals run while holding one cache lock, making this useful for
231    /// atomically clearing secondary-index entries that point at a mutated primary key.
232    pub fn remove_where<F>(&self, mut predicate: F) -> usize
233    where
234        F: FnMut(&K, &V) -> bool,
235    {
236        let mut state = self.state.lock().expect("cache lock poisoned");
237        let keys: Vec<_> = state
238            .entries
239            .iter()
240            .filter(|(key, entry)| predicate(key, &entry.value))
241            .map(|(key, _)| key.clone())
242            .collect();
243        for key in &keys {
244            state.entries.remove(key);
245            remove_from_recency(&mut state.least_to_most_recent, key);
246        }
247        keys.len()
248    }
249
250    pub fn clear(&self) {
251        let mut state = self.state.lock().expect("cache lock poisoned");
252        state.entries.clear();
253        state.least_to_most_recent.clear();
254    }
255
256    pub fn len(&self) -> usize {
257        let now = Instant::now();
258        let mut state = self.state.lock().expect("cache lock poisoned");
259        let expired: Vec<_> = state
260            .entries
261            .iter()
262            .filter(|(_, entry)| entry.expires_at <= now)
263            .map(|(key, _)| key.clone())
264            .collect();
265        for key in &expired {
266            state.entries.remove(key);
267            remove_from_recency(&mut state.least_to_most_recent, key);
268        }
269        self.counters
270            .evictions
271            .fetch_add(expired.len() as u64, Ordering::Relaxed);
272        state.entries.len()
273    }
274
275    pub fn is_empty(&self) -> bool {
276        self.len() == 0
277    }
278
279    pub fn stats(&self) -> CacheStats {
280        CacheStats {
281            hits: self.counters.hits.load(Ordering::Relaxed),
282            misses: self.counters.misses.load(Ordering::Relaxed),
283            insertions: self.counters.insertions.load(Ordering::Relaxed),
284            evictions: self.counters.evictions.load(Ordering::Relaxed),
285        }
286    }
287}
288
289fn remove_from_recency<K>(recency: &mut VecDeque<K>, key: &K)
290where
291    K: Eq,
292{
293    if let Some(index) = recency.iter().position(|candidate| candidate == key) {
294        recency.remove(index);
295    }
296}
297
298/// A bounded cache that coalesces concurrent misses for the same key.
299pub struct ReadThroughCache<K, V, E> {
300    cache: MemoryCache<K, V>,
301    flights: SingleFlight<K, V, E>,
302}
303
304impl<K, V, E> ReadThroughCache<K, V, E>
305where
306    K: Clone + Eq + Hash,
307    V: Clone,
308{
309    pub fn new(capacity: usize) -> Self {
310        Self {
311            cache: MemoryCache::new(capacity),
312            flights: SingleFlight::new(),
313        }
314    }
315
316    pub fn get(&self, key: &K) -> Option<V> {
317        self.cache.get(key)
318    }
319
320    /// Returns a cached value or invokes `fetch` once across concurrent callers.
321    pub async fn take<F, Fut>(
322        &self,
323        key: K,
324        ttl: Duration,
325        fetch: F,
326    ) -> Result<V, SingleFlightError<E>>
327    where
328        F: FnOnce() -> Fut,
329        Fut: Future<Output = Result<V, E>>,
330    {
331        if let Some(value) = self.cache.get(&key) {
332            return Ok(value);
333        }
334
335        self.flights
336            .execute(key.clone(), || async {
337                if let Some(value) = self.cache.get(&key) {
338                    return Ok(value);
339                }
340                let value = fetch().await?;
341                self.cache.insert(key, value.clone(), ttl);
342                Ok(value)
343            })
344            .await
345    }
346
347    pub fn remove(&self, key: &K) -> Option<V> {
348        self.cache.remove(key)
349    }
350
351    pub fn stats(&self) -> CacheStats {
352        self.cache.stats()
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359
360    #[test]
361    fn expiry_jitter_is_bounded_and_varies_by_sequence() {
362        let base = Duration::from_secs(10);
363        let jitter = Duration::from_secs(5);
364        let values: Vec<_> = (0..8)
365            .map(|sequence| jittered_ttl(base, jitter, sequence))
366            .collect();
367
368        assert!(values
369            .iter()
370            .all(|ttl| *ttl >= base && *ttl <= base + jitter));
371        assert!(values.windows(2).any(|pair| pair[0] != pair[1]));
372        assert_eq!(jittered_ttl(base, Duration::ZERO, 99), base);
373    }
374    use std::thread;
375
376    #[test]
377    fn returns_a_value_until_its_ttl_expires() {
378        let cache = TtlCache::new();
379        cache.insert("user-42", "cached", Duration::from_millis(5));
380
381        assert_eq!(cache.get(&"user-42"), Some("cached"));
382        thread::sleep(Duration::from_millis(10));
383        assert_eq!(cache.get(&"user-42"), None);
384        assert!(cache.is_empty());
385    }
386
387    #[test]
388    fn remove_returns_the_cached_value() {
389        let cache = TtlCache::new();
390        cache.insert("user-42", 42, Duration::from_secs(1));
391
392        assert_eq!(cache.remove(&"user-42"), Some(42));
393        assert_eq!(cache.get(&"user-42"), None);
394    }
395
396    #[test]
397    fn bounded_cache_evicts_the_least_recently_used_value() {
398        let cache = MemoryCache::new(2);
399        cache.insert("one", 1, Duration::from_secs(1));
400        cache.insert("two", 2, Duration::from_secs(1));
401        assert_eq!(cache.get(&"one"), Some(1));
402        cache.insert("three", 3, Duration::from_secs(1));
403
404        assert_eq!(cache.get(&"two"), None);
405        assert_eq!(cache.get(&"one"), Some(1));
406        assert_eq!(cache.get(&"three"), Some(3));
407        assert_eq!(cache.stats().evictions, 1);
408    }
409
410    #[test]
411    fn bounded_cache_removes_matching_values_atomically() {
412        let cache = MemoryCache::new(4);
413        cache.insert("email", Some(1), Duration::from_secs(1));
414        cache.insert("phone", Some(1), Duration::from_secs(1));
415        cache.insert("other", Some(2), Duration::from_secs(1));
416        cache.insert("missing", None, Duration::from_secs(1));
417
418        assert_eq!(cache.remove_where(|_, value| *value == Some(1)), 2);
419        assert_eq!(cache.get(&"email"), None);
420        assert_eq!(cache.get(&"phone"), None);
421        assert_eq!(cache.get(&"other"), Some(Some(2)));
422        assert_eq!(cache.get(&"missing"), Some(None));
423    }
424
425    #[tokio::test]
426    async fn read_through_cache_fetches_and_reuses_a_value() {
427        use std::sync::atomic::{AtomicUsize, Ordering};
428
429        let cache = ReadThroughCache::<String, usize, String>::new(10);
430        let calls = AtomicUsize::new(0);
431        for _ in 0..2 {
432            let value = cache
433                .take("answer".to_owned(), Duration::from_secs(1), || async {
434                    calls.fetch_add(1, Ordering::SeqCst);
435                    Ok(42)
436                })
437                .await
438                .unwrap();
439            assert_eq!(value, 42);
440        }
441        assert_eq!(calls.load(Ordering::SeqCst), 1);
442    }
443}
rust_zero_core/cache.rs

rust_zero_core/
cache.rs
