From e1df32f1eacbd4f3e286ae04ed0d4a4172451dc2 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Wed, 14 Oct 2026 12:52:19 +0000 Subject: [PATCH] Weigh Rust cache entries, and let callers fill a cache directly The Rust `Cache` bounded itself by entry count, through `lru`'s capacity. The state resolution caches need a bound on the sum of what their values weigh: a state map counts for its entries, a chain's links for each link, as Python's `size_callback` caches do. So the bound is now a weight, with a `Weigher` chosen at construction and the default weighing every value as one, and the cache enforces it itself over an unbounded `lru`. Those caches are also filled by readers that load a batch of keys in one statement rather than through `get_or_load`, so `Cache::insert` is now public, with `prefill` semantics: it leaves any in-flight load for the key to answer its waiters and cache what it found. A cache key that is a single string or integer now decodes from Python, as the keys of the caches moving here are event IDs and group IDs. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_0194eNQP5weBtVMg8NKtc6VU --- rust/src/caches/mod.rs | 252 +++++++++++++++++++++++++++++------- rust/src/caches/python.rs | 32 ++++- rust/src/caches/registry.rs | 26 +++- 3 files changed, 261 insertions(+), 49 deletions(-) diff --git a/rust/src/caches/mod.rs b/rust/src/caches/mod.rs index f6e723b0d7..1914814874 100644 --- a/rust/src/caches/mod.rs +++ b/rust/src/caches/mod.rs @@ -17,7 +17,14 @@ //! //! 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. +//! coalesced into one database round trip. A cache can also be filled +//! directly, with [`Cache::insert`], by code that loads values its own way. +//! +//! The bound is a weight: the sum of what a [`Weigher`] says each cached value +//! weighs. The default weigher counts every value as one, which makes the +//! bound a number of entries, as for `@cached`. A cache of state maps weighs +//! each by its entries instead, as Python's `size_callback` caches do, so a +//! few huge maps cannot pin more memory than many small ones. //! //! Invalidation deliberately stays in Python. Synapse invalidates caches by //! *name*, resolving `getattr(store, cache_name)` and calling `.invalidate(key)` @@ -51,7 +58,6 @@ use std::{ fmt, future::Future, hash::Hash, - num::NonZeroUsize, sync::{ atomic::{AtomicU64, Ordering}, Arc, Mutex, MutexGuard, @@ -87,16 +93,35 @@ impl fmt::Display for SharedError { impl std::error::Error for SharedError {} +/// How much a cached value counts towards the cache's bound. +pub type Weigher = fn(&V) -> usize; + +/// The default [`Weigher`]: every value weighs one, so the bound is a number +/// of entries. +pub fn weigh_one(_: &V) -> usize { + 1 +} + /// 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, + /// What the value counted for when it was inserted. + weight: usize, } struct Inner { + /// Unbounded as far as `lru` is concerned: the bound is [`Inner::max_weight`], + /// which this module enforces itself. entries: lru::LruCache>, + /// The sum of the weights of the entries. + total_weight: usize, + + /// The bound on `total_weight`, after any cache factor. + max_weight: usize, + /// 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 @@ -108,10 +133,10 @@ struct Inner { 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, + /// The bound this cache was declared with, before any cache factor is + /// applied. `set_cache_factor` scales this rather than the current bound, + /// so that repeated resizes don't compound. + original_max_weight: usize, } /// A bounded, LRU cache of values loaded from somewhere slow. @@ -125,6 +150,8 @@ pub struct Cache { /// under the cache lock. See [`crate::clock`]. clock: Clock, + weigher: Weigher, + hits: AtomicU64, misses: AtomicU64, evicted_size: AtomicU64, @@ -147,20 +174,31 @@ where /// `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 { + Self::new_weighted(name, max_size, weigh_one) + } + + /// Build a cache whose values' weights, as `weigher` measures them, sum + /// to at most `max_weight`. A value heavier than the whole cache is never + /// cached. Otherwise as [`Cache::new`]. + pub fn new_weighted(name: &'static str, max_weight: usize, weigher: Weigher) -> Arc { + let max_weight = non_zero(max_weight); Arc::new(Cache { name, clock: Clock::new(), + weigher, 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)), + entries: lru::LruCache::unbounded(), + total_weight: 0, + max_weight, pending: HashMap::new(), waiters: HashMap::new(), next_load_id: 0, - original_max_size: max_size, + original_max_weight: max_weight, }), }) } @@ -234,7 +272,7 @@ where // changed, so whatever it returns may already be stale. inner.pending.remove(key); - let removed = inner.entries.pop(key).map(|entry| entry.value); + let removed = inner.remove(key); if removed.is_some() { self.evicted_invalidation.fetch_add(1, Ordering::Relaxed); } @@ -250,8 +288,8 @@ where inner.pending.clear(); let mut removed = Vec::with_capacity(inner.entries.len()); - while let Some((_, entry)) = inner.entries.pop_lru() { - removed.push(entry.value); + while let Some(value) = inner.pop_lru() { + removed.push(value); } self.evicted_invalidation @@ -273,6 +311,26 @@ where value } + /// Cache `value` under `key`, replacing what was there. Returns the + /// values this displaced (the old value, and any cold entries evicted to + /// make room) for the caller to drop. + /// + /// For code that loads values its own way, as `prefill` is for a Python + /// cache. It leaves any in-flight load for the key alone: that load + /// still answers its waiters, and caches what it found if nothing + /// invalidated the key meanwhile, which is at least as fresh as `value`. + pub fn insert(&self, key: K, value: V) -> Vec { + let weight = (self.weigher)(&value); + let mut displaced = self.lock().insert(key, value, weight, self.clock.now()); + self.evicted_size + .fetch_add(displaced.evicted.len() as u64, Ordering::Relaxed); + if let Some(old) = displaced.old.take() { + self.evicted_invalidation.fetch_add(1, Ordering::Relaxed); + displaced.evicted.push(old); + } + displaced.evicted + } + /// 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 { @@ -284,7 +342,7 @@ where expired } - /// Rescale the cache to `original_max_size * factor`, evicting from the + /// Rescale the cache's bound to `original * factor`, evicting from the /// cold end if it shrank. Returns the evicted values, for the caller to /// drop. /// @@ -293,19 +351,13 @@ where 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() { + let new_max = non_zero((inner.original_max_weight as f64 * factor) as usize); + if new_max == inner.max_weight { return Vec::new(); } + inner.max_weight = new_max; - // 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); + let evicted = inner.evict_to_fit(0); self.evicted_size .fetch_add(evicted.len() as u64, Ordering::Relaxed); @@ -313,6 +365,7 @@ where evicted } + /// The number of cached entries. pub fn len(&self) -> usize { self.lock().entries.len() } @@ -321,8 +374,15 @@ where self.len() == 0 } + /// The sum of the cached values' weights. The same as [`Cache::len`] for + /// a cache built with [`Cache::new`]. + pub fn size(&self) -> usize { + self.lock().total_weight + } + + /// The bound on [`Cache::size`]. pub fn max_size(&self) -> usize { - self.lock().entries.cap().get() + self.lock().max_weight } pub fn name(&self) -> &'static str { @@ -400,17 +460,17 @@ where 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 + let weight = (self.weigher)(value); + let mut displaced = inner.insert(key, value.clone(), weight, self.clock.now()); + // `Cache::insert` may have filled the key while this load + // was in flight. What it put there is no fresher than + // this, so replacing it is not an invalidation. + displaced.evicted.extend(displaced.old.take()); + self.evicted_size + .fetch_add(displaced.evicted.len() as u64, Ordering::Relaxed); + displaced.evicted } - _ => None, + _ => Vec::new(), }; (waiters, displaced) @@ -445,15 +505,61 @@ where 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 { + /// Insert an entry of `weight` read at `now`, evicting cold entries to + /// make room. Returns whatever it displaced for the caller to drop. + /// + /// A value heavier than the whole cache is not inserted, and any entry + /// under its key is removed, so that a key never maps to a value older + /// than the one its caller just loaded. + fn insert(&mut self, key: K, value: V, weight: usize, now: SystemTime) -> Displaced { + let weight = weight.max(1); + let old = self.remove(&key); + if weight > self.max_weight { + return Displaced { + old, + evicted: Vec::new(), + }; + } + + let evicted = self.evict_to_fit(weight); + self.total_weight += weight; let entry = Entry { value, last_read: now, + weight, }; + // The key was removed above, so nothing is pushed out here. + self.entries.push(key, entry); - self.entries.push(key, entry).map(|(_, entry)| entry.value) + Displaced { old, evicted } + } + + /// Remove the entry under `key`, returning its value for the caller to + /// drop. + fn remove(&mut self, key: &K) -> Option { + let entry = self.entries.pop(key)?; + self.total_weight -= entry.weight; + Some(entry.value) + } + + /// Pop the coldest entry, returning its value for the caller to drop. + fn pop_lru(&mut self) -> Option { + let (_, entry) = self.entries.pop_lru()?; + self.total_weight -= entry.weight; + Some(entry.value) + } + + /// Pop cold entries until `extra` more weight fits under the bound. + /// Returns the values for the caller to drop. + fn evict_to_fit(&mut self, extra: usize) -> Vec { + let mut evicted = Vec::new(); + while self.total_weight + extra > self.max_weight { + match self.pop_lru() { + Some(value) => evicted.push(value), + None => break, + } + } + evicted } /// Pop every entry not read since `cutoff`, returning the values for the @@ -468,14 +574,22 @@ where .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); + let value = self.pop_lru().expect("peek_lru just succeeded"); + expired.push(value); } expired } } +/// What an insert pushed out of the cache: the value that was under the key, +/// if any, and the cold entries evicted to make room. Both are the caller's +/// to drop, with the lock released. +struct Displaced { + old: Option, + evicted: Vec, +} + /// 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. @@ -492,7 +606,8 @@ where /// 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); + let weight = (self.weigher)(&value); + let displaced = self.lock().insert(key, value, weight, now); drop(displaced); } } @@ -562,10 +677,10 @@ pub(crate) fn lock_ignoring_poison(mutex: &Mutex) -> MutexGuard<'_, T> { } /// 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) +/// never what anyone means. Treat it as one entry, as `LruCache` does for a +/// cache factor that rounds to nothing. +fn non_zero(size: usize) -> usize { + size.max(1) } #[cfg(test)] @@ -897,6 +1012,55 @@ mod tests { assert_eq!(cache.get_at(&"a", at(6)), Some(1)); } + #[test] + fn insert_fills_the_cache_directly() { + let cache = cache(2); + let loads = AtomicUsize::new(0); + + assert!(cache.insert("a", 1).is_empty()); + assert_eq!(cache.get(&"a"), Some(1)); + + // A loaded key is replaced, and the old value handed back. + assert_eq!(cache.insert("a", 2), vec![1]); + assert_eq!( + block_on(cache.get_or_load("a", counting_loader(9, &loads))).unwrap(), + 2 + ); + assert_eq!(loads.load(Ordering::SeqCst), 0); + + // Inserting past the bound evicts the coldest entry. + cache.insert("b", 3); + assert_eq!(cache.insert("c", 4), vec![2]); + assert_eq!(cache.len(), 2); + assert_eq!(cache.eviction_counts(), (1, 0, 1)); + } + + #[test] + fn weighted_values_share_the_bound() { + let cache: Arc>> = + Cache::new_weighted("weighted", 10, |value| value.len()); + + cache.insert("a", vec![0; 4]); + cache.insert("b", vec![0; 4]); + assert_eq!((cache.len(), cache.size()), (2, 8)); + + // Touch "a", so "b" is the coldest and goes when "c" needs room. + assert!(cache.get(&"a").is_some()); + assert_eq!(cache.insert("c", vec![0; 4]), vec![vec![0; 4]]); + assert!(cache.get(&"b").is_none()); + assert_eq!((cache.len(), cache.size()), (2, 8)); + + // Too heavy to cache at all, and the key it would have replaced goes. + cache.insert("a", vec![0; 11]); + assert!(cache.get(&"a").is_none()); + assert_eq!((cache.len(), cache.size()), (1, 4)); + + // The factor scales the weight bound. + cache.insert("d", vec![0; 4]); + assert_eq!(cache.set_cache_factor(0.5).len(), 1); + assert_eq!((cache.max_size(), cache.size()), (5, 4)); + } + #[test] fn cache_factor_resizes_and_evicts() { let cache = cache(4); diff --git a/rust/src/caches/python.rs b/rust/src/caches/python.rs index 260100a085..9a8d27745c 100644 --- a/rust/src/caches/python.rs +++ b/rust/src/caches/python.rs @@ -57,6 +57,27 @@ pub trait PyCacheKey: Sized { fn from_py_key(key: &Bound<'_, PyTuple>) -> Option; } +/// A key that is one string, such as a room or user ID: a one-tuple of it in +/// Python. +impl PyCacheKey for Arc { + fn from_py_key(key: &Bound<'_, PyTuple>) -> Option { + let (key,): (String,) = key.extract().ok()?; + Some(Arc::from(key)) + } +} + +/// A key that is one integer, such as a state group ID. Over the invalidation +/// stream it arrives as a string. +impl PyCacheKey for i64 { + fn from_py_key(key: &Bound<'_, PyTuple>) -> Option { + if let Ok((key,)) = key.extract::<(i64,)>() { + return Some(key); + } + let (key,): (String,) = key.extract().ok()?; + key.parse().ok() + } +} + /// A [`Cache`] with its key and value types erased, so that one `#[pyclass]` /// can serve every cache and the [`registry`](super::registry) can hold them /// all. @@ -73,7 +94,10 @@ pub(crate) trait CacheHandle: Send + Sync { /// What `register_cache` reads off a cache at each scrape. pub(crate) struct CacheStats { - pub len: usize, + /// The sum of the entries' weights, which is the entry count for a cache + /// that weighs every value as one. Reported as the cache's length, as a + /// Python `LruCache` with a `size_callback` reports its size. + pub size: usize, pub max_size: usize, pub hits: u64, pub misses: u64, @@ -95,7 +119,7 @@ where let (hits, misses) = self.hit_counts(); let (evicted_size, evicted_time, evicted_invalidation) = self.eviction_counts(); CacheStats { - len: self.len(), + size: self.size(), max_size: self.max_size(), hits, misses, @@ -176,8 +200,10 @@ impl RustCache { self.handle.invalidate_everything(); } + /// The cache's size: its entry count, or for a weighted cache the sum of + /// its entries' weights, as for a Python cache with a `size_callback`. fn __len__(&self) -> usize { - self.handle.stats().len + self.handle.stats().size } #[getter] diff --git a/rust/src/caches/registry.rs b/rust/src/caches/registry.rs index 3ff5101c91..d62dc07771 100644 --- a/rust/src/caches/registry.rs +++ b/rust/src/caches/registry.rs @@ -40,7 +40,7 @@ use std::{ use crate::caches::{ lock_ignoring_poison, python::{CacheHandle, PyCacheKey}, - Cache, + Cache, Weigher, }; /// Every cache built for one homeserver, weakly. @@ -80,8 +80,30 @@ impl CacheRegistry { K: PyCacheKey + Eq + Hash + Clone + Send + Sync + 'static, V: Clone + Send + Sync + 'static, { - let cache = Cache::new(name, max_size); + self.register(Cache::new(name, max_size)) + } + /// Build a cache whose values' weights, as `weigher` measures them, sum + /// to at most `max_weight`, enrolled for this homeserver's age-based + /// expiry. See [`Cache::new_weighted`]. + pub fn new_weighted_cache( + &self, + name: &'static str, + max_weight: usize, + weigher: Weigher, + ) -> Arc> + where + K: PyCacheKey + Eq + Hash + Clone + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, + { + self.register(Cache::new_weighted(name, max_weight, weigher)) + } + + fn register(&self, cache: Arc>) -> Arc> + where + K: PyCacheKey + Eq + Hash + Clone + Send + Sync + 'static, + V: Clone + Send + Sync + 'static, + { // Via a second `Arc` to the same allocation, because the unsizing // coercion to `dyn CacheHandle` has to happen before we downgrade. let handle: Arc = cache.clone();