From 1afdeb20df87e6e446adc236fe3daf032f0a45bc Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Fri, 16 Oct 2026 13:47:30 +0000 Subject: [PATCH] Add the reads persisting events needs, in Rust `PersistReader` wraps a `SyncReader` and adds the reads the persist controller makes before its write transaction: forward and state DAG extremities, the prev event walks of `_get_events_which_are_prevs` and `_get_prevs_before_rejected`, `have_seen_events`, memberships by event ID, `is_local_host_in_room_ignoring_users`, the room version, the sliding sync joined room check, `get_last_event_pos_in_room` and the soft failed sticky events. None of them is cached. Each runs once per chunk. The chain cover reads are free functions over a `Querier`, not reader methods. The chain cover step reads, allocates chain IDs and writes in one transaction, so they must run on that transaction's connection. To allow that, `WriteTxn` now implements `Querier`. A `WriteError` that passes through `anyhow` on the way out comes back as itself, so a conflict is still retried. The persister builds its own `SyncStore` over the shared caches. `tests/storage/test_persist_reader.py` runs each read against the Python method or statement it ports, on the native driver. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01UHKYPSc37kWqGACn9aKxKL --- changelog.d/rust-persist-prepare.misc | 1 + rust/src/database/postgres/write.rs | 210 +++++- rust/src/persist/mod.rs | 3 + rust/src/persist/python.rs | 8 + rust/src/persist/reader.rs | 1002 +++++++++++++++++++++++++ rust/src/storage/sync_store.rs | 13 +- synapse/synapse_rust/persist.pyi | 45 ++ tests/storage/test_persist_reader.py | 670 +++++++++++++++++ 8 files changed, 1935 insertions(+), 17 deletions(-) create mode 100644 changelog.d/rust-persist-prepare.misc create mode 100644 rust/src/persist/reader.rs create mode 100644 tests/storage/test_persist_reader.py diff --git a/changelog.d/rust-persist-prepare.misc b/changelog.d/rust-persist-prepare.misc new file mode 100644 index 0000000000..a1932f745a --- /dev/null +++ b/changelog.d/rust-persist-prepare.misc @@ -0,0 +1 @@ +Add the Rust reads for persisting events before the write transaction, behind `experimental_features.rust_event_persistence`. diff --git a/rust/src/database/postgres/write.rs b/rust/src/database/postgres/write.rs index 36751ad972..c9339e232d 100644 --- a/rust/src/database/postgres/write.rs +++ b/rust/src/database/postgres/write.rs @@ -31,19 +31,26 @@ use std::{ fmt::Write as _, - sync::atomic::{AtomicBool, Ordering}, + sync::atomic::{AtomicU8, Ordering}, }; +use async_trait::async_trait; use futures::{future::try_join_all, future::BoxFuture}; use log::warn; use tokio::runtime::Handle; use tokio_postgres::{types::ToSql, Row}; use super::pool::{ConnectionPool, PooledConnection}; +use crate::storage::events::Querier; /// Python's `MAX_NUMBER_OF_ATTEMPTS`. const MAX_ATTEMPTS: u32 = 5; +/// Values of [`WriteTxn::failed`]. +const NOT_FAILED: u8 = 0; +const FAILED: u8 = 1; +const FAILED_RETRYABLE: u8 = 2; + /// Why a write failed. #[derive(Debug)] pub enum WriteError { @@ -111,8 +118,16 @@ impl std::fmt::Display for WriteError { impl std::error::Error for WriteError {} impl From for WriteError { + /// A [`WriteError`] that went through `anyhow` on its way out, as one + /// from a [`Querier`] call does, comes back as itself, so that a + /// conflict is still retried. Anything else is [`WriteError::Other`]. fn from(error: anyhow::Error) -> Self { - WriteError::Other(error) + match error.downcast_ref::() { + Some(WriteError::Other(_)) | None => WriteError::Other(error), + Some(_) => error + .downcast::() + .expect("checked by downcast_ref"), + } } } @@ -147,8 +162,9 @@ pub struct WriteTxn { in_transaction: bool, /// A statement failed, so Postgres has aborted the transaction. A /// `COMMIT` would then roll back without an error, so the closure must - /// not be able to report success after catching the error. - failed: AtomicBool, + /// not be able to report success after catching the error. Holds one of + /// the `FAILED_*` values, for the first failure. + failed: AtomicU8, } impl WriteTxn { @@ -158,7 +174,7 @@ impl WriteTxn { let txn = WriteTxn { conn: Some(conn), in_transaction: true, - failed: AtomicBool::new(false), + failed: AtomicU8::new(NOT_FAILED), }; txn.conn() .batch_execute("BEGIN ISOLATION LEVEL REPEATABLE READ") @@ -174,8 +190,18 @@ impl WriteTxn { /// The error for a failed statement, noting that the transaction is now /// aborted. fn statement_failed(&self, context: &str, error: tokio_postgres::Error) -> WriteError { - self.failed.store(true, Ordering::Relaxed); - WriteError::from_pg(context, error) + let error = WriteError::from_pg(context, error); + let failed = if error.is_retryable() { + FAILED_RETRYABLE + } else { + FAILED + }; + // Keep the first: later statements fail only because it aborted the + // transaction. + let _ = + self.failed + .compare_exchange(NOT_FAILED, failed, Ordering::Relaxed, Ordering::Relaxed); + error } /// Run a statement that returns rows. @@ -289,11 +315,23 @@ impl WriteTxn { } async fn commit(mut self) -> Result<(), WriteError> { - if self.failed.load(Ordering::Relaxed) { - self.rollback().await; - return Err(WriteError::Other(anyhow::anyhow!( - "a statement failed, so the transaction was rolled back" - ))); + match self.failed.load(Ordering::Relaxed) { + NOT_FAILED => {} + // The closure caught an error that a retry may get past, so + // retry rather than fail for good. + FAILED_RETRYABLE => { + self.rollback().await; + return Err(WriteError::Conflict( + "a statement failed with a retryable error, so the transaction was rolled back" + .to_owned(), + )); + } + _ => { + self.rollback().await; + return Err(WriteError::Other(anyhow::anyhow!( + "a statement failed, so the transaction was rolled back" + ))); + } } self.conn() .batch_execute("COMMIT") @@ -319,6 +357,17 @@ impl WriteTxn { } } +/// For the read helpers shared with the read-only paths, such as the chain +/// cover reads in `persist::reader`. +#[async_trait] +impl Querier for WriteTxn { + async fn query(&self, sql: &str, params: &[&(dyn ToSql + Sync)]) -> anyhow::Result> { + WriteTxn::query(self, sql, params) + .await + .map_err(anyhow::Error::new) + } +} + impl Drop for WriteTxn { fn drop(&mut self) { let Some(conn) = self.conn.take() else { return }; @@ -457,6 +506,7 @@ mod tests { Arc, }; + use anyhow::Context; use futures::FutureExt; use tokio::runtime::Runtime; @@ -728,6 +778,142 @@ mod tests { teardown(&rt, &pool, &table); } + #[test] + fn write_errors_survive_anyhow() { + let through = |error: WriteError| -> WriteError { + let error: anyhow::Result<()> = Err(anyhow::Error::new(error)); + error.context("reading the chain links").unwrap_err().into() + }; + assert!(matches!( + through(WriteError::Conflict("simulated".into())), + WriteError::Conflict(_) + )); + assert!(matches!( + through(WriteError::Constraint { + sqlstate: "23505".into(), + message: "simulated".into() + }), + WriteError::Constraint { sqlstate, .. } if sqlstate == "23505" + )); + assert!(matches!( + through(WriteError::ConnectionLost("simulated".into())), + WriteError::ConnectionLost(_) + )); + assert!(matches!( + through(WriteError::Cancelled("simulated".into())), + WriteError::Cancelled(_) + )); + // Anything else keeps its context. + let WriteError::Other(error) = through(WriteError::Other(anyhow::anyhow!("boom"))) else { + panic!("expected Other"); + }; + assert_eq!(format!("{error:#}"), "reading the chain links: boom"); + } + + #[test] + fn serialization_failure_through_querier_is_retried() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + pool.get() + .await + .unwrap() + .batch_execute(&format!("INSERT INTO {table} (id, n) VALUES ('a', 0)")) + .await + .unwrap(); + + // As in `serialization_failure_is_retried`, with the statements + // going through `Querier` and `anyhow`. + let attempts = Arc::new(AtomicU32::new(0)); + let result = run_write_interaction(&pool, "conflict", |txn, _| { + let (pool, table, attempts) = (&pool, &table, Arc::clone(&attempts)); + async move { + let querier: &dyn Querier = txn; + let attempt = attempts.fetch_add(1, Ordering::SeqCst) + 1; + querier + .query(&format!("SELECT n FROM {table} WHERE id = 'a'"), &[]) + .await?; + if attempt == 1 { + pool.get() + .await + .unwrap() + .batch_execute(&format!("UPDATE {table} SET n = n + 10 WHERE id = 'a'")) + .await + .unwrap(); + } + querier + .query( + &format!("UPDATE {table} SET n = n + 1 WHERE id = 'a' RETURNING n"), + &[], + ) + .await + .context("bumping n")?; + Ok(attempt) + } + .boxed() + }) + .await + .unwrap(); + + assert_eq!(result, 2); + assert_eq!(rows(&pool, &table).await[0].1, 11); + }); + teardown(&rt, &pool, &table); + } + + #[test] + fn caught_serialization_failure_is_retried() { + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + pool.get() + .await + .unwrap() + .batch_execute(&format!("INSERT INTO {table} (id, n) VALUES ('a', 0)")) + .await + .unwrap(); + + // The first attempt hits a serialization failure and swallows + // it. The commit is refused, and as the error was retryable the + // whole interaction runs again. + let attempts = Arc::new(AtomicU32::new(0)); + let result = run_write_interaction(&pool, "caught conflict", |txn, _| { + let (pool, table, attempts) = (&pool, &table, Arc::clone(&attempts)); + async move { + let attempt = attempts.fetch_add(1, Ordering::SeqCst) + 1; + txn.query(&format!("SELECT n FROM {table} WHERE id = 'a'"), &[]) + .await?; + if attempt == 1 { + pool.get() + .await + .unwrap() + .batch_execute(&format!("UPDATE {table} SET n = n + 10 WHERE id = 'a'")) + .await + .unwrap(); + } + let update = txn + .execute(&format!("UPDATE {table} SET n = n + 1 WHERE id = 'a'"), &[]) + .await; + if attempt == 1 { + assert!(matches!(update, Err(WriteError::Conflict(_))), "{update:?}"); + } else { + update?; + } + Ok(attempt) + } + .boxed() + }) + .await + .unwrap(); + + assert_eq!(result, 2); + assert_eq!(rows(&pool, &table).await[0].1, 11); + }); + teardown(&rt, &pool, &table); + } + #[test] fn retries_stop_after_five_attempts() { let Some((rt, pool, table)) = setup() else { diff --git a/rust/src/persist/mod.rs b/rust/src/persist/mod.rs index 50750842e9..9fb628fe2c 100644 --- a/rust/src/persist/mod.rs +++ b/rust/src/persist/mod.rs @@ -23,6 +23,7 @@ //! it commits. //! - [`metrics`]: counts for the Python metrics, flushed when scraped. //! - [`python`]: the [`python::Persister`] and test helpers. +//! - [`reader`]: what persisting reads before its write transaction. //! //! Every GIL hold in this tree goes through [`crate::gil::with_gil`] with a //! site named `persist_*`, so that tests can pin the count per site. @@ -33,12 +34,14 @@ pub mod context; pub mod effects; pub mod metrics; pub mod python; +pub mod reader; /// Called when registering modules with python. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { let child_module = PyModule::new(py, "persist")?; child_module.add_class::()?; child_module.add_class::()?; + child_module.add_class::()?; child_module.add_function(wrap_pyfunction!( crate::gil::persist_gil_stats, &child_module diff --git a/rust/src/persist/python.rs b/rust/src/persist/python.rs index bf86b495c1..03b11e63ec 100644 --- a/rust/src/persist/python.rs +++ b/rust/src/persist/python.rs @@ -29,6 +29,7 @@ use crate::{ native_pool::NativePool, store::Store, stream_ids::{PyStreamIdTracker, StreamIdTracker}, + sync_store::{SyncStore, SyncStoreConfig}, }, types::state_filter::state_map_to_py, }; @@ -50,6 +51,9 @@ pub(super) struct PersisterInner { pub pool: NativePool, pub runtime: RustRuntime, pub store: Arc, + /// The sync store over the same caches, which the + /// [`super::reader::PersistReader`] builds on. + pub sync: Arc, pub events_stream: Arc, pub backfill_stream: Arc, pub cache_stream: Arc, @@ -93,6 +97,10 @@ impl Persister { inner: Arc::new(PersisterInner { pool: NativePool::new(connection_pool, "persist_db_pool"), runtime, + sync: Arc::new(SyncStore::new( + Arc::clone(&store.caches), + SyncStoreConfig::from_config(&homeserver.config(py)?), + )), store, events_stream: tracker(intern!(py, "_stream_id_gen"))?, backfill_stream: tracker(intern!(py, "_backfill_id_gen"))?, diff --git a/rust/src/persist/reader.rs b/rust/src/persist/reader.rs new file mode 100644 index 0000000000..2aa1a851a4 --- /dev/null +++ b/rust/src/persist/reader.rs @@ -0,0 +1,1002 @@ +/* + * 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: + * . + */ + +//! What persisting reads before its write transaction, through the native +//! Postgres driver. +//! +//! A [`PersistReader`] is a [`SyncReader`] (for the state, events and +//! positions the sync store already reads) plus the reads the persist +//! controller needs on top. Each is a port of the store method it is named +//! after, and `tests/storage/test_persist_reader.py` runs both against the +//! same database. +//! +//! None of these reads is cached. Python caches a few of them +//! (`get_latest_event_ids_in_room`, `have_seen_event`, +//! `get_room_version_id`), but each runs once per chunk of up to 100 +//! events, and a cache would need invalidating from both languages. +//! +//! The chain cover reads are free functions over a [`Querier`] in place of +//! reader methods. The chain cover step reads, allocates chain IDs and +//! writes `event_auth_chain_to_calculate` in one transaction, so they run +//! on that transaction's connection. + +use std::{ + collections::{HashMap, HashSet}, + ops::Deref, + sync::Arc, +}; + +use anyhow::Context; +use futures::{future::BoxFuture, FutureExt}; +use pyo3::prelude::*; +use serde::Deserialize; +use serde_json::Value; +use tokio::runtime::Handle; +use tokio_postgres::Row; + +use crate::database::postgres::pool::{ConnectionPool, PyConnectionPool}; +use crate::storage::events::Querier; +use crate::storage::sync_store::{ + PersistedEventPosition, SyncReader, SyncStore, SELECT_MANY_BATCH_SIZE, +}; + +/// How many events `_get_events_which_are_prevs` and +/// `_get_prevs_before_rejected` handle per transaction. +const PREVS_BATCH_SIZE: usize = 100; + +/// `chain_links` fetches the links of this many chains per statement, as +/// `_get_chain_links` does. +const CHAIN_LINKS_BATCH_SIZE: usize = 1000; + +/// `EventIdMembership`: the user and membership of a membership event. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct EventIdMembership { + pub user_id: Arc, + pub membership: Arc, +} + +/// A row of the chain cover step's auth event query: an auth event's type +/// and state key, and its `(chain ID, sequence number)` if it has one. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AuthEventChain { + pub event_id: Arc, + pub event_type: Arc, + pub state_key: Arc, + pub chain: Option<(i64, i64)>, +} + +/// A row of `event_auth_chain_links`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct ChainLink { + pub origin_chain_id: i64, + pub origin_sequence_number: i64, + pub target_chain_id: i64, + pub target_sequence_number: i64, +} + +/// A row of `event_auth_chain_to_calculate`, less the room. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ChainToCalculate { + pub event_id: Arc, + pub event_type: Arc, + pub state_key: Arc, +} + +/// The reads of one persist step. See the module docs. +pub struct PersistReader<'a> { + sync: SyncReader<'a>, +} + +impl<'a> Deref for PersistReader<'a> { + type Target = SyncReader<'a>; + + fn deref(&self) -> &SyncReader<'a> { + &self.sync + } +} + +impl<'a> PersistReader<'a> { + pub fn new(sync: SyncReader<'a>) -> Self { + PersistReader { sync } + } + + /// `get_latest_event_ids_in_room`: the room's forward extremities. + pub async fn latest_event_ids_in_room( + &self, + room_id: &str, + ) -> anyhow::Result>> { + let rows = self + .query_one_shot( + "SELECT event_id FROM event_forward_extremities WHERE room_id = $1", + &[&room_id], + ) + .await?; + first_column_set(&rows) + } + + /// `get_state_dag_extremities`: the room's MSC4242 state DAG forward + /// extremities. + pub async fn state_dag_extremities(&self, room_id: &str) -> anyhow::Result>> { + let rows = self + .query_one_shot( + "SELECT event_id FROM msc4242_state_dag_forward_extremities WHERE room_id = $1", + &[&room_id], + ) + .await?; + first_column_set(&rows) + } + + /// `_get_events_which_are_prevs`: those of the events that are a prev + /// event of a persisted event that is not an outlier, rejected or soft + /// failed. An event appears once per such child, as in Python. + pub async fn events_which_are_prevs( + &self, + event_ids: &[Arc], + ) -> anyhow::Result>> { + let mut results = Vec::new(); + if event_ids.is_empty() { + return Ok(results); + } + let borrowed = self.borrow().await?; + for batch in event_ids.chunks(PREVS_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = borrowed + .query( + "SELECT prev_event_id, internal_metadata + FROM event_edges + INNER JOIN events USING (event_id) + LEFT JOIN rejections USING (event_id) + LEFT JOIN event_json USING (event_id) + WHERE + NOT events.outlier + AND rejections.event_id IS NULL + AND prev_event_id = ANY($1)", + &[&ids], + ) + .await?; + for row in rows { + let metadata: Option = row.try_get(1)?; + if !is_soft_failed(metadata.as_deref())? { + let prev_event_id: String = row.try_get(0)?; + results.push(Arc::from(prev_event_id)); + } + } + } + Ok(results) + } + + /// `_get_prevs_before_rejected`: the soft failed and rejected events + /// among the prev events of `event_ids`, and their prev events, walking + /// up the DAG until it finds none. With `include_soft_failed` false, + /// only rejected events count. + /// + /// Each batch walks in one snapshot, as each Python transaction does. + pub async fn prevs_before_rejected( + &self, + event_ids: &[Arc], + include_soft_failed: bool, + ) -> anyhow::Result>> { + let mut existing_prevs: HashSet> = HashSet::new(); + for batch in event_ids.chunks(PREVS_BATCH_SIZE) { + let txn = self.begin().await?; + let mut to_check: Vec> = batch.to_vec(); + while !to_check.is_empty() { + let ids: Vec<&str> = to_check.iter().map(|id| &**id).collect(); + let rows = txn + .query( + "SELECT + prev_event_id, internal_metadata, + rejections.event_id IS NOT NULL + FROM event_edges + INNER JOIN events USING (event_id) + LEFT JOIN rejections USING (event_id) + LEFT JOIN event_json USING (event_id) + WHERE + NOT events.outlier + AND event_id = ANY($1)", + &[&ids], + ) + .await?; + to_check = Vec::new(); + for row in rows { + let prev_event_id: String = row.try_get(0)?; + if existing_prevs.contains(prev_event_id.as_str()) { + continue; + } + let metadata: Option = row.try_get(1)?; + let rejected: bool = row.try_get(2)?; + if (include_soft_failed && is_soft_failed(metadata.as_deref())?) || rejected { + let prev_event_id: Arc = Arc::from(prev_event_id); + existing_prevs.insert(Arc::clone(&prev_event_id)); + to_check.push(prev_event_id); + } + } + } + txn.commit().await?; + } + Ok(existing_prevs) + } + + /// `have_seen_events`: which of the events are in `events`, in any + /// room. + pub async fn have_seen_events( + &self, + event_ids: &[Arc], + ) -> anyhow::Result>> { + let mut found = HashSet::new(); + if event_ids.is_empty() { + return Ok(found); + } + let borrowed = self.borrow().await?; + for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = borrowed + .query( + "SELECT event_id FROM events WHERE event_id = ANY($1)", + &[&ids], + ) + .await?; + found.extend(first_column_set(&rows)?); + } + Ok(found) + } + + /// `get_membership_from_event_ids`, less the events that are not + /// membership events (which Python maps to `None`). + pub async fn membership_from_event_ids( + &self, + event_ids: &[Arc], + ) -> anyhow::Result, EventIdMembership>> { + let mut memberships = HashMap::new(); + if event_ids.is_empty() { + return Ok(memberships); + } + let borrowed = self.borrow().await?; + for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = borrowed + .query( + "SELECT user_id, membership, event_id FROM room_memberships + WHERE event_id = ANY($1)", + &[&ids], + ) + .await?; + for row in rows { + let user_id: String = row.try_get(0)?; + let membership: String = row.try_get(1)?; + let event_id: String = row.try_get(2)?; + memberships.insert( + Arc::from(event_id), + EventIdMembership { + user_id: Arc::from(user_id), + membership: Arc::from(membership), + }, + ); + } + } + Ok(memberships) + } + + /// `is_local_host_in_room_ignoring_users`: whether a local user other + /// than those in `ignore_users` is joined to the room. + pub async fn is_local_host_in_room_ignoring_users( + &self, + room_id: &str, + ignore_users: &[Arc], + ) -> anyhow::Result { + let ignore: Vec<&str> = ignore_users.iter().map(|id| &**id).collect(); + let rows = self + .query_one_shot( + "SELECT 1 FROM local_current_membership + WHERE + room_id = $1 AND membership = 'join' + AND NOT (user_id = ANY($2)) + LIMIT 1", + &[&room_id, &ignore], + ) + .await?; + Ok(!rows.is_empty()) + } + + /// `get_room_version_id`, with `None` where Python raises + /// `NotFoundError`. + pub async fn room_version_id(&self, room_id: &str) -> anyhow::Result>> { + let rows = self + .query_one_shot( + "SELECT room_version FROM rooms WHERE room_id = $1", + &[&room_id], + ) + .await?; + let Some(row) = rows.first() else { + return Ok(None); + }; + let version: Option = row.try_get(0)?; + Ok(version.map(Arc::from)) + } + + /// Whether the room has a `sliding_sync_joined_rooms` row, as + /// `_calculate_sliding_sync_table_changes` checks. + pub async fn has_sliding_sync_joined_room(&self, room_id: &str) -> anyhow::Result { + let rows = self + .query_one_shot( + "SELECT room_id FROM sliding_sync_joined_rooms WHERE room_id = $1", + &[&room_id], + ) + .await?; + Ok(!rows.is_empty()) + } + + /// `get_last_event_pos_in_room`: the last event in the room that is not + /// an outlier or rejected, of one of `event_types` if given and not + /// empty, and where it was persisted. + pub async fn last_event_pos_in_room( + &self, + room_id: &str, + event_types: Option<&[&str]>, + ) -> anyhow::Result, PersistedEventPosition)>> { + let rows = match event_types { + Some(types) if !types.is_empty() => { + self.query_one_shot( + "SELECT event_id, stream_ordering, instance_name + FROM events + LEFT JOIN rejections USING (event_id) + WHERE room_id = $1 + AND type = ANY($2) + AND NOT outlier + AND rejections.event_id IS NULL + ORDER BY stream_ordering DESC + LIMIT 1", + &[&room_id, &types], + ) + .await? + } + _ => { + self.query_one_shot( + "SELECT event_id, stream_ordering, instance_name + FROM events + LEFT JOIN rejections USING (event_id) + WHERE room_id = $1 + AND NOT outlier + AND rejections.event_id IS NULL + ORDER BY stream_ordering DESC + LIMIT 1", + &[&room_id], + ) + .await? + } + }; + let Some(row) = rows.first() else { + return Ok(None); + }; + let event_id: String = row.try_get(0)?; + let stream: i64 = row.try_get(1)?; + let instance_name: Option = row.try_get(2)?; + Ok(Some(( + Arc::from(event_id), + PersistedEventPosition { + // A null instance name is a row from before sharded persisters. + // Python's `instance_name or "master"` maps "" the same way. + instance_name: Arc::from( + instance_name + .as_deref() + .filter(|name| !name.is_empty()) + .unwrap_or("master"), + ), + stream, + }, + ))) + } + + /// `_get_soft_failed_sticky_events_txn`: the room's sticky events that + /// are soft failed and expire after `now_ms`, sent by one of `senders` + /// if given. + pub async fn soft_failed_sticky_events( + &self, + room_id: &str, + senders: Option<&[Arc]>, + now_ms: i64, + ) -> anyhow::Result>> { + let rows = match senders { + Some([]) => return Ok(Vec::new()), + Some(senders) => { + let senders: Vec<&str> = senders.iter().map(|id| &**id).collect(); + self.query_one_shot( + "SELECT se.event_id + FROM sticky_events se + INNER JOIN event_json ej USING (event_id) + WHERE + se.room_id = $1 + AND $2 < se.expires_at + AND COALESCE(((ej.internal_metadata::jsonb)->>'soft_failed')::boolean, FALSE) + AND se.sender = ANY($3)", + &[&room_id, &now_ms, &senders], + ) + .await? + } + None => { + self.query_one_shot( + "SELECT se.event_id + FROM sticky_events se + INNER JOIN event_json ej USING (event_id) + WHERE + se.room_id = $1 + AND $2 < se.expires_at + AND COALESCE(((ej.internal_metadata::jsonb)->>'soft_failed')::boolean, FALSE)", + &[&room_id, &now_ms], + ) + .await? + } + }; + first_column_vec(&rows) + } +} + +// The chain cover reads. See the module docs. + +/// Whether the chain cover index is kept for the room, as +/// `calculate_chain_cover_index_for_events_txn` checks: not for an unknown +/// room, nor for a legacy room marked `false`. A null counts as kept. +pub async fn has_auth_chain_index(q: &dyn Querier, room_id: &str) -> anyhow::Result { + let rows = q + .query( + "SELECT has_auth_chain_index FROM rooms WHERE room_id = $1", + &[&room_id], + ) + .await?; + let Some(row) = rows.first() else { + return Ok(false); + }; + let has_index: Option = row.try_get(0)?; + Ok(has_index != Some(false)) +} + +/// Which of the events already have a chain ID. +pub async fn events_with_chain_ids( + q: &dyn Querier, + event_ids: &[Arc], +) -> anyhow::Result>> { + let mut found = HashSet::new(); + for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = q + .query( + "SELECT event_id FROM event_auth_chains WHERE event_id = ANY($1)", + &[&ids], + ) + .await?; + found.extend(first_column_set(&rows)?); + } + Ok(found) +} + +/// The `event_auth_chain_to_calculate` rows of the rooms: out of band +/// memberships whose chains were put off until their auth events arrived. +pub async fn chains_to_calculate( + q: &dyn Querier, + room_ids: &[&str], +) -> anyhow::Result> { + if room_ids.is_empty() { + return Ok(Vec::new()); + } + let rows = q + .query( + "SELECT event_id, type, state_key FROM event_auth_chain_to_calculate + WHERE room_id = ANY($1)", + &[&room_ids], + ) + .await?; + rows.iter() + .map(|row| { + Ok(ChainToCalculate { + event_id: Arc::from(row.try_get::<_, String>(0)?), + event_type: Arc::from(row.try_get::<_, String>(1)?), + state_key: Arc::from(row.try_get::<_, String>(2)?), + }) + }) + .collect() +} + +/// The auth event IDs of a persisted event, from `event_auth`. +pub async fn auth_event_ids(q: &dyn Querier, event_id: &str) -> anyhow::Result>> { + let rows = q + .query( + "SELECT auth_id FROM event_auth WHERE event_id = $1", + &[&event_id], + ) + .await?; + first_column_vec(&rows) +} + +/// The type, state key and chain position of persisted state events, for +/// the auth events of the events being persisted. Events that are not +/// persisted state events are left out. +pub async fn auth_event_chains( + q: &dyn Querier, + event_ids: &[Arc], +) -> anyhow::Result> { + let mut chains = Vec::with_capacity(event_ids.len()); + for batch in event_ids.chunks(SELECT_MANY_BATCH_SIZE) { + let ids: Vec<&str> = batch.iter().map(|id| &**id).collect(); + let rows = q + .query( + "SELECT event_id, events.type, se.state_key, chain_id, sequence_number + FROM events + INNER JOIN state_events AS se USING (event_id) + LEFT JOIN event_auth_chains USING (event_id) + WHERE event_id = ANY($1)", + &[&ids], + ) + .await?; + for row in rows { + let chain_id: Option = row.try_get(3)?; + let sequence_number: Option = row.try_get(4)?; + chains.push(AuthEventChain { + event_id: Arc::from(row.try_get::<_, String>(0)?), + event_type: Arc::from(row.try_get::<_, String>(1)?), + state_key: Arc::from(row.try_get::<_, String>(2)?), + chain: chain_id.zip(sequence_number), + }); + } + } + Ok(chains) +} + +/// `_get_chain_links`: the links from the chains, and from the chains they +/// reach, recursively. May return links not reachable from them. +pub async fn chain_links( + q: &dyn Querier, + chains: impl IntoIterator, +) -> anyhow::Result> { + let mut to_fetch: HashSet = chains.into_iter().collect(); + let mut links = Vec::new(); + while !to_fetch.is_empty() { + let batch: Vec = to_fetch + .iter() + .copied() + .take(CHAIN_LINKS_BATCH_SIZE) + .collect(); + for chain_id in &batch { + to_fetch.remove(chain_id); + } + let rows = q + .query( + "WITH RECURSIVE links(chain_id) AS ( + SELECT + DISTINCT origin_chain_id + FROM event_auth_chain_links WHERE origin_chain_id = ANY($1) + UNION + SELECT + target_chain_id + FROM event_auth_chain_links + INNER JOIN links ON (chain_id = origin_chain_id) + ) + SELECT + origin_chain_id, origin_sequence_number, + target_chain_id, target_sequence_number + FROM links + INNER JOIN event_auth_chain_links ON (chain_id = origin_chain_id)", + &[&batch], + ) + .await?; + for row in rows { + let link = ChainLink { + origin_chain_id: row.try_get(0)?, + origin_sequence_number: row.try_get(1)?, + target_chain_id: row.try_get(2)?, + target_sequence_number: row.try_get(3)?, + }; + // A chain reached from this batch needs no fetch of its own. + to_fetch.remove(&link.origin_chain_id); + links.push(link); + } + } + Ok(links) +} + +/// The highest sequence number allocated on each of the chains, for +/// `_allocate_chain_ids`. +pub async fn max_sequence_numbers( + q: &dyn Querier, + chain_ids: &[i64], +) -> anyhow::Result> { + if chain_ids.is_empty() { + return Ok(HashMap::new()); + } + let rows = q + .query( + "SELECT chain_id, MAX(sequence_number) FROM event_auth_chains + WHERE chain_id = ANY($1) + GROUP BY chain_id", + &[&chain_ids], + ) + .await?; + rows.iter() + .map(|row| Ok((row.try_get(0)?, row.try_get(1)?))) + .collect() +} + +/// The first column of each row, as strings. +fn first_column_vec(rows: &[Row]) -> anyhow::Result>> { + rows.iter() + .map(|row| Ok(Arc::from(row.try_get::<_, String>(0)?))) + .collect() +} + +fn first_column_set(rows: &[Row]) -> anyhow::Result>> { + rows.iter() + .map(|row| Ok(Arc::from(row.try_get::<_, String>(0)?))) + .collect() +} + +/// Whether an `event_json.internal_metadata` value marks the event soft +/// failed: `db_to_json(metadata).get("soft_failed")` is truthy. +fn is_soft_failed(metadata: Option<&str>) -> anyhow::Result { + #[derive(Deserialize)] + struct Metadata { + #[serde(default)] + soft_failed: Value, + } + // Python's `db_to_json(None)` raises here, so a missing `event_json` + // row fails the persist rather than counting the event as live. + let metadata = metadata.context("event has no internal metadata")?; + let metadata: Metadata = serde_json::from_str(metadata)?; + Ok(is_truthy(&metadata.soft_failed)) +} + +/// Python's truth value of a JSON value. +fn is_truthy(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Bool(b) => *b, + Value::Number(n) => n.as_f64().is_some_and(|n| n != 0.0), + Value::String(s) => !s.is_empty(), + Value::Array(a) => !a.is_empty(), + Value::Object(o) => !o.is_empty(), + } +} + +// Test-only entry points, so trial can compare the reads against the Python +// store on a real schema. + +fn anyhow_to_py(e: anyhow::Error) -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err(format!("{e:#}")) +} + +fn arcs(ids: Vec) -> Vec> { + ids.into_iter().map(Arc::from).collect() +} + +fn strings(ids: impl IntoIterator>) -> HashSet { + ids.into_iter().map(|id| id.to_string()).collect() +} + +/// The persist reads, each run to completion on the tokio runtime with the +/// GIL released, over the native driver's pool and the persister's sync +/// store. For tests only. +#[pyclass(name = "PersistReaderForTest")] +pub struct PyPersistReaderForTest { + store: Arc, + pool: ConnectionPool, + handle: Handle, +} + +impl PyPersistReaderForTest { + /// Run `f` over a reader, blocking until it is done. + fn run<'a, T, F>(&'a self, py: Python<'_>, f: F) -> PyResult + where + T: Send + 'a, + F: FnOnce(PersistReader<'a>) -> BoxFuture<'a, anyhow::Result> + Send, + { + py.detach(|| { + self.handle + .block_on(f(PersistReader::new(self.store.reader(self.pool.clone())))) + }) + .map_err(anyhow_to_py) + } +} + +#[pymethods] +impl PyPersistReaderForTest { + #[new] + fn new(pool: &Bound<'_, PyConnectionPool>, persister: &super::python::Persister) -> Self { + let pool = pool.get(); + PyPersistReaderForTest { + store: Arc::clone(&persister.inner.sync), + pool: pool.pool.clone(), + handle: pool.handle.clone(), + } + } + + fn latest_event_ids_in_room( + &self, + py: Python<'_>, + room_id: String, + ) -> PyResult> { + self.run(py, |reader| { + async move { reader.latest_event_ids_in_room(&room_id).await }.boxed() + }) + .map(strings) + } + + fn state_dag_extremities(&self, py: Python<'_>, room_id: String) -> PyResult> { + self.run(py, |reader| { + async move { reader.state_dag_extremities(&room_id).await }.boxed() + }) + .map(strings) + } + + fn events_which_are_prevs( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult> { + let ids = arcs(event_ids); + let prevs = self.run(py, |reader| { + async move { reader.events_which_are_prevs(&ids).await }.boxed() + })?; + Ok(prevs.iter().map(|id| id.to_string()).collect()) + } + + fn prevs_before_rejected( + &self, + py: Python<'_>, + event_ids: Vec, + include_soft_failed: bool, + ) -> PyResult> { + let ids = arcs(event_ids); + self.run(py, |reader| { + async move { + reader + .prevs_before_rejected(&ids, include_soft_failed) + .await + } + .boxed() + }) + .map(strings) + } + + fn have_seen_events( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult> { + let ids = arcs(event_ids); + self.run(py, |reader| { + async move { reader.have_seen_events(&ids).await }.boxed() + }) + .map(strings) + } + + /// Event ID to `(user_id, membership)`. + fn membership_from_event_ids( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult> { + let ids = arcs(event_ids); + let memberships = self.run(py, |reader| { + async move { reader.membership_from_event_ids(&ids).await }.boxed() + })?; + Ok(memberships + .into_iter() + .map(|(event_id, m)| { + ( + event_id.to_string(), + (m.user_id.to_string(), m.membership.to_string()), + ) + }) + .collect()) + } + + fn is_local_host_in_room_ignoring_users( + &self, + py: Python<'_>, + room_id: String, + ignore_users: Vec, + ) -> PyResult { + let ignore = arcs(ignore_users); + self.run(py, |reader| { + async move { + reader + .is_local_host_in_room_ignoring_users(&room_id, &ignore) + .await + } + .boxed() + }) + } + + fn room_version_id(&self, py: Python<'_>, room_id: String) -> PyResult> { + let version = self.run(py, |reader| { + async move { reader.room_version_id(&room_id).await }.boxed() + })?; + Ok(version.map(|v| v.to_string())) + } + + fn has_sliding_sync_joined_room(&self, py: Python<'_>, room_id: String) -> PyResult { + self.run(py, |reader| { + async move { reader.has_sliding_sync_joined_room(&room_id).await }.boxed() + }) + } + + /// `(event_id, instance_name, stream)`. + #[pyo3(signature = (room_id, event_types=None))] + fn last_event_pos_in_room( + &self, + py: Python<'_>, + room_id: String, + event_types: Option>, + ) -> PyResult> { + let last = self.run(py, |reader| { + async move { + let types: Option> = event_types + .as_ref() + .map(|types| types.iter().map(String::as_str).collect()); + reader + .last_event_pos_in_room(&room_id, types.as_deref()) + .await + } + .boxed() + })?; + Ok(last.map(|(event_id, pos)| { + ( + event_id.to_string(), + pos.instance_name.to_string(), + pos.stream, + ) + })) + } + + #[pyo3(signature = (room_id, senders, now_ms))] + fn soft_failed_sticky_events( + &self, + py: Python<'_>, + room_id: String, + senders: Option>, + now_ms: i64, + ) -> PyResult> { + let senders = senders.map(arcs); + let events = self.run(py, |reader| { + async move { + reader + .soft_failed_sticky_events(&room_id, senders.as_deref(), now_ms) + .await + } + .boxed() + })?; + Ok(events.iter().map(|id| id.to_string()).collect()) + } + + // The chain cover reads, each on a connection of its own. + + fn has_auth_chain_index(&self, py: Python<'_>, room_id: String) -> PyResult { + self.run(py, |reader| { + async move { has_auth_chain_index(&reader.borrow().await?, &room_id).await }.boxed() + }) + } + + fn events_with_chain_ids( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult> { + let ids = arcs(event_ids); + self.run(py, |reader| { + async move { events_with_chain_ids(&reader.borrow().await?, &ids).await }.boxed() + }) + .map(strings) + } + + /// `(event_id, type, state_key)` rows. + fn chains_to_calculate( + &self, + py: Python<'_>, + room_ids: Vec, + ) -> PyResult> { + let rows = self.run(py, |reader| { + async move { + let rooms: Vec<&str> = room_ids.iter().map(String::as_str).collect(); + chains_to_calculate(&reader.borrow().await?, &rooms).await + } + .boxed() + })?; + Ok(rows + .into_iter() + .map(|r| { + ( + r.event_id.to_string(), + r.event_type.to_string(), + r.state_key.to_string(), + ) + }) + .collect()) + } + + fn auth_event_ids(&self, py: Python<'_>, event_id: String) -> PyResult> { + let ids = self.run(py, |reader| { + async move { auth_event_ids(&reader.borrow().await?, &event_id).await }.boxed() + })?; + Ok(ids.iter().map(|id| id.to_string()).collect()) + } + + /// `(event_id, type, state_key, chain_id, sequence_number)` rows. + #[allow(clippy::type_complexity)] + fn auth_event_chains( + &self, + py: Python<'_>, + event_ids: Vec, + ) -> PyResult, Option)>> { + let ids = arcs(event_ids); + let rows = self.run(py, |reader| { + async move { auth_event_chains(&reader.borrow().await?, &ids).await }.boxed() + })?; + Ok(rows + .into_iter() + .map(|r| { + ( + r.event_id.to_string(), + r.event_type.to_string(), + r.state_key.to_string(), + r.chain.map(|(chain_id, _)| chain_id), + r.chain.map(|(_, seq)| seq), + ) + }) + .collect()) + } + + /// `(origin_chain_id, origin_sequence_number, target_chain_id, + /// target_sequence_number)` rows. + fn chain_links(&self, py: Python<'_>, chains: Vec) -> PyResult> { + let links = self.run(py, |reader| { + async move { chain_links(&reader.borrow().await?, chains).await }.boxed() + })?; + Ok(links + .into_iter() + .map(|l| { + ( + l.origin_chain_id, + l.origin_sequence_number, + l.target_chain_id, + l.target_sequence_number, + ) + }) + .collect()) + } + + fn max_sequence_numbers( + &self, + py: Python<'_>, + chain_ids: Vec, + ) -> PyResult> { + self.run(py, |reader| { + async move { max_sequence_numbers(&reader.borrow().await?, &chain_ids).await }.boxed() + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn soft_failed_follows_python_truthiness() { + assert!(is_soft_failed(None).is_err()); + assert!(!is_soft_failed(Some("{}")).unwrap()); + assert!(!is_soft_failed(Some(r#"{"soft_failed": false}"#)).unwrap()); + assert!(!is_soft_failed(Some(r#"{"soft_failed": null}"#)).unwrap()); + assert!(!is_soft_failed(Some(r#"{"soft_failed": 0}"#)).unwrap()); + assert!(is_soft_failed(Some(r#"{"soft_failed": true}"#)).unwrap()); + assert!(is_soft_failed(Some(r#"{"soft_failed": 1, "outlier": true}"#)).unwrap()); + assert!(is_soft_failed(Some("not json")).is_err()); + } +} diff --git a/rust/src/storage/sync_store.rs b/rust/src/storage/sync_store.rs index bd7836cd83..554c84e4ca 100644 --- a/rust/src/storage/sync_store.rs +++ b/rust/src/storage/sync_store.rs @@ -244,7 +244,7 @@ impl<'a> SyncReader<'a> { /// Borrow a connection for a group of statements that need no common /// snapshot. - async fn borrow(&self) -> anyhow::Result> { + pub(crate) async fn borrow(&self) -> anyhow::Result> { let start = Instant::now(); let conn = self .pool @@ -267,17 +267,20 @@ impl<'a> SyncReader<'a> { /// snapshot. Such a transaction cannot fail with a serialization error /// or a deadlock, so there is no retry loop. Finish with /// [`Borrowed::commit`]. - async fn begin(&self) -> anyhow::Result> { + pub(crate) async fn begin(&self) -> anyhow::Result> { let mut borrowed = self.borrow().await?; + // Set before awaiting: if we are dropped after BEGIN is sent but + // before it returns, the transaction is open and must be rolled back + // before the connection goes back to the pool. + borrowed.in_transaction = true; borrowed .execute("BEGIN ISOLATION LEVEL REPEATABLE READ READ ONLY") .await?; - borrowed.in_transaction = true; Ok(borrowed) } /// Run one statement on a connection borrowed for it. - async fn query_one_shot( + pub(crate) async fn query_one_shot( &self, sql: &str, params: &[&(dyn ToSql + Sync)], @@ -1997,7 +2000,7 @@ impl<'a> SyncReader<'a> { /// How many values to send per `= ANY` statement, as /// `simple_select_many_batch` does. -const SELECT_MANY_BATCH_SIZE: usize = 1000; +pub(crate) const SELECT_MANY_BATCH_SIZE: usize = 1000; /// How many state groups to load per transaction, as /// `_get_state_groups_from_groups` does. diff --git a/synapse/synapse_rust/persist.pyi b/synapse/synapse_rust/persist.pyi index 4da3152209..5ba788f969 100644 --- a/synapse/synapse_rust/persist.pyi +++ b/synapse/synapse_rust/persist.pyi @@ -15,6 +15,7 @@ from typing import Any from synapse.events import EventBase from synapse.events.snapshot import EventContext +from synapse.synapse_rust.database.postgres import ConnectionPool class Persister: """Persists events in Rust, on the events writer, over the native driver. @@ -99,6 +100,50 @@ class PersistEffects: def new_replication_data(self) -> bool: """Whether rows were written to the cache invalidation stream.""" +class PersistReaderForTest: + """The persist reads (`rust/src/persist/reader.rs`), each run to + completion on the tokio runtime over the native driver's pool. The chain + cover reads each get a connection of their own. For + `tests/storage/test_persist_reader.py`.""" + + def __init__(self, pool: ConnectionPool, persister: Persister) -> None: ... + def latest_event_ids_in_room(self, room_id: str) -> set[str]: ... + def state_dag_extremities(self, room_id: str) -> set[str]: ... + def events_which_are_prevs(self, event_ids: list[str]) -> list[str]: ... + def prevs_before_rejected( + self, event_ids: list[str], include_soft_failed: bool + ) -> set[str]: ... + def have_seen_events(self, event_ids: list[str]) -> set[str]: ... + def membership_from_event_ids( + self, event_ids: list[str] + ) -> dict[str, tuple[str, str]]: + """Event ID to `(user_id, membership)`, for membership events only.""" + def is_local_host_in_room_ignoring_users( + self, room_id: str, ignore_users: list[str] + ) -> bool: ... + def room_version_id(self, room_id: str) -> str | None: ... + def has_sliding_sync_joined_room(self, room_id: str) -> bool: ... + def last_event_pos_in_room( + self, room_id: str, event_types: list[str] | None = None + ) -> tuple[str, str, int] | None: + """`(event_id, instance_name, stream)`.""" + def soft_failed_sticky_events( + self, room_id: str, senders: list[str] | None, now_ms: int + ) -> list[str]: ... + def has_auth_chain_index(self, room_id: str) -> bool: ... + def events_with_chain_ids(self, event_ids: list[str]) -> set[str]: ... + def chains_to_calculate(self, room_ids: list[str]) -> list[tuple[str, str, str]]: + """`(event_id, type, state_key)` rows.""" + def auth_event_ids(self, event_id: str) -> list[str]: ... + def auth_event_chains( + self, event_ids: list[str] + ) -> list[tuple[str, str, str, int | None, int | None]]: + """`(event_id, type, state_key, chain_id, sequence_number)` rows.""" + def chain_links(self, chains: list[int]) -> list[tuple[int, int, int, int]]: + """`(origin_chain_id, origin_sequence_number, target_chain_id, + target_sequence_number)` rows.""" + def max_sequence_numbers(self, chain_ids: list[int]) -> dict[int, int]: ... + def persist_gil_stats() -> dict[str, tuple[int, float]]: """`{site: (calls, seconds held)}` for every time event persistence took the GIL from tokio.""" diff --git a/tests/storage/test_persist_reader.py b/tests/storage/test_persist_reader.py new file mode 100644 index 0000000000..a17bf8075d --- /dev/null +++ b/tests/storage/test_persist_reader.py @@ -0,0 +1,670 @@ +# +# 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: +# . +# + +"""The Rust persist reads (`rust/src/persist/reader.rs`) against the Python +store methods they port, on the same database. These need the native Rust +Postgres driver: + + SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.storage.test_persist_reader +""" + +from collections import Counter +from typing import Any +from unittest import SkipTest + +from twisted.internet.testing import MemoryReactor + +from synapse.api.constants import EventTypes +from synapse.api.errors import NotFoundError +from synapse.rest import admin +from synapse.rest.client import login, room +from synapse.server import HomeServer +from synapse.storage.database import LoggingTransaction +from synapse.storage.databases.main.event_federation import EventFederationStore +from synapse.storage.rust_pool import RustConnectionPool +from synapse.synapse_rust.persist import PersistReaderForTest +from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES +from synapse.util.clock import Clock + +from tests import unittest +from tests.test_utils.event_injection import inject_event +from tests.utils import RUST_PERSIST_FOR_TESTS + + +class PersistReaderTestBase(unittest.HomeserverTestCase): + """The fixture: a room with two members, and the reader on the + persister of a homeserver at `shadow`.""" + + servlets = [ + admin.register_servlets, + room.register_servlets, + login.register_servlets, + ] + + def default_config(self) -> dict[str, Any]: + config = super().default_config() + config.setdefault("experimental_features", {})["rust_event_persistence"] = ( + "shadow" + ) + return config + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + if RUST_PERSIST_FOR_TESTS == "off": + raise SkipTest("SYNAPSE_TEST_RUST_PERSIST=off turns the persister off") + hs = self.setup_test_homeserver() + if not isinstance( + hs.get_datastores().main.db_pool._db_pool, RustConnectionPool + ): + raise SkipTest("needs the native Rust Postgres driver") + return hs + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + self.store = hs.get_datastores().main + persist_events_store = hs.get_datastores().persist_events + assert persist_events_store is not None + self.persist_events_store = persist_events_store + db_pool = self.store.db_pool._db_pool + assert isinstance(db_pool, RustConnectionPool) + assert db_pool._pool is not None + persister = hs.get_rust_handlers().persister + assert persister is not None + self.rust = PersistReaderForTest(db_pool._pool, persister) + + self.alice = self.register_user("alice", "pass") + self.alice_tok = self.login("alice", "pass") + self.bob = self.register_user("bob", "pass") + self.bob_tok = self.login("bob", "pass") + self.room_id = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self.helper.join(self.room_id, self.bob, tok=self.bob_tok) + + def _send(self, body: str) -> str: + return self.helper.send(self.room_id, body, tok=self.alice_tok)["event_id"] + + def _sql(self, sql: str, *args: Any) -> list[tuple[Any, ...]]: + return self.get_success(self.store.db_pool.execute("test", sql, *args)) + + def _write(self, sql: str, *args: Any) -> None: + """Run a statement that returns no rows.""" + self.get_success( + self.store.db_pool.runInteraction( + "test", lambda txn: txn.execute(sql, args) + ) + ) + + def _all_event_ids(self) -> list[str]: + return [ + event_id + for (event_id,) in self._sql( + "SELECT event_id FROM events WHERE room_id = ?", self.room_id + ) + ] + + +class PersistReaderTestCase(PersistReaderTestBase): + """Each test runs a Rust read and the Python store method against the + same fixture and compares.""" + + def test_latest_event_ids_in_room(self) -> None: + message = self._send("hi") + # Two children of one parent fork the room. + for body in ("a", "b"): + self.get_success( + inject_event( + self.hs, + prev_event_ids=[message], + room_id=self.room_id, + sender=self.alice, + type=EventTypes.Message, + content={"msgtype": "m.text", "body": body}, + ) + ) + + for room_id in (self.room_id, "!unknown:test"): + # The Python cache was invalidated as the events went in, so both + # read the table. + self.store.get_latest_event_ids_in_room.invalidate((room_id,)) + expected = set( + self.get_success(self.store.get_latest_event_ids_in_room(room_id)) + ) + self.assertEqual(self.rust.latest_event_ids_in_room(room_id), expected) + self.assertEqual(len(self.rust.latest_event_ids_in_room(self.room_id)), 2) + + def test_state_dag_extremities(self) -> None: + # The table is only written in MSC4242 rooms, so fill it by hand. + for event_id in self._all_event_ids()[:2]: + self.get_success( + self.store.db_pool.simple_insert( + "msc4242_state_dag_forward_extremities", + {"room_id": self.room_id, "event_id": event_id}, + ) + ) + + for room_id in (self.room_id, "!unknown:test"): + self.assertEqual( + self.rust.state_dag_extremities(room_id), + set(self.get_success(self.store.get_state_dag_extremities(room_id))), + ) + self.assertEqual(len(self.rust.state_dag_extremities(self.room_id)), 2) + + def _dag_with_bad_events(self) -> list[str]: + """A chain of events in the room, of which the third is rejected, + the fourth soft failed, the fifth an outlier and the seventh both + rejected and an outlier.""" + chain = [self._send(f"m{i}") for i in range(3)] + chain.append( + self.get_success( + inject_event( + self.hs, + prev_event_ids=[chain[-1]], + internal_metadata={"soft_failed": True}, + room_id=self.room_id, + sender=self.alice, + type=EventTypes.Message, + content={"msgtype": "m.text", "body": "soft failed"}, + ) + ).event_id + ) + chain.extend( + self.get_success( + inject_event( + self.hs, + prev_event_ids=[chain[-1]], + room_id=self.room_id, + sender=self.alice, + type=EventTypes.Message, + content={"msgtype": "m.text", "body": body}, + ) + ).event_id + for body in ("outlier", "last", "rejected outlier") + ) + for event_id in (chain[2], chain[6]): + self._reject(event_id) + for event_id in (chain[4], chain[6]): + self._write("UPDATE events SET outlier = TRUE WHERE event_id = ?", event_id) + return chain + + def _reject(self, event_id: str) -> None: + self.get_success( + self.store.db_pool.simple_insert( + "rejections", + {"event_id": event_id, "reason": "test", "last_check": "now"}, + ) + ) + + def test_events_which_are_prevs(self) -> None: + chain = self._dag_with_bad_events() + # A second child of "last", so it is a prev event twice. + self.get_success( + inject_event( + self.hs, + prev_event_ids=[chain[5]], + room_id=self.room_id, + sender=self.alice, + type=EventTypes.Message, + content={"msgtype": "m.text", "body": "second child"}, + ) + ) + # More than a batch of unknown IDs first, so the real ones are in + # the second batch. + ids = [f"$unknown{i}" for i in range(150)] + self._all_event_ids() + + expected = self.get_success( + self.persist_events_store._get_events_which_are_prevs(ids) + ) + # Row order is the database's, so compare as multisets. + self.assertEqual( + Counter(self.rust.events_which_are_prevs(ids)), Counter(expected) + ) + # The prevs of the rejected, soft failed and outlier events don't + # count. + prevs = set(expected) + for event_id in chain[1:4]: + self.assertNotIn(event_id, prevs) + self.assertIn(chain[0], prevs) + self.assertIn(chain[4], prevs) + # The rejected outlier child doesn't count, the other does. + self.assertEqual(Counter(expected)[chain[5]], 1) + self.assertEqual(Counter(self.rust.events_which_are_prevs(ids))[chain[5]], 1) + + def test_prevs_before_rejected(self) -> None: + chain = self._dag_with_bad_events() + + pad = [f"$unknown{i}" for i in range(150)] + for start in [ + [chain[-1]], + [chain[3]], + chain, + self._all_event_ids(), + pad + chain, + ]: + for include_soft_failed in (True, False): + self.assertEqual( + self.rust.prevs_before_rejected(start, include_soft_failed), + self.get_success( + self.persist_events_store._get_prevs_before_rejected( + start, include_soft_failed + ) + ), + (start, include_soft_failed), + ) + # The walk is from the soft failed event, past the rejected one. + self.assertEqual( + self.rust.prevs_before_rejected([chain[3]], True), {chain[1], chain[2]} + ) + self.assertEqual(self.rust.prevs_before_rejected([chain[3]], False), set()) + self.assertEqual(self.rust.prevs_before_rejected([chain[2]], False), {chain[1]}) + # The rejected outlier's edges are not walked. + self.assertEqual(self.rust.prevs_before_rejected([chain[6]], False), set()) + # In the second batch, which keeps what the first found. + self.assertEqual( + self.rust.prevs_before_rejected(pad + [chain[3]], True), + {chain[1], chain[2]}, + ) + + def test_have_seen_events(self) -> None: + # More than a batch of unknown IDs first. + ids = [f"$unknown{i}" for i in range(1001)] + self._all_event_ids() + self.assertEqual( + self.rust.have_seen_events(ids), + self.get_success(self.store.have_seen_events(self.room_id, ids)), + ) + + def test_membership_from_event_ids(self) -> None: + self.helper.leave(self.room_id, self.bob, tok=self.bob_tok) + ids = [f"$unknown{i}" for i in range(1001)] + self._all_event_ids() + expected = self.get_success(self.store.get_membership_from_event_ids(ids)) + self.assertEqual( + self.rust.membership_from_event_ids(ids), + { + event_id: (m.user_id, m.membership) + for event_id, m in expected.items() + if m is not None + }, + ) + self.assertIn( + "leave", {m for _, m in self.rust.membership_from_event_ids(ids).values()} + ) + + def test_is_local_host_in_room_ignoring_users(self) -> None: + # Local users who are invited, or have left, don't count. + carol = self.register_user("carol", "pass") + carol_tok = self.login("carol", "pass") + self.helper.join(self.room_id, carol, tok=carol_tok) + self.helper.leave(self.room_id, carol, tok=carol_tok) + dave = self.register_user("dave", "pass") + self.helper.invite(self.room_id, self.alice, dave, tok=self.alice_tok) + + for room_id in (self.room_id, "!unknown:test"): + for ignore in ([], [self.alice], [self.alice, self.bob], ["@x:test"]): + self.assertEqual( + self.rust.is_local_host_in_room_ignoring_users(room_id, ignore), + self.get_success( + self.store.is_local_host_in_room_ignoring_users(room_id, ignore) + ), + (room_id, ignore), + ) + self.assertTrue( + self.rust.is_local_host_in_room_ignoring_users(self.room_id, [self.alice]) + ) + self.assertFalse( + self.rust.is_local_host_in_room_ignoring_users( + self.room_id, [self.alice, self.bob] + ) + ) + + def test_room_version_id(self) -> None: + self.assertEqual( + self.rust.room_version_id(self.room_id), + self.get_success(self.store.get_room_version_id(self.room_id)), + ) + self.get_failure(self.store.get_room_version_id("!unknown:test"), NotFoundError) + self.assertIsNone(self.rust.room_version_id("!unknown:test")) + + def test_has_sliding_sync_joined_room(self) -> None: + for room_id in (self.room_id, "!unknown:test"): + row = self.get_success( + self.store.db_pool.simple_select_one_onecol( + table="sliding_sync_joined_rooms", + keyvalues={"room_id": room_id}, + retcol="room_id", + allow_none=True, + ) + ) + self.assertEqual( + self.rust.has_sliding_sync_joined_room(room_id), row is not None + ) + self.assertTrue(self.rust.has_sliding_sync_joined_room(self.room_id)) + + def test_last_event_pos_in_room(self) -> None: + chain = self._dag_with_bad_events() + topic = self.helper.send_state( + self.room_id, EventTypes.Topic, {"topic": "t"}, tok=self.alice_tok + )["event_id"] + # Newer events of each type that don't count: a rejected message and + # an outlier topic. + rejected = self._send("rejected") + self._reject(rejected) + outlier = self.helper.send_state( + self.room_id, EventTypes.Topic, {"topic": "outlier"}, tok=self.alice_tok + )["event_id"] + self._write("UPDATE events SET outlier = TRUE WHERE event_id = ?", outlier) + # A row from before sharded persisters. + self._write("UPDATE events SET instance_name = NULL WHERE event_id = ?", topic) + + for room_id in (self.room_id, "!unknown:test"): + for event_types in ( + None, + [], + list(SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES), + [EventTypes.Topic], + ["m.nothing"], + ): + expected = self.get_success( + self.store.get_last_event_pos_in_room(room_id, event_types) + ) + self.assertEqual( + self.rust.last_event_pos_in_room(room_id, event_types), + None + if expected is None + else (expected[0], expected[1].instance_name, expected[1].stream), + (room_id, event_types), + ) + + def last(event_types: list[str] | None) -> tuple[str, str] | None: + pos = self.rust.last_event_pos_in_room(self.room_id, event_types) + return None if pos is None else (pos[0], pos[1]) + + self.assertEqual(last(None), (topic, "master")) + self.assertEqual(last([EventTypes.Topic]), (topic, "master")) + self.assertEqual(last([EventTypes.Message]), (chain[5], "master")) + + def test_soft_failed_sticky_events(self) -> None: + # Sticky events need MSC4354, so fill the table by hand: a soft + # failed event from each user, and one that is not soft failed, with + # different expiries. + start_ms = self.clock.time_msec() + events: list[str] = [] + for sender, soft_failed, expires_in in ( + (self.alice, True, 2000), + (self.bob, True, 4000), + (self.alice, False, 4000), + ): + event = self.get_success( + inject_event( + self.hs, + internal_metadata={"soft_failed": soft_failed}, + room_id=self.room_id, + sender=sender, + type=EventTypes.Message, + content={"msgtype": "m.text", "body": f"sticky {len(events)}"}, + ) + ) + events.append(event.event_id) + self.get_success( + self.store.db_pool.simple_insert( + "sticky_events", + { + "stream_id": len(events), + "instance_name": "master", + "event_id": event.event_id, + "room_id": self.room_id, + "event_stream_ordering": event.internal_metadata.stream_ordering, + "sender": sender, + "expires_at": start_ms + expires_in, + }, + ) + ) + + # A soft failed sticky event in another room, which never shows. + other_room = self.helper.create_room_as(self.alice, tok=self.alice_tok) + other = self.get_success( + inject_event( + self.hs, + internal_metadata={"soft_failed": True}, + room_id=other_room, + sender=self.alice, + type=EventTypes.Message, + content={"msgtype": "m.text", "body": "elsewhere"}, + ) + ) + self.get_success( + self.store.db_pool.simple_insert( + "sticky_events", + { + "stream_id": 100, + "instance_name": "master", + "event_id": other.event_id, + "room_id": other_room, + "event_stream_ordering": other.internal_metadata.stream_ordering, + "sender": self.alice, + "expires_at": start_ms + 10000, + }, + ) + ) + + # Python reads the clock, so move it to each time. An event expires + # at its `expires_at`. + expected_all = { + 1000: events[:2], + 2000: events[1:2], + 3000: events[1:2], + 4000: [], + 5000: [], + } + for offset_ms, expected_now in expected_all.items(): + self.reactor.advance((start_ms + offset_ms - self.clock.time_msec()) / 1000) + now_ms = self.clock.time_msec() + for senders in (None, set(), {self.alice}, {self.bob, "@x:test"}): + expected = self.get_success( + self.store.db_pool.runInteraction( + "test", + self.store._get_soft_failed_sticky_events_txn, + self.room_id, + senders=senders, + ) + ) + self.assertEqual( + sorted( + self.rust.soft_failed_sticky_events( + self.room_id, + None if senders is None else list(senders), + now_ms, + ) + ), + sorted(expected), + (offset_ms, senders), + ) + self.assertEqual( + sorted(self.rust.soft_failed_sticky_events(self.room_id, None, now_ms)), + sorted(expected_now), + offset_ms, + ) + + +class ChainCoverReadsTestCase(PersistReaderTestBase): + """The chain cover reads, against the statements and helpers of + `_calculate_chain_cover_index`, which are inline in Python.""" + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + super().prepare(reactor, clock, hs) + # Some state changes by both users, so there are several chains with + # links between them. + for i in range(3): + self.helper.send_state( + self.room_id, EventTypes.Topic, {"topic": f"{i}"}, tok=self.alice_tok + ) + self.helper.send_state( + self.room_id, EventTypes.Name, {"name": f"{i}"}, tok=self.alice_tok + ) + self.helper.leave(self.room_id, self.bob, tok=self.bob_tok) + self.helper.join(self.room_id, self.bob, tok=self.bob_tok) + + def _state_event_ids(self) -> list[str]: + return [ + event_id + for (event_id,) in self._sql( + "SELECT event_id FROM state_events WHERE room_id = ?", self.room_id + ) + ] + + def _chain_ids(self) -> list[int]: + return [ + chain_id + for (chain_id,) in self._sql( + """SELECT DISTINCT chain_id FROM event_auth_chains + INNER JOIN events USING (event_id) WHERE room_id = ?""", + self.room_id, + ) + ] + + def test_has_auth_chain_index(self) -> None: + other = self.helper.create_room_as(self.alice, tok=self.alice_tok) + legacy = self.helper.create_room_as(self.alice, tok=self.alice_tok) + self._write( + "UPDATE rooms SET has_auth_chain_index = NULL WHERE room_id = ?", other + ) + self._write( + "UPDATE rooms SET has_auth_chain_index = FALSE WHERE room_id = ?", legacy + ) + + for room_id in (self.room_id, other, legacy, "!unknown:test"): + row = self.get_success( + self.store.db_pool.simple_select_one( + table="rooms", + keyvalues={"room_id": room_id}, + retcols=("room_id", "has_auth_chain_index"), + allow_none=True, + ) + ) + # As `calculate_chain_cover_index_for_events_txn` decides. + expected = not (row is None or row[1] is False) + self.assertEqual(self.rust.has_auth_chain_index(room_id), expected, room_id) + self.assertTrue(self.rust.has_auth_chain_index(other)) + self.assertFalse(self.rust.has_auth_chain_index(legacy)) + + def test_events_with_chain_ids(self) -> None: + ids = self._all_event_ids() + ["$unknown"] + rows = self.get_success( + self.store.db_pool.simple_select_many_batch( + table="event_auth_chains", + column="event_id", + iterable=ids, + keyvalues={}, + retcols=("event_id",), + ) + ) + self.assertEqual( + self.rust.events_with_chain_ids(ids), {event_id for (event_id,) in rows} + ) + self.assertEqual( + self.rust.events_with_chain_ids(ids), set(self._state_event_ids()) + ) + + def test_chains_to_calculate(self) -> None: + other = self.helper.create_room_as(self.alice, tok=self.alice_tok) + for room_id, event_id in ((self.room_id, "$a"), (other, "$b"), (other, "$c")): + self.get_success( + self.store.db_pool.simple_insert( + "event_auth_chain_to_calculate", + { + "event_id": event_id, + "room_id": room_id, + "type": EventTypes.Member, + "state_key": self.bob, + }, + ) + ) + + for room_ids in ([self.room_id], [self.room_id, other], [], ["!unknown:test"]): + rows = self.get_success( + self.store.db_pool.simple_select_many_batch( + table="event_auth_chain_to_calculate", + keyvalues={}, + column="room_id", + iterable=set(room_ids), + retcols=("event_id", "type", "state_key"), + ) + ) + self.assertEqual( + sorted(self.rust.chains_to_calculate(room_ids)), sorted(rows), room_ids + ) + + def test_auth_event_ids(self) -> None: + for event_id in self._all_event_ids() + ["$unknown"]: + expected = self.get_success( + self.store.db_pool.simple_select_onecol( + "event_auth", keyvalues={"event_id": event_id}, retcol="auth_id" + ) + ) + self.assertEqual( + sorted(self.rust.auth_event_ids(event_id)), sorted(expected), event_id + ) + + def test_auth_event_chains(self) -> None: + # An event without a chain, as an out of band membership has. + state = self._state_event_ids() + self._write("DELETE FROM event_auth_chains WHERE event_id = ?", state[0]) + ids = self._all_event_ids() + ["$unknown"] + + # The statement `_calculate_chain_cover_index` runs, through Python. + rows = self._sql( + """SELECT event_id, events.type, se.state_key, chain_id, sequence_number + FROM events + INNER JOIN state_events AS se USING (event_id) + LEFT JOIN event_auth_chains USING (event_id) + WHERE event_id = ANY(?)""", + ids, + ) + actual = self.rust.auth_event_chains(ids) + self.assertEqual(sorted(actual), sorted(rows)) + self.assertEqual(len(actual), len(state)) + self.assertIn((state[0], None, None), {(r[0], r[3], r[4]) for r in actual}) + + def test_chain_links(self) -> None: + chains = self._chain_ids() + self.assertGreater(len(chains), 1) + for fetch in ([chains[0]], chains, [], [-1]): + + def python_links(txn: LoggingTransaction) -> list[tuple[int, ...]]: + return [ + (origin, origin_seq, target, target_seq) + for links in EventFederationStore._get_chain_links(txn, set(fetch)) + for origin, inner in links.items() + for origin_seq, target, target_seq in inner + ] + + expected = self.get_success( + self.store.db_pool.runInteraction("test", python_links) + ) + self.assertEqual( + Counter(self.rust.chain_links(fetch)), Counter(expected), fetch + ) + self.assertTrue(self.rust.chain_links(chains)) + + def test_max_sequence_numbers(self) -> None: + chains = self._chain_ids() + for chain_ids in (chains, chains[:1], [], [-1]): + rows = ( + self._sql( + """SELECT chain_id, MAX(sequence_number) FROM event_auth_chains + WHERE chain_id = ANY(?) + GROUP BY chain_id""", + chain_ids, + ) + if chain_ids + else [] + ) + self.assertEqual( + self.rust.max_sequence_numbers(chain_ids), dict(rows), chain_ids + )