diff --git a/changelog.d/rust-event-fetch.misc b/changelog.d/rust-event-fetch.misc new file mode 100644 index 0000000000..bf5346b3b0 --- /dev/null +++ b/changelog.d/rust-event-fetch.misc @@ -0,0 +1 @@ +Load events from the database in Rust, for Python callers as well as the Rust readers, when the main database uses the native Postgres driver. diff --git a/rust/src/storage/events.rs b/rust/src/storage/events.rs index 41683cadb2..ee8cacb278 100644 --- a/rust/src/storage/events.rs +++ b/rust/src/storage/events.rs @@ -28,16 +28,20 @@ use std::{ collections::{HashMap, HashSet}, - sync::{Arc, OnceLock}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, Mutex, OnceLock, + }, time::{Duration, Instant}, }; -use anyhow::{bail, Context}; +use anyhow::Context; use async_trait::async_trait; use log::{debug, error, info, warn}; use pyo3::{Py, PyResult, Python}; use tokio_postgres::{types::ToSql, Row}; +use crate::caches::lock_ignoring_poison; use crate::event_auth::get_domain_from_id; use crate::events::{ constants::event_type::{M_ROOM_CREATE, M_ROOM_MEMBER, M_ROOM_REDACTION}, @@ -152,6 +156,39 @@ fn auth_event_ids_of(event: &Event) -> anyhow::Result]>> { Ok(ids.into_iter().map(Arc::from).collect()) } +/// Why an event in the database cannot be used. Each kind is a Python +/// error of the same name, which some callers catch to skip a broken room. +#[derive(Debug)] +pub enum UnusableEvent { + /// `InvalidEventError`. + Invalid(String), + /// `DatabaseCorruptionError`: the event ID computed from the stored + /// event doesn't match the stored ID. + Corrupt { + room_id: String, + persisted_event_id: String, + computed_event_id: String, + }, +} + +impl std::fmt::Display for UnusableEvent { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + UnusableEvent::Invalid(message) => f.write_str(message), + UnusableEvent::Corrupt { + room_id, + persisted_event_id, + computed_event_id, + } => write!( + f, + "Database corruption: event {persisted_event_id} in room {room_id} has event ID {computed_event_id}" + ), + } + } +} + +impl std::error::Error for UnusableEvent {} + /// What to do with a redacted event. `EventRedactBehaviour`. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RedactBehaviour { @@ -187,6 +224,23 @@ pub struct EventLoadStats { pub parse_time: Duration, } +impl EventLoadStats { + fn add(&mut self, other: &EventLoadStats) { + self.fetched += other.fetched; + self.cached += other.cached; + self.parse_time += other.parse_time; + } +} + +/// `ids` without repeats, in order. +fn dedup(ids: &[Arc]) -> Vec> { + let mut seen = HashSet::with_capacity(ids.len()); + ids.iter() + .filter(|id| seen.insert(Arc::clone(id))) + .cloned() + .collect() +} + /// Loads events through `querier`, consulting `cache` first and filling it /// with what it loads. pub struct EventLoader<'a> { @@ -239,38 +293,69 @@ impl<'a> EventLoader<'a> { /// database. Events that are not found are left out. The result can /// hold events that were not asked for: the redaction events read to /// check an unconfirmed redaction. + /// + /// With a cache, an event another caller is loading is waited for + /// rather than read again, and the rest are read in one go. Those that + /// were asked for are cached, unless invalidated while being read. The + /// extra redaction events are not, as nothing guards them against an + /// invalidation during the read. pub async fn load( &self, ids: &[Arc], stats: &mut EventLoadStats, ) -> anyhow::Result, CachedEvent>> { - let mut events = HashMap::with_capacity(ids.len()); - let mut missing: Vec> = Vec::new(); - for id in ids { - if events.contains_key(id) || missing.contains(id) { - continue; - } - match self.cache.and_then(|cache| cache.get_or_count_miss(id)) { - Some(event) => { - events.insert(Arc::clone(id), event); - } - None => missing.push(Arc::clone(id)), - } - } - stats.cached += events.len() as u64; - - if !missing.is_empty() { - let fetched = self.fetch_from_db(missing, stats).await?; + let Some(cache) = self.cache else { + let fetched = self.fetch_from_db(dedup(ids), stats).await?; stats.fetched += fetched.len() as u64; - for (event_id, event) in fetched { - if let Some(cache) = self.cache { - drop(cache.insert(Arc::clone(&event_id), event.clone())); + return Ok(fetched); + }; + + // `get_or_load_many` takes a closure it may call more than once, so + // the closure cannot hold `stats` itself. + let load_stats = Mutex::new(EventLoadStats::default()); + let loaded_requested = AtomicU64::new(0); + let events = cache + .get_or_load_many(ids, |missing| { + let load_stats = &load_stats; + let loaded_requested = &loaded_requested; + async move { + let mut batch_stats = EventLoadStats::default(); + let fetched = self + .fetch_from_db(missing.clone(), &mut batch_stats) + .await?; + let requested = missing.iter().filter(|id| fetched.contains_key(*id)); + loaded_requested.fetch_add(requested.count() as u64, Ordering::Relaxed); + batch_stats.fetched += fetched.len() as u64; + lock_ignoring_poison(load_stats).add(&batch_stats); + Ok(fetched) } - events.insert(event_id, event); + }) + .await + .map_err(anyhow::Error::new)?; + + let load_stats = lock_ignoring_poison(&load_stats); + stats.add(&load_stats); + let requested_found = ids.iter().filter(|id| events.contains_key(*id)).count() as u64; + stats.cached += requested_found.saturating_sub(loaded_requested.into_inner()); + Ok(events) + } + + /// [`EventLoader::load`] reading every event from the database, as + /// `_get_events_from_db`: for callers that must not trust the cache. + /// What it reads replaces what was cached. + pub async fn reload( + &self, + ids: &[Arc], + stats: &mut EventLoadStats, + ) -> anyhow::Result, CachedEvent>> { + if let Some(cache) = self.cache { + // Invalidating also stops loads already in flight from caching + // what they read before ours. + for id in ids { + drop(cache.invalidate(id)); } } - - Ok(events) + self.load(ids, stats).await } /// `get_events_as_list`: the events in the order of `ids`, each as @@ -777,17 +862,22 @@ fn event_from_row(row: &Row) -> anyhow::Result, Event)>> { }; if room_version_id.is_none() && event.type_() != M_ROOM_MEMBER { - bail!("Room for event {event_id} is unknown"); + return Err(UnusableEvent::Invalid(format!( + "Room {} for event {event_id} is unknown", + event.room_id() + )) + .into()); } // If the content was modified in the database, the computed event ID // will not match. if event.event_id() != event_id { - bail!( - "Database corruption: event {event_id} in room {} has event ID {}", - event.room_id(), - event.event_id() - ); + return Err(UnusableEvent::Corrupt { + room_id: event.room_id().to_owned(), + persisted_event_id: event_id, + computed_event_id: event.event_id().to_owned(), + } + .into()); } Ok(Some((Arc::from(event_id), event))) diff --git a/rust/src/storage/mod.rs b/rust/src/storage/mod.rs index 7756685db8..13e2a5c0e5 100644 --- a/rust/src/storage/mod.rs +++ b/rust/src/storage/mod.rs @@ -17,6 +17,7 @@ pub mod caches; pub mod db; pub mod events; pub mod federation_store; +pub mod pool_querier; pub mod python; pub mod store; pub mod sync_store; diff --git a/rust/src/storage/pool_querier.rs b/rust/src/storage/pool_querier.rs new file mode 100644 index 0000000000..14554c33ee --- /dev/null +++ b/rust/src/storage/pool_querier.rs @@ -0,0 +1,59 @@ +/* + * 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: + * . + */ + +//! A [`Querier`] over the native driver's pool, outside any transaction. + +use anyhow::Context; +use async_trait::async_trait; +use tokio_postgres::{types::ToSql, Row}; + +use crate::database::postgres::pool::{ConnectionPool, PooledConnection}; +use crate::storage::events::Querier; + +/// Runs each statement on its own, in autocommit, on one connection checked +/// out of the pool, which goes back when this is dropped. +/// +/// The connection is checked out up front, before the caller claims any +/// event loads in the event cache. A load claimed first and then left +/// waiting for a connection could stall until the checkout times out: the +/// Rust readers that join it hold pooled connections while they wait, and +/// enough of them leave none for the load. +pub struct PoolQuerier { + conn: PooledConnection, +} + +impl PoolQuerier { + pub async fn checkout(pool: &ConnectionPool) -> anyhow::Result { + let conn = pool + .get() + .await + .map_err(|e| anyhow::anyhow!("checking out a database connection: {e}"))?; + Ok(PoolQuerier { conn }) + } +} + +#[async_trait] +impl Querier for PoolQuerier { + async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> anyhow::Result> { + let statement = self + .conn + .prepare_cached(sql, false) + .await + .with_context(|| format!("preparing {}", sql.trim()))?; + self.conn + .query(&statement, params) + .await + .with_context(|| format!("running {}", sql.trim())) + } +} diff --git a/rust/src/storage/python.rs b/rust/src/storage/python.rs index 4119504b61..3a7e6f14f3 100644 --- a/rust/src/storage/python.rs +++ b/rust/src/storage/python.rs @@ -34,33 +34,86 @@ //! `register_cache` through [`RustStore::caches`], for metrics and cache //! factors. //! +//! When the main database uses the native driver, events are loaded from +//! the database here too, for Python as for Rust: `EventsWorkerStore` calls +//! [`RustStore::load_events`] for the events its caches miss. So every event +//! load goes through the one [`EventLoader`], and loads of the same event +//! from either side coalesce in the event cache. +//! //! The caches that mirror a Python `@cached` function reach the datastore //! through [`RustStore::mirrored_caches`], which attaches each to the //! Python cache it names, so that every invalidation of the Python cache //! reaches it. -use std::sync::Arc; +use std::{collections::HashMap, future::Future, sync::Arc}; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, - pyclass, pymethods, - types::{PyAnyMethods, PyModule, PyModuleMethods, PyTuple}, - Bound, Py, PyAny, PyClassInitializer, PyResult, Python, + import_exception, intern, pyclass, pymethods, + types::{ + PyAnyMethods, PyDict, PyDictMethods, PyModule, PyModuleMethods, PyTuple, PyWeakrefMethods, + PyWeakrefReference, + }, + Bound, IntoPyObject, Py, PyAny, PyClassInitializer, PyErr, PyResult, Python, }; use crate::{ - caches::python::{PyCacheKey, RustCache}, + caches::{ + python::{PyCacheKey, RustCache}, + SharedError, + }, + database::postgres::pool::{ConnectionPool, PyConnectionPool}, deferred::{create_deferred, succeed}, events::Event, runtime::RustRuntime, storage::{ caches::{new_event_cache, EventCache, StoreCaches}, db::python_db_pool::PythonDatabasePoolWrapper, - events::CachedEvent, + events::{CachedEvent, EventLoadStats, EventLoader, UnusableEvent}, + pool_querier::PoolQuerier, store::{IsFeatureEnabledKey, PerUserExperimentalFeature, Store}, }, }; +import_exception!( + synapse.storage.databases.main.events_worker, + InvalidEventError +); +import_exception!( + synapse.storage.databases.main.events_worker, + DatabaseCorruptionError +); + +/// The Python error for a failed event load: `InvalidEventError` or +/// `DatabaseCorruptionError` for an event that cannot be used, as the +/// Python loader raises, and `RuntimeError` for anything else. +fn event_load_error_to_py(error: anyhow::Error) -> PyErr { + match unusable_event(&error) { + Some(UnusableEvent::Invalid(message)) => InvalidEventError::new_err(message.clone()), + Some(UnusableEvent::Corrupt { + room_id, + persisted_event_id, + computed_event_id, + }) => DatabaseCorruptionError::new_err(( + room_id.clone(), + persisted_event_id.clone(), + computed_event_id.clone(), + )), + None => PyRuntimeError::new_err(format!("loading events failed: {error:#}")), + } +} + +/// The [`UnusableEvent`] behind `error`, which may have reached us through +/// another caller's load in the event cache. +fn unusable_event(error: &anyhow::Error) -> Option<&UnusableEvent> { + if let Some(unusable) = error.downcast_ref::() { + return Some(unusable); + } + error + .downcast_ref::() + .and_then(|shared| unusable_event(shared.inner())) +} + /// The Rust [`Store`], as the main datastore holds it. #[pyclass(frozen, name = "RustStore")] pub struct RustStore { @@ -72,12 +125,132 @@ pub struct RustStore { caches: Vec>, /// The caches that mirror a Python cache, with the Python cache's name. mirrored_caches: Vec<(&'static str, Py)>, + runtime: RustRuntime, + /// The datastore's `DatabasePool`, when it uses the native driver. Its + /// `_db_pool._pool` is read per call: `close()` and `start()` replace + /// the `_pool`, and the replication tests swap a worker's `_db_pool` for + /// the main process's. A weak reference, as the `DatabasePool` references + /// the homeserver, which owns this: see `PythonDatabasePoolWrapper`. + native_db_pool: Option>, } impl RustStore { pub fn store(&self) -> &Arc { &self.store } + + /// The native driver's pool, or an error if the database does not use + /// that driver or the pool is closed. + fn pool(&self, py: Python<'_>) -> PyResult { + let Some(db_pool) = &self.native_db_pool else { + return Err(PyRuntimeError::new_err( + "loading events in Rust needs the native Postgres driver", + )); + }; + let db_pool = db_pool + .bind(py) + .upgrade() + .ok_or_else(|| PyRuntimeError::new_err("the database pool has been dropped"))?; + let pool = db_pool + .getattr(intern!(py, "_db_pool"))? + .getattr(intern!(py, "_pool"))?; + if pool.is_none() { + return Err(PyRuntimeError::new_err("the database pool is closed")); + } + Ok(pool.cast::()?.get().pool.clone()) + } + + fn start_load<'py>( + &self, + py: Python<'py>, + event_ids: Vec, + reload: bool, + ) -> PyResult> { + let load = self.load(py, event_ids, reload)?; + // Shared with other callers through the cache, so a cancelled caller + // must leave the load running for them. + create_deferred(py, &self.runtime, load) + } + + /// [`RustStore::start_load`], blocking until the load is done, with the + /// GIL released. + fn load_blocking( + &self, + py: Python<'_>, + event_ids: Vec, + reload: bool, + ) -> PyResult { + let load = self.load(py, event_ids, reload)?; + let handle = self.runtime.tokio_handle()?; + py.detach(|| handle.block_on(load)) + } + + /// The future that loads `event_ids`, to run on the tokio runtime. + fn load( + &self, + py: Python<'_>, + event_ids: Vec, + reload: bool, + ) -> PyResult> + Send + 'static> { + let pool = self.pool(py)?; + let store = Arc::clone(&self.store); + let ids: Vec> = event_ids.into_iter().map(Arc::from).collect(); + Ok(async move { + let mut stats = EventLoadStats::default(); + let events = async { + let querier = PoolQuerier::checkout(&pool).await?; + let loader = EventLoader::new(&querier, Some(&store.caches.events)); + if reload { + loader.reload(&ids, &mut stats).await + } else { + loader.load(&ids, &mut stats).await + } + } + .await; + events + .map(|events| LoadedEvents { + events, + fetched: stats.fetched, + }) + .map_err(event_load_error_to_py) + }) + } +} + +/// `db_pool`, a `DatabasePool`, if it runs on the native driver. +fn native_db_pool(db_pool: &Bound<'_, PyAny>) -> PyResult>> { + let py = db_pool.py(); + let connection_pool = db_pool.getattr(intern!(py, "_db_pool"))?; + let rust_connection_pool_class = py + .import("synapse.storage.rust_pool")? + .getattr("RustConnectionPool")?; + if !connection_pool.is_instance(&rust_connection_pool_class)? { + return Ok(None); + } + Ok(Some(PyWeakrefReference::new(db_pool)?.unbind())) +} + +/// What [`RustStore::load_events`] answers: for each event ID, the event +/// and its redacted form, as new Python objects sharing the cached data, +/// and how many events this load read from the database. Events read by +/// another caller's load that this one joined are not counted. +struct LoadedEvents { + events: HashMap, CachedEvent>, + fetched: u64, +} + +impl<'py> IntoPyObject<'py> for LoadedEvents { + type Target = PyTuple; + type Output = Bound<'py, PyTuple>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let dict = PyDict::new(py); + for (event_id, event) in self.events { + dict.set_item(&*event_id, (event.to_py(py)?, event.redacted_to_py(py)?))?; + } + (dict, self.fetched).into_pyobject(py) + } } #[pymethods] @@ -92,6 +265,7 @@ impl RustStore { db_pool: &Bound<'_, PyAny>, event_cache_size: usize, ) -> PyResult { + let native_db_pool = native_db_pool(db_pool)?; let db_pool = PythonDatabasePoolWrapper::new(db_pool, runtime.clone())?; let store = Arc::new(Store::new( Box::new(db_pool), @@ -141,9 +315,67 @@ impl RustStore { event_cache, caches, mirrored_caches, + runtime: runtime.clone(), + native_db_pool, }) } + /// Whether [`RustStore::load_events`] can be used: whether the database + /// uses the native driver. + #[getter] + fn loads_events(&self) -> bool { + self.native_db_pool.is_some() + } + + /// `_get_events_from_db` behind the event cache: a deferred of the events + /// by ID, each with its redacted form if it has a valid redaction, + /// rejected events included and unknown ones left out, and the number of + /// events this call read from the database. Each event comes + /// from the cache, from a load already in flight, or from the database, + /// and those read from the database are cached. + /// + /// The result can hold events that were not asked for: the redaction + /// events read to check an unconfirmed redaction. + fn load_events<'py>( + &self, + py: Python<'py>, + event_ids: Vec, + ) -> PyResult> { + self.start_load(py, event_ids, false) + } + + /// [`RustStore::load_events`] reading every event from the database, + /// for a caller that must not trust the cache. What it reads replaces + /// what was cached. + fn reload_events<'py>( + &self, + py: Python<'py>, + event_ids: Vec, + ) -> PyResult> { + self.start_load(py, event_ids, true) + } + + /// As [`RustStore::load_events`], but blocks and returns the result. + /// For tests, which run Python's database transactions synchronously + /// too: see `_BlockingRustCalls`. + fn load_events_blocking( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult { + self.load_blocking(py, event_ids, false) + } + + /// As [`RustStore::reload_events`], but blocks and returns the result. + /// For tests. + fn reload_events_blocking( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult { + self.load_blocking(py, event_ids, true) + } + /// The caches that mirror a Python `@cached` function, each with the /// name of the function, for the datastore to attach as a mirror of /// that function's cache and to register for metrics. @@ -384,7 +616,10 @@ impl RustEventCache { }; match &self.store_caches { Some(store_caches) => store_caches.invalidate_room(&room_id), - None => drop(self.events.invalidate_index(&room_id)), + None => { + drop(self.events.invalidate_index(&room_id)); + self.events.abandon_loads(); + } } } diff --git a/synapse/storage/databases/main/events_worker.py b/synapse/storage/databases/main/events_worker.py index bfab52fdcc..519f8f2778 100644 --- a/synapse/storage/databases/main/events_worker.py +++ b/synapse/storage/databases/main/events_worker.py @@ -90,7 +90,7 @@ from synapse.storage.util.id_generators import ( MultiWriterIdGenerator, ) from synapse.storage.util.sequence import build_sequence_generator -from synapse.synapse_rust.storage import RustEventCache +from synapse.synapse_rust.storage import RustEventCache, RustStore from synapse.types import JsonDict, get_domain_from_id from synapse.types.state import StateFilter from synapse.types.storage import _BackgroundUpdates @@ -307,8 +307,15 @@ class EventsWorkerStore(SQLBaseStore): # state resolution read the same events. A datastore without a Rust # store, such as the database port script's, gets a cache of its own. self._get_event_cache: RustEventCache + # With the native Postgres driver, events are loaded from the database + # in Rust too, through the same loader as the Rust readers. Loads of + # the same event from either side then share one database read. + self._rust_event_loader: RustStore | None = None if isinstance(self, RustWorkerStore): - self._get_event_cache = self.get_or_build_rust_store(hs).event_cache + rust_store = self.get_or_build_rust_store(hs) + self._get_event_cache = rust_store.event_cache + if rust_store.loads_events: + self._rust_event_loader = rust_store else: self._get_event_cache = RustEventCache( hs.get_rust_runtime(), hs.config.caches.event_cache_size @@ -317,7 +324,8 @@ class EventsWorkerStore(SQLBaseStore): # Map from event ID to a deferred that will result in a map from event # ID to cache entry. Note that the returned dict may not have the - # requested event in it if the event isn't in the DB. + # requested event in it if the event isn't in the DB. Unused when + # loading in Rust, where the event cache coalesces loads instead. self._current_event_fetches: dict[ str, ObservableDeferred[dict[str, EventCacheEntry]] ] = {} @@ -878,6 +886,15 @@ class EventsWorkerStore(SQLBaseStore): missing_events_ids = {e for e in event_ids if e not in event_entry_map} + if self._rust_event_loader is not None: + if missing_events_ids: + event_entry_map.update( + await self._load_events_in_rust( + self._rust_event_loader, missing_events_ids, reload=False + ) + ) + return self._filter_rejected(event_entry_map, allow_rejected) + # We now look up if we're already fetching some of the events in the DB, # if so we wait for those lookups to finish instead of pulling the same # events out of the DB multiple times. @@ -977,14 +994,56 @@ class EventsWorkerStore(SQLBaseStore): if event_id in already_fetching_ids ) - if not allow_rejected: - event_entry_map = { - event_id: entry - for event_id, entry in event_entry_map.items() - if not entry.event.rejected_reason - } + return self._filter_rejected(event_entry_map, allow_rejected) - return event_entry_map + @staticmethod + def _filter_rejected( + event_entry_map: dict[str, EventCacheEntry], allow_rejected: bool + ) -> dict[str, EventCacheEntry]: + if allow_rejected: + return event_entry_map + return { + event_id: entry + for event_id, entry in event_entry_map.items() + if not entry.event.rejected_reason + } + + async def _load_events_in_rust( + self, rust_store: RustStore, event_ids: Collection[str], reload: bool + ) -> dict[str, EventCacheEntry]: + """Load events through the Rust store, which checks the event cache, + joins loads already in flight and reads the rest from the database, + caching what it reads. With `reload`, every event is read from the + database. + + May return rejected events, and events that weren't asked for. + """ + ids = list(event_ids) + # Cancelling this leaves the load running in Rust for any other + # callers sharing it. It never uses our logcontext, so we don't wait + # for it. + if reload: + loaded, fetched = await rust_store.reload_events(ids) + else: + loaded, fetched = await rust_store.load_events(ids) + current_context().record_event_fetch(fetched) + + result_map: dict[str, EventCacheEntry] = {} + for event_id, (event, redacted_event) in loaded.items(): + if redacted_event is None: + # Keep handing out one object per event: a caller that joined + # another's load gets a new object, so use the one already + # handed out if there is one. + live = None if reload else self._event_ref.get(event_id) + if live is not None: + event = live + else: + # We only cache references to unredacted events. + self._event_ref[event_id] = event + result_map[event_id] = EventCacheEntry( + event=event, redacted_event=redacted_event + ) + return result_map def invalidate_get_event_cache_after_txn( self, txn: LoggingTransaction, event_id: str @@ -1363,6 +1422,11 @@ class EventsWorkerStore(SQLBaseStore): str(len(event_ids)), ) + if self._rust_event_loader is not None: + return await self._load_events_in_rust( + self._rust_event_loader, event_ids, reload=True + ) + fetched_event_ids: set[str] = set() fetched_events: dict[str, _EventRow] = {} diff --git a/synapse/synapse_rust/storage.pyi b/synapse/synapse_rust/storage.pyi index 69d74b8de6..4066918750 100644 --- a/synapse/synapse_rust/storage.pyi +++ b/synapse/synapse_rust/storage.pyi @@ -83,6 +83,40 @@ class RustStore: @property def is_feature_enabled_for_user(self) -> IsFeatureEnabledForUser: ... @property + def loads_events(self) -> bool: + """Whether `load_events` can be used: whether the database uses the + native Postgres driver.""" + + def load_events( + self, event_ids: list[str] + ) -> "Deferred[tuple[dict[str, tuple[Event, Event | None]], int]]": + """`_get_events_from_db` behind the event cache: the events by ID, each + with its redacted form if it has a valid redaction, rejected events + included and unknown ones left out. Each comes from the cache, from a + load already in flight, or from the database, and what is read from + the database is cached. Also the number of events this call read from + the database, which leaves out those of a load it joined. + + May return events that weren't asked for: the redaction events read + to check an unconfirmed redaction. Cancelling the Deferred leaves the + load running for the other callers sharing it.""" + + def reload_events( + self, event_ids: list[str] + ) -> "Deferred[tuple[dict[str, tuple[Event, Event | None]], int]]": + """`load_events`, reading every event from the database. What it reads + replaces what was cached.""" + + def load_events_blocking( + self, event_ids: list[str] + ) -> tuple[dict[str, tuple[Event, Event | None]], int]: + """`load_events`, blocking. For tests.""" + + def reload_events_blocking( + self, event_ids: list[str] + ) -> tuple[dict[str, tuple[Event, Event | None]], int]: + """`reload_events`, blocking. For tests.""" + @property def event_cache(self) -> RustEventCache: """The event cache, for `EventsWorkerStore` to bind as `_get_event_cache`.""" diff --git a/tests/replication/_base.py b/tests/replication/_base.py index 310049bf89..3bb0c9fa5d 100644 --- a/tests/replication/_base.py +++ b/tests/replication/_base.py @@ -83,8 +83,11 @@ class BaseStreamTestCase(unittest.HomeserverTestCase): ) # Since we use sqlite in memory databases we need to make sure the - # databases objects are the same. - self.worker_hs.get_datastores().main.db_pool = hs.get_datastores().main.db_pool + # databases objects are the same. The worker's Rust store keeps the + # worker's own `DatabasePool`, so point that at the same database too. + worker_store = self.worker_hs.get_datastores().main + worker_store.db_pool._db_pool = hs.get_datastores().main.db_pool._db_pool + worker_store.db_pool = hs.get_datastores().main.db_pool # Normally we'd pass in the handler to `setup_test_homeserver`, which would # eventually hit "Install @cache_in_self attributes" in tests/utils.py. diff --git a/tests/server.py b/tests/server.py index 45439187cb..e99c2cc195 100644 --- a/tests/server.py +++ b/tests/server.py @@ -110,6 +110,7 @@ from synapse.synapse_rust.runtime import RustRuntime, set_virtual_time_msec from synapse.synapse_rust.state import ( StateResolutionHandler as RustStateResolutionHandler, ) +from synapse.synapse_rust.storage import RustStore from synapse.synapse_rust.sync import SyncHandler as RustSyncHandler from synapse.types import ISynapseReactor, JsonDict from synapse.util.clock import Clock @@ -1530,10 +1531,21 @@ def setup_test_homeserver( sync_handler = hs.get_sync_handler() if sync_handler._rust_sync is not None: sync_handler._rust_sync = cast( - Any, _BlockingRustSyncHandler(sync_handler._rust_sync) + Any, _BlockingRustCalls(sync_handler._rust_sync) ) sync_handler._rust_sync_strict = RUST_SYNC_STRICT_FOR_TESTS + # Likewise for event loads, when they run in Rust: code that loads events + # in the background would otherwise not have finished when a test checks. + # `GetEventCancellationTestCase` exercises the deferreds. + # (Some tests build a homeserver that is never set up.) + if hs.datastores is not None: + main_store = hs.get_datastores().main + if main_store._rust_event_loader is not None: + main_store._rust_event_loader = cast( + Any, _BlockingRustCalls(main_store._rust_event_loader) + ) + return hs @@ -1584,16 +1596,16 @@ class _BlockingRustStateResolution: return fail() -class _BlockingRustSyncHandler: - """Wraps the Rust `SyncHandler` so that each of its entry points blocks, - for the reason `_BlockingRustStateResolution` gives. +class _BlockingRustCalls: + """Wraps the Rust `SyncHandler` or `RustStore` so that each of its entry + points blocks, for the reason `_BlockingRustStateResolution` gives. - Every entry point `name` on the Rust handler that returns a Deferred has a - `name_blocking` twin that returns the result directly. This wrapper - answers `name` with that twin's result in an already-fired Deferred. + Every entry point `name` that returns a Deferred has a `name_blocking` + twin that returns the result directly. This wrapper answers `name` with + that twin's result in an already-fired Deferred. """ - def __init__(self, handler: "RustSyncHandler") -> None: + def __init__(self, handler: "RustSyncHandler | RustStore") -> None: self._handler = handler def __getattr__(self, name: str) -> Any: diff --git a/tests/storage/databases/main/test_events_worker.py b/tests/storage/databases/main/test_events_worker.py index 108b3b8b6c..ea4e9866df 100644 --- a/tests/storage/databases/main/test_events_worker.py +++ b/tests/storage/databases/main/test_events_worker.py @@ -32,10 +32,12 @@ from synapse.logging.context import LoggingContext from synapse.rest import admin from synapse.rest.client import login, room from synapse.server import HomeServer +from synapse.storage.database import make_conn from synapse.storage.databases.main.events_worker import ( EVENT_QUEUE_THREADS, EventsWorkerStore, ) +from synapse.storage.rust_pool import RustConnectionPool from synapse.storage.types import Connection from synapse.util.async_helpers import yieldable_gather_results from synapse.util.clock import Clock @@ -497,6 +499,23 @@ class DatabaseOutageTestCase(unittest.HomeserverTestCase): """ connection_pool = self.store.db_pool._db_pool + if self.store._rust_event_loader is not None: + # Events load through the native pool, which doesn't use the + # connection factory. Point it at a server that isn't there. + assert isinstance(connection_pool, RustConnectionPool) + dsn = connection_pool._dsn + connection_pool.close() + connection_pool._dsn = "host=127.0.0.1 port=1 connect_timeout=1" + connection_pool.start() + try: + yield + finally: + connection_pool.close() + connection_pool._dsn = dsn + connection_pool.start() + self._populate_events() + return + # Close all connections and shut down the database `ThreadPool`. connection_pool.close() @@ -524,10 +543,28 @@ class DatabaseOutageTestCase(unittest.HomeserverTestCase): failure = self.get_failure( self.store.get_event(self.event_ids[0]), Exception ) - self.assertEqual(str(failure.value), "Could not connect to the database.") + if self.store._rust_event_loader is not None: + self.assertIn("loading events failed", str(failure.value)) + else: + self.assertEqual( + str(failure.value), "Could not connect to the database." + ) def test_recovery(self) -> None: """Test that event fetchers recover after a database outage.""" + if self.store._rust_event_loader is not None: + # There are no fetch threads to exhaust: each fetch fails on its + # own, and the next one after the outage succeeds. + with self._outage(): + event_deferreds = [ + ensureDeferred(self.store.get_event(event_id)) + for event_id in self.event_ids + ] + for event_deferred in event_deferreds: + self.get_failure(event_deferred, Exception) + self.get_success(self.store.get_event(self.event_ids[0])) + return + with self.assertLogs( "synapse.metrics.background_process_metrics", level="ERROR" ): @@ -579,6 +616,11 @@ class GetEventCancellationTestCase(unittest.HomeserverTestCase): # Reset the event cache so the tests start with it empty self.store._get_event_cache.clear() + # The test harness makes Rust event loads block. These tests are about + # the deferreds, so use the store itself. + if self.store._rust_event_loader is not None: + self.store._rust_event_loader = hs.get_rust_store() + @contextmanager def blocking_get_event_calls( self, @@ -596,6 +638,11 @@ class GetEventCancellationTestCase(unittest.HomeserverTestCase): * A cancellable `Deferred` for the first `get_event` call. * A cancellable `Deferred` for the second `get_event` call. """ + if self.store._rust_event_loader is not None: + with self._blocking_rust_get_event_calls() as calls: + yield calls + return + # Patch `DatabasePool.runWithConnection` to block. unblock: "Deferred[None]" = Deferred() original_runWithConnection = self.store.db_pool.runWithConnection @@ -630,6 +677,58 @@ class GetEventCancellationTestCase(unittest.HomeserverTestCase): self.assertEqual(ctx1.get_resource_usage().evt_db_fetch_count, 1) self.assertEqual(ctx2.get_resource_usage().evt_db_fetch_count, 0) + @contextmanager + def _blocking_rust_get_event_calls( + self, + ) -> Generator[ + tuple["Deferred[None]", "Deferred[None]", "Deferred[None]"], None, None + ]: + """`blocking_get_event_calls` when events load in Rust, through the + native pool rather than `runWithConnection`: another connection holds + a lock on `event_json` until `unblock` fires, which blocks the load's + query.""" + lock_conn = make_conn( + db_config=self.store.db_pool._database_config, + engine=self.store.db_pool.engine, + default_txn_name="lock_event_json", + server_name=self.hs.hostname, + ) + cur = lock_conn.cursor() + cur.execute("LOCK TABLE event_json IN ACCESS EXCLUSIVE MODE") + + unblock: "Deferred[None]" = Deferred() + unblock.addCallback(lambda _: lock_conn.rollback()) + + ctx1 = LoggingContext(name="get_event1", server_name=self.hs.hostname) + ctx2 = LoggingContext(name="get_event2", server_name=self.hs.hostname) + + async def get_event(ctx: LoggingContext) -> None: + with ctx: + await self.store.get_event(self.event_id) + + try: + get_event1 = ensureDeferred(get_event(ctx1)) + get_event2 = ensureDeferred(get_event(ctx2)) + + # Both `get_event` calls ought to be blocked. + self.assertNoResult(get_event1) + self.assertNoResult(get_event2) + + yield unblock, get_event1, get_event2 + finally: + if not unblock.called: + unblock.callback(None) + lock_conn.close() + + # Confirm that the two `get_event` calls shared the same database + # fetch. Each call's load runs as its own tokio task, so either may be + # the one that reads, and a cancelled call records nothing. + self.assertLessEqual( + ctx1.get_resource_usage().evt_db_fetch_count + + ctx2.get_resource_usage().evt_db_fetch_count, + 1, + ) + def test_first_get_event_cancelled(self) -> None: """Test cancellation of the first `get_event` call sharing a database fetch. @@ -638,6 +737,17 @@ class GetEventCancellationTestCase(unittest.HomeserverTestCase): call must not abort before the fetch is complete, otherwise the fetch will be using a finished logging context. """ + if self.store._rust_event_loader is not None: + # The load runs in Rust, on no logcontext of ours, so the first call + # is cancelled at once and the load carries on for the second. + with self.blocking_get_event_calls() as (unblock, get_event1, get_event2): + get_event1.cancel() + self.get_failure(get_event1, CancelledError) + self.assertNoResult(get_event2) + unblock.callback(None) + self.get_success(get_event2) + return + with self.blocking_get_event_calls() as (unblock, get_event1, get_event2): # Cancel the first `get_event` call. get_event1.cancel() diff --git a/tests/storage/test_sticky_events.py b/tests/storage/test_sticky_events.py index e81ce129d4..1a11d708ca 100644 --- a/tests/storage/test_sticky_events.py +++ b/tests/storage/test_sticky_events.py @@ -1302,11 +1302,13 @@ class StickyEventsFederationBacklogTestCase(unittest.HomeserverTestCase): """ room_id = self.helper.create_room_as(self.user_id, tok=self.token) - # A sticky event that expires almost immediately. + # A sticky event that expires soon. Not at once: making requests can + # advance the clock, and the background cleanup must not delete the + # event before we've read its stream ID. short_lived_event_id = self.helper.send_sticky_event( room_id, EventTypes.Message, - duration=Duration(milliseconds=1), + duration=Duration(minutes=10), content={"body": "short lived", "msgtype": "m.text"}, tok=self.token, )["event_id"] @@ -1326,7 +1328,7 @@ class StickyEventsFederationBacklogTestCase(unittest.HomeserverTestCase): ) # Advance the reactor and trigger the deletion of expired sticky events - self.reactor.advance(0.002) + self.reactor.advance(Duration(minutes=10).as_secs() + 1) self.get_success(self.store._delete_expired_sticky_events()) self.assertIsNone(