diff --git a/Cargo.lock b/Cargo.lock index f53859d7ae..7a5289a6f8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,12 @@ dependencies = [ "memchr", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "anyhow" version = "1.0.104" @@ -337,6 +343,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "form_urlencoded" version = "1.2.1" @@ -496,6 +508,17 @@ version = "0.15.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + [[package]] name = "headers" version = "0.4.1" @@ -822,7 +845,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.15.4", ] [[package]] @@ -925,6 +948,15 @@ version = "0.4.34" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f9f8bd3e56ce4dfc153cf470fffbfa98c7620958b312ca5c3a4b8d5181fd13c6" +[[package]] +name = "lru" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" +dependencies = [ + "hashbrown 0.16.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -1710,6 +1742,7 @@ dependencies = [ "libc", "libloading", "log", + "lru", "mime", "once_cell", "postgres-protocol", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index e7cabf2ef3..1c3f10c455 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -94,6 +94,7 @@ rustls-pemfile = "2" tokio-postgres-rustls = { version = "0.14.0", default-features = false, features = [ "ring", ] } +lru = "0.16.4" [build-dependencies] blake2 = "0.11.0" diff --git a/rust/src/caches/mod.rs b/rust/src/caches/mod.rs new file mode 100644 index 0000000000..966f17d1f9 --- /dev/null +++ b/rust/src/caches/mod.rs @@ -0,0 +1,917 @@ +/* + * This file is licensed under the Affero General Public License (AGPL) version 3. + * + * Copyright (C) 2026 Element Creations Ltd + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * See the GNU Affero General Public License for more details: + * . + * + */ + +//! Caching for database functions that live in Rust. +//! +//! This is the Rust counterpart to `synapse.util.caches.descriptors.cached`: +//! a bounded LRU of loaded values, with in-flight requests for the same key +//! coalesced into one database round trip. +//! +//! Invalidation deliberately stays in Python. Synapse invalidates caches by +//! *name*, resolving `getattr(store, cache_name)` and calling `.invalidate(key)` +//! on whatever it finds (`synapse/storage/_base.py`), so a cache here only has +//! to be bound onto the datastore under the right name and expose that method. +//! See [`python`]. Everything else (the replication stream, `txn.call_after`, +//! the `*_fake` names) then works unchanged. +//! +//! # Locking +//! +//! **Never hold [`Cache`]'s mutex across an `.await` or a GIL acquisition.** +//! +//! Unlike Python's caches, which are single-threaded by assertion, this one is +//! reached both from Tokio workers (which hold no GIL) and from the reactor +//! thread by way of invalidation (which does). Taking the GIL under the cache +//! lock would invert the order against those Tokio readers and deadlock. The +//! rule holds for value drops too. A cached value may own a `Py`, whose +//! drop needs the GIL, which is why every eviction path here *returns* the +//! values it removed rather than dropping them in place. +//! +//! Keys get no such treatment: they are dropped under the lock. So a key type +//! must not own Python objects. Every key so far is built from strings, and a +//! key that needs to hold a `Py` should extract what it hashes on +//! instead. + +use std::{ + collections::HashMap, + fmt, + future::Future, + hash::Hash, + num::NonZeroUsize, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, MutexGuard, + }, + time::SystemTime, +}; + +use futures::channel::oneshot; + +use crate::clock::Clock; + +/// An error from a load, shared with everyone who was waiting on it. +/// +/// `anyhow::Error` isn't `Clone`, and a single failed load has to be handed to +/// an arbitrary number of waiters, so we share one behind an `Arc`. +#[derive(Debug, Clone)] +pub struct SharedError(Arc); + +impl SharedError { + fn new(error: anyhow::Error) -> Self { + SharedError(Arc::new(error)) + } +} + +impl fmt::Display for SharedError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + // `{:#}` so the whole context chain survives. Otherwise wrapping this + // back up in an `anyhow::Error` would lose everything but the outermost + // message. + write!(f, "{:#}", self.0) + } +} + +impl std::error::Error for SharedError {} + +/// A cached value, and when it was last read. +struct Entry { + value: V, + /// From [`Cache::clock`]. Used by the age-based expiry in [`registry`]. + last_read: SystemTime, +} + +struct Inner { + entries: lru::LruCache>, + + /// Keys with a load in flight, and the id of that load. + /// + /// The id is what makes invalidating a key mid-load safe: a load only gets + /// to populate the cache if the id it started with is still the one here. + pending: HashMap, + + /// Callers waiting on an in-flight load, keyed by that load's id. + waiters: HashMap>>>, + + next_load_id: u64, + + /// The size this cache was declared with, before any cache factor is + /// applied. `set_cache_factor` scales this rather than the current size, so + /// that repeated resizes don't compound. + original_max_size: usize, +} + +/// A bounded, LRU cache of values loaded from somewhere slow. +/// +/// Create one with [`registry::CacheRegistry::new_cache`] and read through it +/// with [`Cache::get_or_load`]. See the module docs for the locking rule. +pub struct Cache { + name: &'static str, + + /// Stamps entries with their last access time. Read without the GIL, from + /// under the cache lock. See [`crate::clock`]. + clock: Clock, + + hits: AtomicU64, + misses: AtomicU64, + evicted_size: AtomicU64, + evicted_time: AtomicU64, + evicted_invalidation: AtomicU64, + + inner: Mutex>, +} + +impl Cache +where + K: Eq + Hash + Clone + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, +{ + /// Build a cache holding at most `max_size` entries. + /// + /// This does not enrol the cache for age-based expiry. Go through + /// [`registry::CacheRegistry::new_cache`] for that. + /// + /// `name` is the name Synapse knows the cache by: the datastore attribute, + /// the metrics label, and what travels over the cache invalidation stream. + pub fn new(name: &'static str, max_size: usize) -> Arc { + Arc::new(Cache { + name, + clock: Clock::new(), + hits: AtomicU64::new(0), + misses: AtomicU64::new(0), + evicted_size: AtomicU64::new(0), + evicted_time: AtomicU64::new(0), + evicted_invalidation: AtomicU64::new(0), + inner: Mutex::new(Inner { + entries: lru::LruCache::new(non_zero(max_size)), + pending: HashMap::new(), + waiters: HashMap::new(), + next_load_id: 0, + original_max_size: max_size, + }), + }) + } + + /// Look `key` up, loading and caching it if it isn't here. + /// + /// Concurrent callers for the same key share a single load: the first + /// caller runs `load`, the rest wait for its result. + /// + /// Failures are never cached. Every waiter on a failed load gets the + /// error, and the next caller tries again. This matches + /// `DeferredCache.set`, which evicts the key when the deferred errbacks. + pub async fn get_or_load(&self, key: K, load: F) -> Result + where + F: Fn() -> Fut, + Fut: Future>, + { + loop { + let load_id = match self.start_load(&key) { + Started::Lead(load_id) => load_id, + Started::Cached(value) => return Ok(value), + Started::Waiting(receiver) => match receiver.await { + // Someone else's load answered for us, so this is a hit, + // as joining an in-flight load is in `DeferredCache.get`. + // It is counted here rather than when we joined, because + // a leader can go away before it answers, and then we go + // round again and get charged as whatever we become. + Ok(result) => { + self.hits.fetch_add(1, Ordering::Relaxed); + return result; + } + // The sender went away, which means the caller that was + // loading this key was dropped or panicked before it could + // finish. Go round again. One of us will pick the work up. + Err(_) => continue, + }, + }; + + // From here until `finish_load` we own the pending marker for this + // key, and must clear it however we leave. The guard covers the + // paths we can't: a panic in `load`, or this future being dropped. + let guard = LoadGuard { + cache: self, + key: &key, + load_id, + armed: true, + }; + + let result = load().await.map_err(SharedError::new); + + // We made it, so hand the bookkeeping over to `finish_load`. + guard.finish(); + + return self.finish_load(key, load_id, result); + } + } + + /// Drop `key` from the cache, and stop any in-flight load from caching what + /// it finds. + /// + /// Returns the value that was cached, for the caller to drop. + /// + /// Callers already waiting on an in-flight load still get their result. + /// It just doesn't get cached. That matches Python, where invalidating a + /// key drops the pending deferred from the cache without disturbing the + /// `ObservableDeferred` its observers are waiting on. + pub fn invalidate(&self, key: &K) -> Option { + let mut inner = self.lock(); + + // Any load in flight for this key started before we knew the value had + // changed, so whatever it returns may already be stale. + inner.pending.remove(key); + + let removed = inner.entries.pop(key).map(|entry| entry.value); + if removed.is_some() { + self.evicted_invalidation.fetch_add(1, Ordering::Relaxed); + } + + removed + } + + /// Empty the cache, and stop every in-flight load from caching what it + /// finds. Returns the values that were cached, for the caller to drop. + pub fn invalidate_all(&self) -> Vec { + let mut inner = self.lock(); + + inner.pending.clear(); + + let mut removed = Vec::with_capacity(inner.entries.len()); + while let Some((_, entry)) = inner.entries.pop_lru() { + removed.push(entry.value); + } + + self.evicted_invalidation + .fetch_add(removed.len() as u64, Ordering::Relaxed); + + removed + } + + /// The cached value for `key`, counted as a hit, or `None`. + /// + /// `None` is not counted as a miss. This is the synchronous fast path for + /// a caller that will go on to [`Cache::get_or_load`], which counts that + /// lookup, if the value isn't here. + pub fn get(&self, key: &K) -> Option { + let value = self.lock().get(key, self.clock.now()); + if value.is_some() { + self.hits.fetch_add(1, Ordering::Relaxed); + } + value + } + + /// Evict every entry that hasn't been read since `cutoff`. Returns the + /// evicted values, for the caller to drop. + pub fn expire_before(&self, cutoff: SystemTime) -> Vec { + let expired = self.lock().expire_before(cutoff); + + self.evicted_time + .fetch_add(expired.len() as u64, Ordering::Relaxed); + + expired + } + + /// Rescale the cache to `original_max_size * factor`, evicting from the + /// cold end if it shrank. Returns the evicted values, for the caller to + /// drop. + /// + /// This is what `SYNAPSE_CACHE_FACTOR_*` and `caches.per_cache_factors` + /// drive, via `synapse.config.cache.add_resizable_cache`. + pub fn set_cache_factor(&self, factor: f64) -> Vec { + let mut inner = self.lock(); + + let new_size = non_zero((inner.original_max_size as f64 * factor) as usize); + if new_size == inner.entries.cap() { + return Vec::new(); + } + + // Evict by hand rather than letting `resize` do it, so that the values + // are dropped by our caller with the lock released. + let mut evicted = Vec::new(); + while inner.entries.len() > new_size.get() { + let (_, entry) = inner.entries.pop_lru().expect("cache is over capacity"); + evicted.push(entry.value); + } + inner.entries.resize(new_size); + + self.evicted_size + .fetch_add(evicted.len() as u64, Ordering::Relaxed); + + evicted + } + + pub fn len(&self) -> usize { + self.lock().entries.len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn max_size(&self) -> usize { + self.lock().entries.cap().get() + } + + pub fn name(&self) -> &'static str { + self.name + } + + /// `(hits, misses)`. + pub fn hit_counts(&self) -> (u64, u64) { + ( + self.hits.load(Ordering::Relaxed), + self.misses.load(Ordering::Relaxed), + ) + } + + /// Eviction counts as `(size, time, invalidation)`, matching Python's + /// `EvictionReason`. + pub fn eviction_counts(&self) -> (u64, u64, u64) { + ( + self.evicted_size.load(Ordering::Relaxed), + self.evicted_time.load(Ordering::Relaxed), + self.evicted_invalidation.load(Ordering::Relaxed), + ) + } + + fn lock(&self) -> MutexGuard<'_, Inner> { + lock_ignoring_poison(&self.inner) + } + + /// Decide what this caller has to do to get a value for `key`: nothing (it + /// is cached), wait for someone else's load, or run the load itself. + fn start_load(&self, key: &K) -> Started { + let mut inner = self.lock(); + + if let Some(value) = inner.get(key, self.clock.now()) { + self.hits.fetch_add(1, Ordering::Relaxed); + return Started::Cached(value); + } + + if let Some(&load_id) = inner.pending.get(key) { + // Counted as a hit once the load answers, in `get_or_load`. + let (sender, receiver) = oneshot::channel(); + inner.waiters.entry(load_id).or_default().push(sender); + return Started::Waiting(receiver); + } + + self.misses.fetch_add(1, Ordering::Relaxed); + + let load_id = inner.next_load_id; + inner.next_load_id += 1; + inner.pending.insert(key.clone(), load_id); + + Started::Lead(load_id) + } + + /// Publish the result of a load: cache it if it is still wanted, and hand + /// it to everyone who was waiting. + fn finish_load( + &self, + key: K, + load_id: u64, + result: Result, + ) -> Result { + let (waiters, displaced) = { + let mut inner = self.lock(); + + // If this is no longer the current load for the key, something + // invalidated it while we were away and what we have may be stale. + // Hand it to our waiters, but don't cache it. + let still_current = inner.pending.get(&key) == Some(&load_id); + if still_current { + inner.pending.remove(&key); + } + + let waiters = inner.waiters.remove(&load_id).unwrap_or_default(); + + let displaced = match &result { + Ok(value) if still_current => { + // The key can't already be present: it wasn't when this + // load started, and nothing else populates a key while a + // load for it is current. So anything `insert` hands back + // is a cold entry evicted to make room. + let displaced = inner.insert(key, value.clone(), self.clock.now()); + if displaced.is_some() { + self.evicted_size.fetch_add(1, Ordering::Relaxed); + } + displaced + } + _ => None, + }; + + (waiters, displaced) + }; + + // With the lock released, so that dropping a value that owns Python + // objects can take the GIL. See the module docs. + drop(displaced); + + for waiter in waiters { + // An error means the waiter gave up, which is fine. + let _ = waiter.send(result.clone()); + } + + result + } +} + +/// The parts of the cache that depend on the time take it as an argument, +/// rather than reading [`Cache::clock`], so that the tests can drive them with +/// explicit timestamps and stay clear of the process-wide virtual clock. +impl Inner +where + K: Eq + Hash, + V: Clone, +{ + /// Look `key` up, stamping its entry as read at `now` and moving it to the + /// head of the LRU. + fn get(&mut self, key: &K, now: SystemTime) -> Option { + let entry = self.entries.get_mut(key)?; + entry.last_read = now; + Some(entry.value.clone()) + } + + /// Insert an entry read at `now`, returning whatever it displaced for the + /// caller to drop. + fn insert(&mut self, key: K, value: V, now: SystemTime) -> Option { + let entry = Entry { + value, + last_read: now, + }; + + self.entries.push(key, entry).map(|(_, entry)| entry.value) + } + + /// Pop every entry not read since `cutoff`, returning the values for the + /// caller to drop. + fn expire_before(&mut self, cutoff: SystemTime) -> Vec { + let mut expired = Vec::new(); + // `last_read` is stamped on read and reads move the entry to the + // head, so the LRU order is also age order: we can stop at the first + // entry that is still live rather than walking the whole cache. + while self + .entries + .peek_lru() + .is_some_and(|(_, entry)| entry.last_read <= cutoff) + { + let (_, entry) = self.entries.pop_lru().expect("peek_lru just succeeded"); + expired.push(entry.value); + } + + expired + } +} + +/// Entry points that take the time as an argument, for tests of the age-based +/// expiry. Production code reads the clock, which tests must not pin: it is +/// process-wide, and `cargo test` runs them in parallel. +#[cfg(test)] +impl Cache +where + K: Eq + Hash + Clone + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, +{ + /// A hit at `now`, or `None` if the key isn't cached. + fn get_at(&self, key: &K, now: SystemTime) -> Option { + self.lock().get(key, now) + } + + /// Insert `value` as if it had been loaded at `now`. + fn insert_at(&self, key: K, value: V, now: SystemTime) { + let displaced = self.lock().insert(key, value, now); + drop(displaced); + } +} + +/// What [`Cache::start_load`] decided. +enum Started { + /// Nobody is loading this key, so the caller runs the load, under this id. + Lead(u64), + Cached(V), + /// Someone else is loading this key. The result arrives here. + Waiting(oneshot::Receiver>), +} + +/// Releases a key's pending marker if the load never completes. +/// +/// A load that is cancelled (its future dropped) or panics would otherwise +/// leave the key marked as in flight forever, and every later caller would +/// queue up behind a load that is never going to finish. +struct LoadGuard<'a, K: Eq + Hash + Clone + Send + Sync + 'static, V: Clone + Send + Sync + 'static> +{ + cache: &'a Cache, + key: &'a K, + load_id: u64, + armed: bool, +} + +impl + LoadGuard<'_, K, V> +{ + /// The load completed, so the pending marker is `finish_load`'s to clear. + /// Consuming the guard also releases its borrow of the key. + fn finish(mut self) { + self.armed = false; + } +} + +impl Drop + for LoadGuard<'_, K, V> +{ + fn drop(&mut self) { + if !self.armed { + return; + } + + let mut inner = self.cache.lock(); + + if inner.pending.get(self.key) == Some(&self.load_id) { + inner.pending.remove(self.key); + } + + // Dropping the senders wakes everyone waiting on us with a cancelled + // channel, which sends them back around to retry. + inner.waiters.remove(&self.load_id); + } +} + +/// Take a lock, recovering from poisoning. +/// +/// Nothing in a locked region in this module runs caller-supplied code, so a +/// panic while holding one means a bug here rather than corrupt state. Failing +/// every subsequent cache read would turn that into an outage. The worst that +/// carrying on can do is serve a stale entry. +pub(crate) fn lock_ignoring_poison(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +/// A cache of size zero would be a cache that never caches anything, which is +/// never what anyone means, and `lru` can't represent it in any case. Treat it +/// as one entry, as `LruCache` does for a cache factor that rounds to nothing. +fn non_zero(size: usize) -> NonZeroUsize { + NonZeroUsize::new(size).unwrap_or(NonZeroUsize::MIN) +} + +#[cfg(test)] +mod tests { + use std::{ + pin::pin, + sync::atomic::{AtomicUsize, Ordering}, + task::Poll, + time::{Duration, UNIX_EPOCH}, + }; + + use futures::{ + executor::block_on, + future::{poll_fn, FutureExt}, + Future, + }; + + use super::*; + + /// `secs` seconds after the epoch, for the expiry tests. + pub(super) fn at(secs: u64) -> SystemTime { + UNIX_EPOCH + Duration::from_secs(secs) + } + + fn cache(max_size: usize) -> Arc> { + Cache::new("test_cache", max_size) + } + + /// A loader that returns `value`, counting how many times it ran. + fn counting_loader( + value: u32, + count: &AtomicUsize, + ) -> impl Fn() -> std::future::Ready> + '_ { + move || { + count.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(value)) + } + } + + #[test] + fn loads_once_then_hits() { + let cache = cache(10); + let loads = AtomicUsize::new(0); + + let first = block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + let second = block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + + assert_eq!((first, second), (1, 1)); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(cache.hit_counts(), (1, 1)); + } + + #[test] + fn evicts_the_least_recently_used() { + let cache = cache(2); + let loads = AtomicUsize::new(0); + + block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + block_on(cache.get_or_load("b", counting_loader(2, &loads))).unwrap(); + // Touch "a" so that "b" becomes the coldest entry. + block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + block_on(cache.get_or_load("c", counting_loader(3, &loads))).unwrap(); + + assert_eq!(cache.len(), 2); + assert_eq!(cache.eviction_counts().0, 1); + + // "b" went, so looking it up loads again. "a" and "c" are still here. + let loads_before = loads.load(Ordering::SeqCst); + block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + block_on(cache.get_or_load("c", counting_loader(3, &loads))).unwrap(); + assert_eq!(loads.load(Ordering::SeqCst), loads_before); + block_on(cache.get_or_load("b", counting_loader(2, &loads))).unwrap(); + assert_eq!(loads.load(Ordering::SeqCst), loads_before + 1); + } + + #[test] + fn invalidate_drops_the_entry() { + let cache = cache(10); + let loads = AtomicUsize::new(0); + + block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + assert_eq!(cache.invalidate(&"a"), Some(1)); + block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + + assert_eq!(loads.load(Ordering::SeqCst), 2); + assert_eq!(cache.eviction_counts().2, 1); + } + + #[test] + fn invalidate_all_empties_the_cache() { + let cache = cache(10); + let loads = AtomicUsize::new(0); + + block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + block_on(cache.get_or_load("b", counting_loader(2, &loads))).unwrap(); + + let mut dropped = cache.invalidate_all(); + dropped.sort_unstable(); + + assert_eq!(dropped, vec![1, 2]); + assert_eq!(cache.len(), 0); + } + + #[test] + fn errors_are_not_cached() { + let cache = cache(10); + let loads = AtomicUsize::new(0); + + let failing = || { + loads.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err(anyhow::anyhow!("nope"))) + }; + + assert!(block_on(cache.get_or_load("a", failing)).is_err()); + assert_eq!(cache.len(), 0); + + // The next caller tries again rather than getting the failure back. + let value = block_on(cache.get_or_load("a", counting_loader(1, &loads))).unwrap(); + assert_eq!(value, 1); + assert_eq!(loads.load(Ordering::SeqCst), 2); + } + + /// A loader that blocks until released, so a test can hold a load open. + struct GatedLoader<'a> { + starts: &'a AtomicUsize, + gate: futures::future::Shared>, + value: u32, + } + + impl<'a> GatedLoader<'a> { + fn new(value: u32, starts: &'a AtomicUsize) -> (Self, oneshot::Sender<()>) { + let (release, gate) = oneshot::channel(); + ( + GatedLoader { + starts, + gate: gate.shared(), + value, + }, + release, + ) + } + + fn load(&self) -> impl Future> { + self.starts.fetch_add(1, Ordering::SeqCst); + + let gate = self.gate.clone(); + let value = self.value; + async move { + let _ = gate.await; + Ok(value) + } + } + } + + /// The highest-risk case: a key invalidated while its load is in flight + /// must not end up cached, or the cache holds a stale value that nothing + /// will ever evict. + #[test] + fn invalidating_an_in_flight_load_does_not_cache_it() { + let cache = cache(10); + let starts = AtomicUsize::new(0); + let (loader, release) = GatedLoader::new(7, &starts); + + block_on(async { + let mut load = pin!(cache.get_or_load("a", || loader.load())); + + // Get as far as the loader, then leave it hanging. + poll_fn(|cx| { + assert!(load.poll_unpin(cx).is_pending()); + Poll::Ready(()) + }) + .await; + assert_eq!(starts.load(Ordering::SeqCst), 1); + + cache.invalidate(&"a"); + release.send(()).unwrap(); + + // The caller still gets its value... + assert_eq!(load.await.unwrap(), 7); + }); + + // ...but the value is not cached, because it may already be stale. + assert_eq!(cache.len(), 0); + } + + #[test] + fn concurrent_misses_share_one_load() { + let cache = cache(10); + let starts = AtomicUsize::new(0); + let (loader, release) = GatedLoader::new(7, &starts); + + let values = block_on(async { + let mut loads = vec![ + cache.get_or_load("a", || loader.load()).boxed_local(), + cache.get_or_load("a", || loader.load()).boxed_local(), + cache.get_or_load("a", || loader.load()).boxed_local(), + ]; + + // Let all three reach the cache: the first becomes the leader, the + // other two queue up behind it. + poll_fn(|cx| { + for load in &mut loads { + assert!(load.poll_unpin(cx).is_pending()); + } + Poll::Ready(()) + }) + .await; + + release.send(()).unwrap(); + futures::future::join_all(loads).await + }); + + assert_eq!( + values.into_iter().map(Result::unwrap).collect::>(), + vec![7, 7, 7] + ); + // Exactly one load ran, and only the leader was charged a miss. + assert_eq!(starts.load(Ordering::SeqCst), 1); + assert_eq!(cache.hit_counts(), (2, 1)); + } + + /// If the caller doing the load goes away, whoever is waiting on it has to + /// pick the work up rather than wait forever. + #[test] + fn a_dropped_load_does_not_wedge_the_key() { + let cache = cache(10); + let starts = AtomicUsize::new(0); + let (loader, _release) = GatedLoader::new(7, &starts); + let loads = AtomicUsize::new(0); + + block_on(async { + { + let mut load = pin!(cache.get_or_load("a", || loader.load())); + poll_fn(|cx| { + assert!(load.poll_unpin(cx).is_pending()); + Poll::Ready(()) + }) + .await; + // ...and now the leader is dropped, still holding the key. + } + + let value = cache + .get_or_load("a", counting_loader(1, &loads)) + .await + .unwrap(); + assert_eq!(value, 1); + }); + + assert_eq!(loads.load(Ordering::SeqCst), 1); + } + + /// A waiter whose leader goes away becomes the leader itself, and the + /// counters see one miss for it, not a hit and then a miss. + #[test] + fn a_waiter_on_a_dropped_load_is_charged_once() { + let cache = cache(10); + let starts = AtomicUsize::new(0); + let (gated, _release) = GatedLoader::new(7, &starts); + let loads = AtomicUsize::new(0); + + let value = block_on(async { + let mut waiter = pin!(cache.get_or_load("a", counting_loader(1, &loads))); + { + let mut leader = pin!(cache.get_or_load("a", || gated.load())); + poll_fn(|cx| { + assert!(leader.poll_unpin(cx).is_pending()); + assert!(waiter.poll_unpin(cx).is_pending()); + Poll::Ready(()) + }) + .await; + // The leader's future is dropped here, with the waiter queued + // behind it. (`pin!` puts it on the stack, so leaving the block + // is what drops it.) + } + + waiter.await.unwrap() + }); + + assert_eq!(value, 1); + assert_eq!(loads.load(Ordering::SeqCst), 1); + assert_eq!(cache.hit_counts(), (0, 2)); + } + + #[test] + fn get_answers_only_from_the_cache() { + let cache = cache(10); + + // A miss here is not counted: the caller goes on to `get_or_load`, + // which counts it. + assert_eq!(cache.get(&"a"), None); + assert_eq!(cache.hit_counts(), (0, 0)); + + cache.insert_at("a", 1, at(1)); + assert_eq!(cache.get(&"a"), Some(1)); + assert_eq!(cache.hit_counts(), (1, 0)); + } + + #[test] + fn expiry_takes_only_entries_past_the_cutoff() { + let cache = cache(10); + + cache.insert_at("old", 1, at(1)); + cache.insert_at("new", 2, at(5)); + + assert_eq!(cache.expire_before(at(3)), vec![1]); + assert_eq!(cache.len(), 1); + assert_eq!(cache.eviction_counts().1, 1); + + // Nothing else is old enough yet, and the walk stops at the first live + // entry rather than clearing the cache. + assert!(cache.expire_before(at(3)).is_empty()); + assert_eq!(cache.len(), 1); + } + + #[test] + fn reading_an_entry_keeps_it_from_expiring() { + let cache = cache(10); + + cache.insert_at("a", 1, at(1)); + cache.insert_at("b", 2, at(2)); + + // Reading `a` restamps it and moves it to the head, so `b` is now the + // coldest entry even though it was inserted later. + assert_eq!(cache.get_at(&"a", at(5)), Some(1)); + + assert_eq!(cache.expire_before(at(3)), vec![2]); + assert_eq!(cache.len(), 1); + assert!(cache.expire_before(at(3)).is_empty()); + assert_eq!(cache.get_at(&"a", at(6)), Some(1)); + } + + #[test] + fn cache_factor_resizes_and_evicts() { + let cache = cache(4); + let loads = AtomicUsize::new(0); + + for key in ["a", "b", "c", "d"] { + block_on(cache.get_or_load(key, counting_loader(1, &loads))).unwrap(); + } + assert_eq!(cache.len(), 4); + + // Halving the size evicts the two coldest entries. + assert_eq!(cache.set_cache_factor(0.5).len(), 2); + assert_eq!(cache.max_size(), 2); + assert_eq!(cache.len(), 2); + + // The factor applies to the declared size, not the current one, so + // going back to 1.0 restores the original. + assert!(cache.set_cache_factor(1.0).is_empty()); + assert_eq!(cache.max_size(), 4); + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index b9a82561eb..453992f42b 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -8,6 +8,7 @@ use pyo3::{exceptions::PyRuntimeError, prelude::*}; use pyo3_log::ResetHandle; pub mod acl; +pub mod caches; pub mod canonical_json; pub mod clock; pub mod config;