diff --git a/rust/src/state/conflict_cache.rs b/rust/src/state/lru.rs similarity index 52% rename from rust/src/state/conflict_cache.rs rename to rust/src/state/lru.rs index 2a4834a652..b61f515049 100644 --- a/rust/src/state/conflict_cache.rs +++ b/rust/src/state/lru.rs @@ -12,44 +12,69 @@ * . */ -//! The cache of resolved conflicted sets, the Rust counterpart of -//! `StateResolutionHandler._conflict_resolution_cache`. +//! The caches on the Rust state resolution handler. //! -//! Keyed on the digest from `_conflict_cache_key` (see -//! [`super::v2::resolve`]). Bounded by the total number of state entries -//! across the cached maps, evicting the least recently used entry first. -//! Entries also expire after a period without use. +//! [`SizedLru`] is a bounded, least recently used cache whose bound is the +//! sum of the sizes its caller assigns to the entries. Entries also expire +//! after a period without use, so a room that goes quiet does not pin its +//! events in memory for ever. +//! +//! - [`ConflictCache`]: resolved conflicted sets, the Rust counterpart of +//! `StateResolutionHandler._conflict_resolution_cache`. Keyed on the digest +//! from `_conflict_cache_key` (see [`super::v2::resolve`]), sized by the +//! number of state entries in each cached map. +//! - [`EventCache`]: events loaded from the database for state resolution, +//! keyed on event ID and sized by count. The state resolution reads events +//! as they are (unredacted, rejected ones included), so a cached event only +//! goes stale when its rejection status or internal metadata changes. +//! `EventsWorkerStore` invalidates it alongside its own event cache, by +//! event ID and on a room purge, so the two stay in step. use std::{ collections::{BTreeMap, HashMap}, + hash::Hash, sync::{Arc, Mutex}, time::{Duration, Instant}, }; -use super::StateMap; +use super::{StateMap, StateResEvent}; -/// A cache key: the sha256 digest of the conflicted set and base state. +/// A conflict cache key: the sha256 digest of the conflicted set and base state. pub type CacheKey = [u8; 32]; -struct Entry { - value: Arc, +pub type ConflictCache = SizedLru>; + +pub type EventCache = SizedLru, StateResEvent>; + +struct Entry { + value: V, /// Position in the recency order. seq: u64, last_used: Instant, size: usize, } -#[derive(Default)] -struct Inner { - entries: HashMap, +struct Inner { + entries: HashMap>, /// Recency order: sequence number to key, oldest first. - order: BTreeMap, + order: BTreeMap, next_seq: u64, total_size: usize, } -impl Inner { - fn remove(&mut self, key: &CacheKey) { +impl Default for Inner { + fn default() -> Self { + Inner { + entries: HashMap::new(), + order: BTreeMap::new(), + next_seq: 0, + total_size: 0, + } + } +} + +impl Inner { + fn remove(&mut self, key: &K) { if let Some(entry) = self.entries.remove(key) { self.order.remove(&entry.seq); self.total_size -= entry.size; @@ -63,24 +88,26 @@ impl Inner { } } -pub struct ConflictCache { - inner: Mutex, +/// A bounded cache evicting the least recently used entry first. See the +/// module docs. +pub struct SizedLru { + inner: Mutex>, max_size: usize, expiry: Duration, } -impl ConflictCache { - /// A cache holding at most `max_size` state entries in total, dropping +impl SizedLru { + /// A cache whose entries' sizes sum to at most `max_size`, dropping /// entries unused for `expiry`. pub fn new(max_size: usize, expiry: Duration) -> Self { - ConflictCache { - inner: Mutex::default(), + SizedLru { + inner: Mutex::new(Inner::default()), max_size, expiry, } } - pub fn get(&self, key: &CacheKey) -> Option> { + pub fn get(&self, key: &K) -> Option { let mut inner = self.inner.lock().expect("not poisoned"); let now = Instant::now(); let expired = now.duration_since(inner.entries.get(key)?.last_used) > self.expiry; @@ -93,15 +120,16 @@ impl ConflictCache { let entry = inner.entries.get_mut(key).expect("checked above"); let old_seq = std::mem::replace(&mut entry.seq, seq); entry.last_used = now; - let value = Arc::clone(&entry.value); + let value = entry.value.clone(); inner.order.remove(&old_seq); - inner.order.insert(seq, *key); + inner.order.insert(seq, key.clone()); Some(value) } - pub fn insert(&self, key: CacheKey, value: Arc) { - // As `_state_map_size`, but never 0. - let size = value.len().max(1); + /// Insert `value` under `key`, counting `size` (at least 1) towards the + /// bound. A value bigger than the whole cache is not cached. + pub fn insert(&self, key: K, value: V, size: usize) { + let size = size.max(1); if size > self.max_size { return; } @@ -119,7 +147,7 @@ impl ConflictCache { } let seq = inner.bump(); - inner.order.insert(seq, key); + inner.order.insert(seq, key.clone()); inner.total_size += size; inner.entries.insert( key, @@ -132,7 +160,12 @@ impl ConflictCache { ); } - /// The number of cached resolutions. + /// Drop the entry under `key`, if any. + pub fn remove(&self, key: &K) { + self.inner.lock().expect("not poisoned").remove(key); + } + + /// The number of cached entries. pub fn len(&self) -> usize { self.inner.lock().expect("not poisoned").entries.len() } @@ -146,6 +179,15 @@ impl ConflictCache { } } +impl ConflictCache { + /// Cache a resolved conflicted set, sized by its number of state entries + /// as `_state_map_size` does. + pub fn insert_state(&self, key: CacheKey, value: Arc) { + let size = value.len(); + self.insert(key, value, size); + } +} + #[cfg(test)] mod tests { use super::*; @@ -162,17 +204,17 @@ mod tests { #[test] fn evicts_least_recently_used_by_size() { let cache = ConflictCache::new(10, Duration::from_secs(60)); - cache.insert([1; 32], map(4)); - cache.insert([2; 32], map(4)); + cache.insert_state([1; 32], map(4)); + cache.insert_state([2; 32], map(4)); // Use 1, so 2 is the oldest. assert!(cache.get(&[1; 32]).is_some()); - cache.insert([3; 32], map(4)); + cache.insert_state([3; 32], map(4)); assert!(cache.get(&[1; 32]).is_some()); assert!(cache.get(&[2; 32]).is_none()); assert!(cache.get(&[3; 32]).is_some()); // Too big to cache at all. - cache.insert([4; 32], map(11)); + cache.insert_state([4; 32], map(11)); assert!(cache.get(&[4; 32]).is_none()); assert_eq!(cache.len(), 2); } @@ -180,9 +222,20 @@ mod tests { #[test] fn expires_unused_entries() { let cache = ConflictCache::new(10, Duration::ZERO); - cache.insert([1; 32], map(1)); + cache.insert_state([1; 32], map(1)); std::thread::sleep(Duration::from_millis(2)); assert!(cache.get(&[1; 32]).is_none()); assert!(cache.is_empty()); } + + #[test] + fn counts_entries_by_one() { + let cache: SizedLru = SizedLru::new(2, Duration::from_secs(60)); + cache.insert(1, 1, 1); + cache.insert(2, 2, 1); + cache.insert(3, 3, 1); + assert!(cache.get(&1).is_none()); + assert_eq!(cache.get(&2), Some(2)); + assert_eq!(cache.get(&3), Some(3)); + } } diff --git a/rust/src/state/mod.rs b/rust/src/state/mod.rs index 62333dda87..2c3ea2690e 100644 --- a/rust/src/state/mod.rs +++ b/rust/src/state/mod.rs @@ -19,7 +19,7 @@ //! with a Postgres implementation and an in-memory one for tests. //! - [`auth_chain`]: the chain cover auth chain difference. //! - [`v2`]: the algorithm. -//! - [`conflict_cache`]: the cache of resolved conflicted sets. +//! - [`lru`]: the conflict and event caches. //! //! [`StateResolutionHandler`] is the Python entry point, used by //! `synapse.state.StateResolutionHandler` when `rust_state_res` is enabled. @@ -51,11 +51,11 @@ use crate::room_versions::{RoomVersion, StateResolutionVersions}; use crate::runtime::RustRuntime; pub mod auth_chain; -pub mod conflict_cache; +pub mod lru; pub mod store; pub mod v2; -use conflict_cache::ConflictCache; +use lru::{ConflictCache, EventCache}; use store::{NoChainCoverIndex, PostgresStateResStore}; create_exception!( @@ -318,6 +318,13 @@ const CONFLICT_CACHE_MAX_SIZE: usize = 100_000; /// `EVICTION_TIMEOUT_SECONDS`. const CONFLICT_CACHE_EXPIRY: Duration = Duration::from_secs(60 * 60); +/// The bound on the event cache, in events. +/// TODO: size this from `event_cache_size`, as the Python event cache is. +const EVENT_CACHE_MAX_SIZE: usize = 10_000; + +/// How long an unused event cache entry lives. +const EVENT_CACHE_EXPIRY: Duration = Duration::from_secs(60 * 60); + /// Timings and counts for one resolution, for metrics. Times are in seconds. #[pyclass(frozen, get_all, skip_from_py_object)] #[derive(Debug, Default, Clone)] @@ -337,6 +344,12 @@ pub struct StateResStats { pub db_rows: u64, /// Events loaded from the database. pub db_events: u64, + /// Events found in the handler's event cache. + pub event_cache_hits: u64, + /// Time converting the state sets from Python, under the GIL. + pub convert_in_time: f64, + /// Time converting the resolved state to Python, under the GIL. + pub convert_out_time: f64, /// Events in the full conflicted set. pub conflicted_events: usize, /// The resolution of the conflicted set came from the cache. @@ -362,8 +375,11 @@ impl<'py> IntoPyObject<'py> for ResolveOutput { type Error = PyErr; fn into_pyobject(self, py: Python<'py>) -> PyResult { + let start = Instant::now(); let state = state_map_to_py(py, &self.state)?; - (state, Bound::new(py, self.stats)?).into_pyobject(py) + let mut stats = self.stats; + stats.convert_out_time = start.elapsed().as_secs_f64(); + (state, Bound::new(py, stats)?).into_pyobject(py) } } @@ -391,6 +407,7 @@ pub struct StateResolutionHandler { db_pool: Py, runtime: RustRuntime, cache: Arc, + event_cache: Arc, } impl StateResolutionHandler { @@ -402,6 +419,7 @@ impl StateResolutionHandler { CONFLICT_CACHE_MAX_SIZE, CONFLICT_CACHE_EXPIRY, )), + event_cache: Arc::new(EventCache::new(EVENT_CACHE_MAX_SIZE, EVENT_CACHE_EXPIRY)), } } @@ -430,6 +448,7 @@ impl StateResolutionHandler { } let pool = pool.cast::()?.get().pool.clone(); + let convert_start = Instant::now(); let mut interner = Interner::default(); let state_sets = state_sets .iter() @@ -439,6 +458,7 @@ impl StateResolutionHandler { // Without conflicts there is nothing to read or resolve, so answer // straight away rather than on the runtime. let (unconflicted, conflicted) = v2::separate(&state_sets); + let convert_in_time = convert_start.elapsed().as_secs_f64(); if conflicted.is_empty() { return Ok(Started::Done(ResolveOutput { state: unconflicted, @@ -450,10 +470,11 @@ impl StateResolutionHandler { let handle = self.runtime.tokio_handle()?; let cache = Arc::clone(&self.cache); + let event_cache = Arc::clone(&self.event_cache); Ok(Started::Pending(Box::pin(async move { let start = Instant::now(); - let store = PostgresStateResStore::begin(&pool, handle) + let store = PostgresStateResStore::begin(&pool, handle, Some(event_cache)) .await .map_err(resolve_error_to_py)?; let prepared = v2::prepare( @@ -490,6 +511,9 @@ impl StateResolutionHandler { db_queries: db.queries, db_rows: db.rows, db_events: db.events_fetched, + event_cache_hits: db.events_cached, + convert_in_time, + convert_out_time: 0.0, conflicted_events: core.conflicted_events, cache_hit: core.cache_hit, }; @@ -560,6 +584,29 @@ impl StateResolutionHandler { }; output.into_pyobject(py) } + + /// Empty the conflict cache, keeping the event cache. For benchmarks and + /// tests that want to time the algorithm again on loaded events. + fn clear_conflict_cache_for_test(&self) { + self.cache.clear(); + } + + /// The number of events in the event cache. + fn event_cache_len(&self) -> usize { + self.event_cache.len() + } + + /// Drop an event from the event cache. Called whenever the Python event + /// cache drops it, so a change to its rejection status or internal + /// metadata reaches the next resolution. + fn invalidate_event(&self, event_id: &str) { + self.event_cache.remove(&Arc::from(event_id)); + } + + /// Empty the event cache. Called when a room is purged. + fn clear_event_cache(&self) { + self.event_cache.clear(); + } } /// Called when registering modules with python. diff --git a/rust/src/state/store.rs b/rust/src/state/store.rs index 3b253521e6..83e8f6c13b 100644 --- a/rust/src/state/store.rs +++ b/rust/src/state/store.rs @@ -26,7 +26,7 @@ use std::{ use anyhow::{bail, Context}; use async_trait::async_trait; -use log::{error, warn}; +use log::{debug, error, warn}; use pyo3::{ prelude::*, types::{PyDict, PyModule, PySet}, @@ -34,7 +34,7 @@ use pyo3::{ use tokio::runtime::Handle; use tokio_postgres::{types::ToSql, Row}; -use super::{auth_chain, EventMap, StateResEvent}; +use super::{auth_chain, lru::EventCache, EventMap, StateResEvent}; use crate::database::postgres::pool::{ConnectionPool, PooledConnection, PyConnectionPool}; use crate::events::{ constants::event_type::M_ROOM_MEMBER, internal_metadata::EventInternalMetadata, Event, @@ -93,8 +93,10 @@ pub struct DbStats { pub queries: u64, /// Rows returned. pub rows: u64, - /// Events loaded by [`StateResStore::get_events`]. + /// Events loaded from the database by [`StateResStore::get_events`]. pub events_fetched: u64, + /// Events [`StateResStore::get_events`] found in the event cache. + pub events_cached: u64, /// Time spent waiting on statements. pub db_time: Duration, /// Time spent building events from rows: parsing the JSON and computing @@ -115,6 +117,8 @@ pub struct DbStats { pub struct PostgresStateResStore { conn: Option, handle: Handle, + /// Events loaded earlier, consulted before the database. + event_cache: Option>, db_stats: Mutex, } @@ -133,7 +137,13 @@ const LOAD_EVENTS_SQL: &str = " impl PostgresStateResStore { /// Check a connection out of `pool` and open the read-only transaction. - pub async fn begin(pool: &ConnectionPool, handle: Handle) -> anyhow::Result { + /// Events are read from `event_cache` first, and loaded events are added + /// to it. + pub async fn begin( + pool: &ConnectionPool, + handle: Handle, + event_cache: Option>, + ) -> anyhow::Result { let conn = pool .get() .await @@ -141,6 +151,7 @@ impl PostgresStateResStore { let store = PostgresStateResStore { conn: Some(conn), handle, + event_cache, db_stats: Mutex::default(), }; // If this fails, dropping the store sends a ROLLBACK, which is @@ -197,6 +208,14 @@ impl PostgresStateResStore { .await .with_context(|| format!("running {}", sql.trim()))?; self.record(start, rows.len()); + if log::log_enabled!(log::Level::Debug) { + let head: String = sql.split_whitespace().take(6).collect::>().join(" "); + debug!( + "State res query [{head}] returned {} rows in {:?}", + rows.len(), + start.elapsed() + ); + } Ok(rows) } } @@ -220,18 +239,43 @@ impl Drop for PostgresStateResStore { impl StateResStore for PostgresStateResStore { async fn get_events(&self, ids: &[Arc]) -> anyhow::Result { let mut events = EventMap::with_capacity(ids.len()); - for batch in ids.chunks(EVENT_BATCH_SIZE) { + + let missing: Vec> = match &self.event_cache { + Some(cache) => { + let mut missing = Vec::new(); + for id in ids { + match cache.get(id) { + Some(event) => { + events.insert(Arc::clone(id), event); + } + None => missing.push(Arc::clone(id)), + } + } + missing + } + None => ids.to_vec(), + }; + let cached = events.len() as u64; + + for batch in missing.chunks(EVENT_BATCH_SIZE) { let batch: Vec<&str> = batch.iter().map(|id| &**id).collect(); let rows = self.query(LOAD_EVENTS_SQL, &[&batch]).await?; let start = Instant::now(); for row in rows { if let Some((event_id, event)) = event_from_row(&row)? { - events.insert(event_id, StateResEvent::new(event)?); + let event = StateResEvent::new(event)?; + if let Some(cache) = &self.event_cache { + cache.insert(Arc::clone(&event_id), event.clone(), 1); + } + events.insert(event_id, event); } } self.db_stats.lock().expect("not poisoned").parse_time += start.elapsed(); } - self.db_stats.lock().expect("not poisoned").events_fetched += events.len() as u64; + + let mut stats = self.db_stats.lock().expect("not poisoned"); + stats.events_cached += cached; + stats.events_fetched += events.len() as u64 - cached; Ok(events) } @@ -455,7 +499,7 @@ fn get_auth_chain_difference_for_test<'py>( let difference = py .detach(|| { handle.block_on(async { - let store = PostgresStateResStore::begin(&pool, handle.clone()).await?; + let store = PostgresStateResStore::begin(&pool, handle.clone(), None).await?; let difference = auth_chain::get_auth_chain_difference_using_cover_index( &store, &room_id, @@ -494,7 +538,7 @@ fn load_events_for_test<'py>( let events = py .detach(|| { handle.block_on(async { - let store = PostgresStateResStore::begin(&pool, handle.clone()).await?; + let store = PostgresStateResStore::begin(&pool, handle.clone(), None).await?; let events = store.get_events(&ids).await?; store.finish().await?; anyhow::Ok(events) diff --git a/rust/src/state/v2.rs b/rust/src/state/v2.rs index a6e55d254d..38304bb26e 100644 --- a/rust/src/state/v2.rs +++ b/rust/src/state/v2.rs @@ -46,8 +46,8 @@ use serde_json::Value; use sha2::{Digest, Sha256}; use super::{ - conflict_cache::ConflictCache, lookup, + lru::ConflictCache, store::{StateDifference, StateResStore}, EventMap, StateKey, StateMap, StateResEvent, }; @@ -178,7 +178,7 @@ impl PendingResolution { let resolved = Arc::new(resolved); if let (Some(cache), Some(key)) = (conflict_cache, cache_key) { - cache.insert(key, Arc::clone(&resolved)); + cache.insert_state(key, Arc::clone(&resolved)); } Ok(Resolution { diff --git a/synapse/storage/databases/main/events_worker.py b/synapse/storage/databases/main/events_worker.py index a3ce4a883b..a3d0c63f26 100644 --- a/synapse/storage/databases/main/events_worker.py +++ b/synapse/storage/databases/main/events_worker.py @@ -1017,7 +1017,8 @@ class EventsWorkerStore(SQLBaseStore): def _invalidate_local_get_event_cache(self, event_id: str) -> None: """ - Invalidates an event in local in-memory get event caches. + Invalidates an event in local in-memory get event caches, including + the Rust state resolution's. Arguments: event_id: the event ID to invalidate @@ -1027,6 +1028,10 @@ class EventsWorkerStore(SQLBaseStore): self._event_ref.pop(event_id, None) self._current_event_fetches.pop(event_id, None) + rust_state_res = self.hs.get_rust_handlers().state_resolution + if rust_state_res is not None: + rust_state_res.invalidate_event(event_id) + def _invalidate_local_get_event_cache_room_id(self, room_id: str) -> None: """Clears the in-memory get event caches for a room. @@ -1036,6 +1041,10 @@ class EventsWorkerStore(SQLBaseStore): self._event_ref.clear() self._current_event_fetches.clear() + rust_state_res = self.hs.get_rust_handlers().state_resolution + if rust_state_res is not None: + rust_state_res.clear_event_cache() + def _invalidate_async_get_event_cache_room_id(self, room_id: str) -> None: """ Clears the async get_event cache for a room. Currently a no-op until diff --git a/synapse/synapse_rust/state.pyi b/synapse/synapse_rust/state.pyi index 273b86c858..66bc9cc305 100644 --- a/synapse/synapse_rust/state.pyi +++ b/synapse/synapse_rust/state.pyi @@ -35,6 +35,7 @@ class StateResStats: db_queries: int db_rows: int db_events: int + event_cache_hits: int conflicted_events: int cache_hit: bool @@ -70,6 +71,18 @@ class StateResolutionHandler: """As `resolve`, but blocks and returns the result. For tests, where the fake reactor does not wait for a deferred fired from another thread.""" + def clear_conflict_cache_for_test(self) -> None: + """Empty the conflict cache, keeping the event cache.""" + + def event_cache_len(self) -> int: + """The number of events in the event cache.""" + + def invalidate_event(self, event_id: str) -> None: + """Drop an event from the event cache.""" + + def clear_event_cache(self) -> None: + """Empty the event cache.""" + def get_auth_chain_difference_for_test( pool: ConnectionPool, room_id: str, diff --git a/tests/state/test_rust_handler.py b/tests/state/test_rust_handler.py index 20d6b4ee5e..dc1d458119 100644 --- a/tests/state/test_rust_handler.py +++ b/tests/state/test_rust_handler.py @@ -160,3 +160,51 @@ class RustStateResHandlerTestCase(unittest.HomeserverTestCase): any("Falling back to Python state res" in line for line in logs.output) ) self.assertIn(self._current_topic(room_id), topics) + + def test_event_cache(self) -> None: + """A second resolution reads its events from the handler's event + cache, which the store's event cache invalidation empties.""" + rust = self.hs.get_rust_handlers().state_resolution + assert rust is not None + + room_id = self.helper.create_room_as( + self.user_id, room_version=RoomVersions.V10.identifier, tok=self.tok + ) + topic_a, topic_b = self._fork(room_id) + groups = self.get_success( + self.store._get_state_group_for_events([topic_a, topic_b]) + ) + state_controller = self.hs.get_storage_controllers().state + state_sets = [ + self.get_success(state_controller.get_state_ids_for_group(group)) + for group in groups.values() + ] + + # Persisting the fork may already have resolved it here, depending on + # the test's `rust_state_res` mode. + rust.clear_event_cache() + rust.clear_conflict_cache_for_test() + + first, first_stats = self.get_success( + rust.resolve(room_id, RoomVersions.V10.identifier, state_sets) + ) + self.assertGreater(first_stats.db_events, 0) + self.assertEqual(first_stats.event_cache_hits, 0) + cached = rust.event_cache_len() + self.assertEqual(cached, first_stats.db_events) + + # Clear the conflict cache, so the resolution runs again rather than + # being answered from it. + rust.clear_conflict_cache_for_test() + second, second_stats = self.get_success( + rust.resolve(room_id, RoomVersions.V10.identifier, state_sets) + ) + self.assertEqual(second, first) + self.assertEqual(second_stats.db_events, 0) + self.assertEqual(second_stats.event_cache_hits, first_stats.db_events) + + self.store._invalidate_local_get_event_cache(topic_a) + self.assertEqual(rust.event_cache_len(), cached - 1) + + self.store._invalidate_local_get_event_cache_room_id(room_id) + self.assertEqual(rust.event_cache_len(), 0)