diff --git a/changelog.d/rust-persist-txn.misc b/changelog.d/rust-persist-txn.misc new file mode 100644 index 0000000000..ad0926b363 --- /dev/null +++ b/changelog.d/rust-persist-txn.misc @@ -0,0 +1 @@ +Write persisted events to the database in Rust, behind `experimental_features.rust_event_persistence`. diff --git a/rust/src/database/postgres/write.rs b/rust/src/database/postgres/write.rs index 286cc35d8c..8be001c41b 100644 --- a/rust/src/database/postgres/write.rs +++ b/rust/src/database/postgres/write.rs @@ -44,7 +44,7 @@ use super::pool::{ConnectionPool, PooledConnection}; use crate::storage::events::Querier; /// Python's `MAX_NUMBER_OF_ATTEMPTS`. -const MAX_ATTEMPTS: u32 = 5; +pub const MAX_ATTEMPTS: u32 = 5; /// Values of [`WriteTxn::failed`]. const NOT_FAILED: u8 = 0; @@ -73,24 +73,58 @@ pub enum WriteError { Other(anyhow::Error), } +/// How [`WriteError::from_pg`] classifies a database error. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum PgErrorKind { + Conflict, + Cancelled, + Constraint, + ConnectionLost, + Other, +} + +impl PgErrorKind { + fn of(error: &tokio_postgres::Error) -> Self { + match error.code().map(|code| code.code()) { + Some("40001" | "40P01") => PgErrorKind::Conflict, + Some("57014") => PgErrorKind::Cancelled, + Some(code) if code.starts_with("23") => PgErrorKind::Constraint, + Some(code) if ["08", "57", "58"].iter().any(|c| code.starts_with(c)) => { + PgErrorKind::ConnectionLost + } + None if error.is_closed() => PgErrorKind::ConnectionLost, + _ => PgErrorKind::Other, + } + } + + fn is_retryable(self) -> bool { + matches!( + self, + PgErrorKind::Conflict | PgErrorKind::Cancelled | PgErrorKind::ConnectionLost + ) + } +} + impl WriteError { fn from_pg(context: &str, error: tokio_postgres::Error) -> Self { let message = match error.as_db_error() { Some(db) => format!("{context}: {} ({})", db.message(), db.code().code()), None => format!("{context}: {error}"), }; - match error.code().map(|code| code.code()) { - Some("40001" | "40P01") => WriteError::Conflict(message), - Some("57014") => WriteError::Cancelled(message), - Some(code) if code.starts_with("23") => WriteError::Constraint { - sqlstate: code.to_owned(), + match PgErrorKind::of(&error) { + PgErrorKind::Conflict => WriteError::Conflict(message), + PgErrorKind::Cancelled => WriteError::Cancelled(message), + PgErrorKind::Constraint => WriteError::Constraint { + sqlstate: error + .code() + .map(|code| code.code().to_owned()) + .unwrap_or_default(), message, }, - Some(code) if ["08", "57", "58"].iter().any(|c| code.starts_with(c)) => { - WriteError::ConnectionLost(message) + PgErrorKind::ConnectionLost => WriteError::ConnectionLost(message), + PgErrorKind::Other => { + WriteError::Other(anyhow::Error::new(error).context(context.to_owned())) } - None if error.is_closed() => WriteError::ConnectionLost(message), - _ => WriteError::Other(anyhow::Error::new(error).context(context.to_owned())), } } @@ -103,6 +137,28 @@ impl WriteError { } } +/// Whether running again the work that failed with `error` may succeed: a +/// cause in its chain is a database error of a kind [`run_write_interaction`] +/// retries, such as a lost connection. For work outside a write transaction, +/// such as a group of reads, which Python runs in a `runInteraction` that +/// retries them. +/// +/// A [`WriteError`] in the chain does not count: it comes from a write +/// interaction that has already run its own retries. +pub fn is_retryable_error(error: &anyhow::Error) -> bool { + error.chain().any(|cause| { + if let Some(error) = cause.downcast_ref::() { + PgErrorKind::of(error).is_retryable() + } else if let Some(deadpool::managed::PoolError::Backend(error)) = + cause.downcast_ref::>() + { + PgErrorKind::of(error).is_retryable() + } else { + false + } + }) +} + impl std::fmt::Display for WriteError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { @@ -843,6 +899,37 @@ mod tests { assert_eq!(format!("{error:#}"), "reading the chain links: boom"); } + #[test] + fn retryable_errors_are_found_in_the_chain() { + // A write interaction has already retried its own errors. + let wrapped = |error: WriteError| anyhow::Error::new(error).context("reading"); + assert!(!is_retryable_error(&wrapped(WriteError::Conflict( + "x".into() + )))); + assert!(!is_retryable_error(&wrapped(WriteError::Cancelled( + "x".into() + )))); + assert!(!is_retryable_error(&anyhow::anyhow!("x"))); + + let Some((rt, pool, table)) = setup() else { + return; + }; + rt.block_on(async { + let conn = pool.get().await.unwrap(); + conn.batch_execute("SET statement_timeout = 1") + .await + .unwrap(); + let cancelled = conn.query("SELECT pg_sleep(1)", &[]).await.unwrap_err(); + conn.batch_execute("RESET statement_timeout").await.unwrap(); + let cancelled = anyhow::Error::new(cancelled).context("running a query"); + assert!(is_retryable_error(&cancelled), "{cancelled:#}"); + + let error = conn.query("SELECT 1/0", &[]).await.unwrap_err(); + assert!(!is_retryable_error(&anyhow::Error::new(error))); + }); + teardown(&rt, &pool, &table); + } + #[test] fn serialization_failure_through_querier_is_retried() { let Some((rt, pool, table)) = setup() else { diff --git a/rust/src/events/constants.rs b/rust/src/events/constants.rs index 51f3b572f9..596f10c1b6 100644 --- a/rust/src/events/constants.rs +++ b/rust/src/events/constants.rs @@ -102,6 +102,8 @@ pub mod unsigned_field { /// Relation types (the `rel_type` of an `m.relates_to`). pub mod relation_type { + /// Relation type: m.annotation + pub const ANNOTATION: &str = "m.annotation"; /// Relation type: m.reference pub const REFERENCE: &str = "m.reference"; /// Relation type: m.replace diff --git a/rust/src/homeserver.rs b/rust/src/homeserver.rs index 5f3273d388..f782ae8ddc 100644 --- a/rust/src/homeserver.rs +++ b/rust/src/homeserver.rs @@ -145,6 +145,15 @@ impl HomeServer { .getattr(intern!(py, "db_pool")) } + /// The state database pool (`hs.get_datastores().state.db_pool`). + pub fn state_database_pool<'py>(&self, py: Python<'py>) -> PyResult> { + self.0 + .bind(py) + .call_method0(intern!(py, "get_datastores"))? + .getattr(intern!(py, "state"))? + .getattr(intern!(py, "db_pool")) + } + /// The Rust `Store` (`hs.get_rust_store()`), which the main datastore /// builds and owns. See [`crate::storage::python`]. pub fn get_rust_store(&self, py: Python<'_>) -> PyResult> { diff --git a/rust/src/persist/mod.rs b/rust/src/persist/mod.rs index 2551a7af08..3c0207067b 100644 --- a/rust/src/persist/mod.rs +++ b/rust/src/persist/mod.rs @@ -31,6 +31,7 @@ //! - [`sliding_sync`]: what a chunk changes in the sliding sync tables. //! - [`state_delta`]: the room's new current state after a chunk. //! - [`sticky`]: the soft-failed sticky events a chunk lets through. +//! - [`txn`]: the chunk's 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. @@ -48,6 +49,7 @@ pub mod reader; pub mod sliding_sync; pub mod state_delta; pub mod sticky; +pub mod txn; /// Called when registering modules with python. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { diff --git a/rust/src/persist/python.rs b/rust/src/persist/python.rs index 25c12bf82a..e7d7e63cab 100644 --- a/rust/src/persist/python.rs +++ b/rust/src/persist/python.rs @@ -27,13 +27,16 @@ use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyDict}; use super::{ context::PersistContext, + effects::PersistEffects, extremities::MissingPrevStateEvents, metrics::{PersistCounts, PersistMetrics}, prepare::{self, PrepareInput, PrepareOptions, Prepared}, reader::{self, PersistReader, PyDelta, PySlidingSync}, state_delta::NativeStateResolver, + txn::{self, TxnEvent, TxnInput}, }; use crate::{ + database::postgres::write::{is_retryable_error, run_write_interaction, MAX_ATTEMPTS}, deferred::{create_deferred_with, OnCancel}, errors::SynapseError, event_auth::AuthFailure, @@ -86,7 +89,9 @@ pub(super) struct PersisterInner { impl Persister { /// Build the persister for `homeserver`, whose main database uses the /// native driver through `connection_pool` (a `RustConnectionPool`). - /// `None` if this instance does not write events. + /// `None` if this instance does not write events, or if the state tables + /// are in another database: persisting reads and writes them in the + /// main database's transactions. pub fn for_homeserver( py: Python<'_>, homeserver: &HomeServer, @@ -100,6 +105,16 @@ impl Persister { return Ok(None); } + if !homeserver + .state_database_pool(py)? + .is(&homeserver.main_database_pool(py)?) + { + log::warn!( + "The state tables are not in the main database: persisting events in Python" + ); + return Ok(None); + } + let config = homeserver.config(py)?; let main = homeserver.main_datastore(py)?; let tracker = |attr: &Bound<'_, pyo3::types::PyString>| -> PyResult> { @@ -165,15 +180,37 @@ impl Persister { sticky: inner.msc4354_enabled, dry_run: shadow, }; - let mut counts = PersistCounts::default(); - let prepared = - prepare::prepare_chunk(&reader, &resolver, &pool, &input, &options, &mut counts) - .await - .map_err(persist_error_to_py)?; - Ok(PreparedChunk { - prepared, - counts: Mutex::new(Some(counts)), - }) + // Python reads in `runInteraction`s, which retry on a lost + // connection or a cancelled statement. Here the whole chunk is + // prepared again. That is safe as its only write, bumping the + // state groups it uses, can be repeated, and runs in a write + // interaction with its own retries. + let mut attempt = 1; + loop { + let mut counts = PersistCounts::default(); + match prepare::prepare_chunk( + &reader, + &resolver, + &pool, + &input, + &options, + &mut counts, + ) + .await + { + Ok(prepared) => { + return Ok(PreparedChunk { + prepared, + counts: Mutex::new(Some(counts)), + }) + } + Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => { + log::warn!("[TXN RETRY] {{prepare_chunk}} {e:#} {attempt}/{MAX_ATTEMPTS}"); + attempt += 1; + } + Err(e) => return Err(persist_error_to_py(e)), + } + } } .boxed()) } @@ -231,6 +268,47 @@ impl Persister { } } + /// Run the part of the write transaction ported so far over the chunk, + /// whose events must have their stream orderings set, and commit it. + /// Returns what it collected for Python to apply and the positions in + /// the chunk of the events left for the Python steps after it. For + /// `tests/storage/test_rust_persist_txn.py`. + fn persist_events_txn_for_test( + &self, + py: Python<'_>, + room_id: String, + events_and_contexts: Vec<(Py, PersistContext)>, + ) -> PyResult<(PersistEffects, Vec)> { + let events = events_and_contexts + .into_iter() + .map(|(event, context)| TxnEvent::new(event.bind(py).get().shallow_clone(), context)) + .collect::>>() + .map_err(persist_error_to_py)?; + let pool = self.inner.pool.get()?; + let future = async move { + let input = TxnInput { + room_id: &room_id, + events: &events, + }; + run_write_interaction(&pool, "persist_events", |txn, _| { + let input = &input; + async move { + let mut effects = PersistEffects::default(); + let left = txn::persist_events_txn(txn, input, &mut effects).await?; + let left = left + .iter() + .filter_map(|e| input.events.iter().position(|x| std::ptr::eq(x, *e))) + .collect(); + Ok((effects, left)) + } + .boxed() + }) + .await + .map_err(|e| persist_error_to_py(anyhow::Error::new(e))) + }; + self.block_on(py, future.boxed()) + } + /// As [`Self::prepare_chunk`], but blocks. #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))] fn prepare_chunk_blocking( diff --git a/rust/src/persist/state_delta.rs b/rust/src/persist/state_delta.rs index 451a23f113..217ab5baf7 100644 --- a/rust/src/persist/state_delta.rs +++ b/rust/src/persist/state_delta.rs @@ -37,7 +37,8 @@ //! is the lowest group ID. Python takes the first in its set's order. //! - [`NativeStateResolver`] reads and bumps the state groups on the main //! database's pool. Python uses the state database. Like the Rust state -//! resolution store, this needs the two to be the same database. +//! resolution store, this needs the two to be the same database, which +//! `Persister::for_homeserver` checks. //! //! A step can fail with [`StateResFallback`] after it has added to the //! caller's [`PersistCounts`] and bumped the input groups. A caller that diff --git a/rust/src/persist/txn/mod.rs b/rust/src/persist/txn/mod.rs new file mode 100644 index 0000000000..1a75b1e76d --- /dev/null +++ b/rust/src/persist/txn/mod.rs @@ -0,0 +1,285 @@ +/* + * 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 write transaction of a chunk: `_persist_events_txn`. +//! +//! [`persist_events_txn`] runs inside a [`WriteTxn`] and takes what +//! [`super::prepare`] computed. It writes the rows Python's transaction +//! writes, in the same order, and collects in a [`PersistEffects`] the +//! cache updates Python queues with `txn.call_after`. +//! +//! This is being ported a group of steps at a time, from the start of the +//! transaction. So far it covers the steps up to dropping duplicate events, +//! and returns the events left at that point. The Python steps after it run +//! only in `tests/storage/test_rust_persist_txn.py`, which persists the +//! rest of a chunk in Python and compares the rows with Python's. + +use std::collections::HashMap; + +use anyhow::anyhow; +use serde_json::Value; + +use super::{context::PersistContext, effects::EventInvalidation, effects::PersistEffects}; +use crate::{ + database::postgres::write::{WriteError, WriteTxn}, + events::{ + constants::{event_field::M_RELATES_TO, relation_type}, + Event, + }, +}; + +/// One event of the chunk, with its stream ordering already set on its +/// internal metadata. +pub struct TxnEvent { + pub event: Event, + pub context: PersistContext, + pub stream_ordering: i64, +} + +impl TxnEvent { + pub fn new(event: Event, context: PersistContext) -> anyhow::Result { + let stream_ordering = event + .internal_metadata() + .stream_ordering() + .ok_or_else(|| anyhow!("{} has no stream ordering", event.event_id()))? + .get(); + Ok(TxnEvent { + event, + context, + stream_ordering, + }) + } + + fn event_id(&self) -> &str { + self.event.event_id() + } + + fn is_outlier(&self) -> Result { + self.event.internal_metadata().is_outlier().map_err(|e| { + anyhow::Error::new(e) + .context(format!("reading {}", self.event_id())) + .into() + }) + } +} + +/// What the transaction takes besides the chunk. +pub struct TxnInput<'a> { + pub room_id: &'a str, + /// The chunk, in order. It may hold the same event twice. + pub events: &'a [TxnEvent], +} + +/// Write the chunk, as `_persist_events_txn` does, adding to `effects` what +/// Python would queue with `txn.call_after`. Returns the events left once +/// duplicates are dropped, which the steps still to be ported take. +/// +/// `effects` must start empty: a retried transaction starts again. +pub async fn persist_events_txn<'a>( + txn: &WriteTxn, + input: &TxnInput<'a>, + effects: &mut PersistEffects, +) -> Result, WriteError> { + // Protect against a race with deleting the room: the share lock holds + // off a delete until we commit. + let rows = txn + .query( + "SELECT room_version FROM rooms WHERE room_id = $1 FOR SHARE", + &[&input.room_id], + ) + .await?; + if rows.is_empty() { + return Err(anyhow!("Room does not exist {}", input.room_id).into()); + } + + // Recorded for every event of the chunk, before any are dropped. + for event in input.events { + add_event_effects(event, effects); + } + + filter_duplicates(input.events) +} + +/// The cache updates Python queues for each event of the chunk at the start +/// of the transaction. +fn add_event_effects(event: &TxnEvent, effects: &mut PersistEffects) { + let redacts = redacts(&event.event); + effects.add_event_cache_invalidation(event.event_id().to_owned()); + if let Some(redacts) = redacts.filter(|r| !r.is_empty()) { + effects.add_event_cache_invalidation(redacts.to_owned()); + } + effects.add_event(EventInvalidation { + stream_ordering: event.stream_ordering, + event_id: event.event_id().to_owned(), + room_id: event.event.room_id().to_owned(), + event_type: event.event.type_().to_owned(), + state_key: event.event.state_key().map(str::to_owned), + redacts: redacts.map(str::to_owned), + relates_to: relation(&event.event).map(|r| r.parent_id.to_owned()), + // Python passes `False` for backfilled events too. + backfilled: false, + }); +} + +/// `_filter_events_and_contexts_for_duplicates`: one of each event, the +/// first that is not an outlier if there is one, else the first, in the +/// order each event was first seen. +fn filter_duplicates(events: &[TxnEvent]) -> Result, WriteError> { + // Slots in first-seen order. A replaced event leaves an empty slot. + let mut slots: Vec> = Vec::with_capacity(events.len()); + let mut index_of: HashMap<&str, usize> = HashMap::with_capacity(events.len()); + for event in events { + match index_of.get(event.event_id()) { + None => { + index_of.insert(event.event_id(), slots.len()); + slots.push(Some(event)); + } + Some(&index) => { + let kept = slots[index].expect("the slot of a kept event is full"); + // Python pops the outlier and appends the event, so it moves + // to the end. + if !event.is_outlier()? && kept.is_outlier()? { + slots[index] = None; + index_of.insert(event.event_id(), slots.len()); + slots.push(Some(event)); + } + } + } + } + Ok(slots.into_iter().flatten().collect()) +} + +/// `event.redacts`, if it is a string. Python passes on any value, but an +/// event whose `redacts` is not a string fails validation before it gets +/// here. +pub(super) fn redacts(event: &Event) -> Option<&str> { + event.redacts().and_then(Value::as_str) +} + +/// What `relation_from_event` parses from an event's `m.relates_to`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct Relation<'a> { + pub parent_id: &'a str, + pub rel_type: &'a str, + /// For an annotation, its `key` if it is a string. + pub aggregation_key: Option<&'a str>, +} + +/// `relation_from_event`: the event's relation, if it is well formed. +pub(super) fn relation(event: &Event) -> Option> { + let relation = event.content().get_field(M_RELATES_TO)?.as_object()?; + let rel_type = relation.get("rel_type")?.as_str()?; + let parent_id = relation.get("event_id")?.as_str()?; + let aggregation_key = if rel_type == relation_type::ANNOTATION { + relation.get("key").and_then(Value::as_str) + } else { + None + }; + Some(Relation { + parent_id, + rel_type, + aggregation_key, + }) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + use crate::{events::internal_metadata::EventInternalMetadata, room_versions::RoomVersion}; + + fn event(body: &str, content: Value, outlier: bool) -> Event { + let json = json!({ + "auth_events": [], "prev_events": [], "type": "m.room.message", + "sender": "@a:example.com", "content": content, "depth": 3, + "room_id": "!r:example.com", "origin_server_ts": 5, + "hashes": {"sha256": body}, "signatures": {}, "unsigned": {}, + }); + let metadata = EventInternalMetadata::from_json_str("{}").unwrap(); + metadata.set_outlier_flag(outlier).unwrap(); + metadata + .set_persisted_position(Some(1), None, None) + .unwrap(); + Event::from_json_str(&json.to_string(), &RoomVersion::V10, metadata, None).unwrap() + } + + fn txn_event(body: &str, outlier: bool) -> TxnEvent { + TxnEvent::new(event(body, json!({}), outlier), PersistContext::default()).unwrap() + } + + #[test] + fn duplicates_keep_the_first_non_outlier() { + // The same event three times, as an outlier first, and another. + let events = [ + txn_event("a", true), + txn_event("b", false), + txn_event("a", false), + txn_event("a", false), + ]; + let kept = filter_duplicates(&events).unwrap(); + // The first non-outlier copy of `a` replaces the outlier and moves + // to the end. + assert_eq!(kept.len(), 2); + assert!(std::ptr::eq(kept[0], &events[1])); + assert!(std::ptr::eq(kept[1], &events[2])); + + // Without a non-outlier copy, the first is kept in place. + let events = [txn_event("a", true), txn_event("a", true)]; + let kept = filter_duplicates(&events).unwrap(); + assert_eq!(kept.len(), 1); + assert!(std::ptr::eq(kept[0], &events[0])); + } + + #[test] + fn relation_is_parsed_as_python_does() { + let relation_of = |relates_to: Value| { + let event = event("a", json!({ "m.relates_to": relates_to }), false); + relation(&event).map(|r| { + ( + r.parent_id.to_owned(), + r.rel_type.to_owned(), + r.aggregation_key.map(str::to_owned), + ) + }) + }; + assert_eq!( + relation_of(json!({"rel_type": "m.thread", "event_id": "$p"})), + Some(("$p".into(), "m.thread".into(), None)) + ); + assert_eq!( + relation_of(json!({"rel_type": "m.annotation", "event_id": "$p", "key": "👍"})), + Some(("$p".into(), "m.annotation".into(), Some("👍".into()))) + ); + // Only an annotation has a key, and only a string one. + assert_eq!( + relation_of(json!({"rel_type": "m.reference", "event_id": "$p", "key": "k"})), + Some(("$p".into(), "m.reference".into(), None)) + ); + assert_eq!( + relation_of(json!({"rel_type": "m.annotation", "event_id": "$p", "key": 1})), + Some(("$p".into(), "m.annotation".into(), None)) + ); + // Malformed relations are ignored. + assert_eq!(relation_of(json!({"rel_type": "m.thread"})), None); + assert_eq!(relation_of(json!({"rel_type": 1, "event_id": "$p"})), None); + assert_eq!( + relation_of(json!({"rel_type": "m.thread", "event_id": 1})), + None + ); + assert_eq!(relation_of(json!(["$p"])), None); + assert_eq!(relation_of(json!({})), None); + } +} diff --git a/rust/src/state/store.rs b/rust/src/state/store.rs index 354b9a3cdf..26f8b5b0df 100644 --- a/rust/src/state/store.rs +++ b/rust/src/state/store.rs @@ -137,7 +137,7 @@ impl PostgresStateResStore { let conn = pool .get() .await - .map_err(|e| anyhow::anyhow!("checking out a database connection: {e}"))?; + .map_err(|e| anyhow::Error::new(e).context("checking out a database connection"))?; let store = PostgresStateResStore { conn: Some(conn), handle, diff --git a/rust/src/storage/sync_store.rs b/rust/src/storage/sync_store.rs index 554c84e4ca..ca43c406f8 100644 --- a/rust/src/storage/sync_store.rs +++ b/rust/src/storage/sync_store.rs @@ -250,7 +250,9 @@ impl<'a> SyncReader<'a> { .pool .get() .await - .map_err(|e| anyhow::anyhow!("checking out a database connection: {e}"))?; + // Keep the pool's error as the cause, so that a caller can tell + // a lost connection, which is worth retrying. + .map_err(|e| anyhow::Error::new(e).context("checking out a database connection"))?; self.update_stats(|stats| { stats.checkouts += 1; stats.checkout_time += start.elapsed(); diff --git a/synapse/storage/controllers/persist_events.py b/synapse/storage/controllers/persist_events.py index 9df5f0387f..02a7ebb8fa 100644 --- a/synapse/storage/controllers/persist_events.py +++ b/synapse/storage/controllers/persist_events.py @@ -598,7 +598,8 @@ class EventsPersistenceStorageController: self.hs = hs # The Rust persister, at `rust_event_persistence` levels other than - # "off". It needs the native Rust Postgres driver. + # "off". It needs the native Rust Postgres driver, with the state + # tables in the main database. self._rust_persister: Persister | None = None rust_level = hs.config.experimental.rust_event_persistence # Whether a Rust error, or in shadow mode a mismatch between the two @@ -609,7 +610,8 @@ class EventsPersistenceStorageController: if self._rust_persister is None: logger.warning( "rust_event_persistence is %r but the main database does not " - "use the native Rust driver: persisting events in Python", + "use the native Rust driver or does not hold the state tables: " + "persisting events in Python", rust_level, ) else: @@ -924,12 +926,15 @@ class EventsPersistenceStorageController: ) -> _PreparedChunk: """Compute what persisting the chunk needs, in Python or Rust as `rust_event_persistence` says.""" - if self._rust_persister is None: + # Read per chunk, so that it can be lowered at runtime, as tests do. + # Raising it from "off" has no effect: the persister only exists if + # the level was not "off" at startup. + level = self.hs.config.experimental.rust_event_persistence + if self._rust_persister is None or level == "off": return await self._prepare_chunk_in_python( room_id, chunk, backfilled, table_changes=False ) - level = self.hs.config.experimental.rust_event_persistence if level == "shadow": return await self._prepare_chunk_in_shadow_mode(room_id, chunk, backfilled) diff --git a/synapse/synapse_rust/persist.pyi b/synapse/synapse_rust/persist.pyi index 6ea135ee4e..a3a5019461 100644 --- a/synapse/synapse_rust/persist.pyi +++ b/synapse/synapse_rust/persist.pyi @@ -51,6 +51,15 @@ class Persister: shadow: bool, ) -> PreparedChunk: """As `prepare_chunk`, but blocks. For tests.""" + def persist_events_txn_for_test( + self, + room_id: str, + events_and_contexts: list[tuple[EventBase, EventContext]], + ) -> tuple[PersistEffects, list[int]]: + """Run the part of the write transaction ported to Rust so far over + the chunk, whose events have their stream orderings set, and commit + it. Returns what it collected for Python to apply and the positions + in the chunk of the events left for the Python steps after it.""" class PreparedChunk: """What `Persister.prepare_chunk` computed. If `state_res_fallback` is diff --git a/tests/storage/test_rust_persist_prepare.py b/tests/storage/test_rust_persist_prepare.py index 9a206a1ebf..ffa11e747c 100644 --- a/tests/storage/test_rust_persist_prepare.py +++ b/tests/storage/test_rust_persist_prepare.py @@ -412,6 +412,15 @@ class PrepareTestCase(_PrepareTestBase): assert stream_ordering is not None self.assertLess(stream_ordering, 0) + def test_off_uses_python(self) -> None: + """The level is read per chunk, and "off" skips Rust even though + the persister exists.""" + self.hs.config.experimental.rust_event_persistence = "off" + persister = Mock() + self.persistence._rust_persister = persister + self._send("hello") + persister.prepare_chunk.assert_not_called() + def test_rejection_is_raised(self) -> None: """A chunk Rust rejects is not computed again in Python.""" self.persistence._rust_persist_strict = False diff --git a/tests/storage/test_rust_persist_txn.py b/tests/storage/test_rust_persist_txn.py new file mode 100644 index 0000000000..605ff09d62 --- /dev/null +++ b/tests/storage/test_rust_persist_txn.py @@ -0,0 +1,852 @@ +# +# 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 write transaction (`rust/src/persist/txn/`) against Python's. + +Each test builds the same DAG in two rooms, the twins. Room A's chunks are +persisted in Python. Room B's are prepared in Rust, then the part of the +transaction ported to Rust so far runs and commits, then the rest of +Python's transaction steps run. The test then compares the rows of every +table the transaction writes, after normalising the IDs that differ between +the rooms, and the cache updates each room's persists asked for. + +These need the native Rust Postgres driver: + + SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.storage.test_rust_persist_txn +""" + +import json +import re +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Any +from unittest import SkipTest + +import attr + +from twisted.internet.testing import MemoryReactor + +from synapse.api.constants import EventTypes +from synapse.api.room_versions import KNOWN_ROOM_VERSIONS +from synapse.events import make_event_from_dict +from synapse.events.snapshot import EventContext, EventPersistencePair +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.events import PersistEventsStore +from synapse.storage.rust_pool import RustConnectionPool +from synapse.synapse_rust.persist import PersistEffects +from synapse.types import create_requester +from synapse.util.caches.deferred_cache import DeferredCache +from synapse.util.clock import Clock + +from tests import unittest + +# The tables `_persist_events_txn` writes. Each is compared on the rows of +# the room, by its `room_id` column if it has one, else by `event_id`. +_TABLES = ( + "current_state_delta_stream", + "current_state_events", + "destination_rooms_sticky_events_backlog", + "event_auth", + "event_auth_chain_links", + "event_auth_chain_to_calculate", + "event_auth_chains", + "event_backward_extremities", + "event_edges", + "event_expiry", + "event_failed_pull_attempts", + "event_forward_extremities", + "event_json", + "event_labels", + "event_push_actions", + "event_push_actions_staging", + "event_relations", + "event_search", + "event_to_state_groups", + "event_txn_id_device_id", + "events", + "ex_outlier_stream", + "local_current_membership", + "msc4242_state_dag_edges", + "msc4242_state_dag_forward_extremities", + "partial_state_events", + "redactions", + "rejections", + "room_depth", + "room_memberships", + "room_retention", + "rooms", + "sliding_sync_joined_rooms", + "sliding_sync_membership_snapshots", + "state_events", + "sticky_events", + "stream_ordering_to_exterm", + "threads", + "timeline_gaps", +) + +# Integer columns holding a position on a stream or a state group, by +# column name or `(table, column)`, and which sequence the value is from. +# They differ between the twins, so they are compared by their rank among +# the twin's values from that sequence. Columns not listed are compared as +# they are. +_STREAM_COLUMNS: dict[str | tuple[str, str], str] = { + "bump_stamp": "events", + "event_stream_ordering": "events", + "stream_ordering": "events", + ("current_state_delta_stream", "stream_id"): "events", + ("ex_outlier_stream", "event_stream_ordering"): "ex_outlier", + ("sticky_events", "stream_id"): "sticky", + "state_group": "state_groups", +} + +# Chain IDs, compared by the lowest event ID in the chain. +_CHAIN_COLUMNS = { + ("event_auth_chains", "chain_id"), + ("event_auth_chain_links", "origin_chain_id"), + ("event_auth_chain_links", "target_chain_id"), +} + +# Parsed and compared as JSON. The hashes and signatures of an event cover +# its room ID, so they differ between the twins. +_JSON_COLUMNS = { + ("event_json", "json"), + ("event_json", "internal_metadata"), +} + +# Which positional arguments of a `call_after` callback hold a stream +# position or state group, by callback name. +_EFFECT_ARGS = { + "_invalidate_caches_for_event": ("events",), + "entity_has_changed": (None, "events"), + "prefill": (None, "state_groups"), + "_mark_ids_as_finished": ("cache",), +} + +# Marks each twin's create event, so the rooms get different IDs in room +# versions that derive the room ID from it. +_TWIN_KEY = "org.matrix.synapse.test.twin" + + +@attr.s(auto_attribs=True) +class _Twin: + """One of the two rooms.""" + + name: str + room_id: str = "" + # The scenario's events, by index. + events: list[str] = attr.Factory(list) + # The cache updates its persists asked for: `(callback, owner, args, + # kwargs)`, with stream positions as they are. + effects: list[tuple[str, str | None, tuple, dict[str, Any]]] = attr.Factory(list) + # The cache stream rows its persists wrote. + cache_rows: list[tuple] = attr.Factory(list) + + +class _Normaliser: + """Rewrites a twin's values so they compare equal with the other + twin's: the room ID, event IDs, stream positions, state groups and + chain IDs.""" + + def __init__( + self, + twin: _Twin, + event_ids: dict[str, str], + ranks: dict[str, dict[int, int]], + ) -> None: + self._ranks = ranks + names = {twin.room_id: "!room", f"twin-{twin.name}": "twin", **event_ids} + self._names = names + self._pattern = re.compile( + "|".join(re.escape(k) for k in sorted(names, key=len, reverse=True)) + ) + self.chains: dict[int, str] = {} + + def string(self, value: str) -> str: + return self._pattern.sub(lambda m: self._names[m.group(0)], value) + + def json(self, value: Any) -> Any: + if isinstance(value, str): + return self.string(value) + if isinstance(value, list): + return [self.json(v) for v in value] + if isinstance(value, dict): + return {self.string(k): self.json(v) for k, v in value.items()} + return value + + def position(self, kind: str | None, value: Any) -> Any: + if kind is None or value is None: + return value + if isinstance(value, (list, tuple)): + return tuple(self.position(kind, v) for v in value) + rank = self._ranks.get(kind, {}).get(value) + return f"{kind}#{rank}" if rank is not None else f"{kind}?{value}" + + def value(self, table: str, column: str, value: Any) -> Any: + if (table, column) in _CHAIN_COLUMNS and value is not None: + return self.chains.get(value, f"chain?{value}") + if (table, column) in _JSON_COLUMNS and value is not None: + parsed = json.loads(value) + if column == "json": + parsed.pop("hashes", None) + parsed.pop("signatures", None) + return json.dumps(self.json(parsed), sort_keys=True) + kind = _STREAM_COLUMNS.get((table, column), _STREAM_COLUMNS.get(column)) + if isinstance(value, int) and not isinstance(value, bool) and kind: + return self.position(kind, value) + if isinstance(value, (list, tuple)): + return tuple(self.value(table, column, v) for v in value) + if isinstance(value, str): + return self.string(value) + return value + + +class _TwinTestBase(unittest.HomeserverTestCase): + """Builds DAGs in twin rooms and compares them. See the module docs.""" + + maxDiff = None + + servlets = [ + admin.register_servlets, + room.register_servlets, + login.register_servlets, + ] + + room_version = "12" + + def default_config(self) -> dict[str, Any]: + config = super().default_config() + experimental = config.setdefault("experimental_features", {}) + experimental["rust_event_persistence"] = "prepare" + experimental["msc4354_enabled"] = True + return config + + def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer: + 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") + persistence = hs.get_storage_controllers().persistence + assert persistence is not None + if persistence._rust_persister is None: + raise SkipTest("SYNAPSE_TEST_RUST_PERSIST=off turns the persister off") + return hs + + def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None: + self.store = hs.get_datastores().main + persistence = hs.get_storage_controllers().persistence + assert persistence is not None + self.persistence = persistence + self.persistence._rust_persist_strict = True + persister = persistence._rust_persister + assert persister is not None + self.persister = persister + events_store = hs.get_datastores().persist_events + assert events_store is not None + self.events_store = events_store + + self.alice = self.register_user("alice", "pass") + self.bob = self.register_user("bob", "pass") + self.twins = (_Twin("a"), _Twin("b")) + for twin in self.twins: + with self._persisting(twin): + twin.room_id, _, _ = self.get_success( + self.hs.get_room_creation_handler().create_room( + create_requester(self.alice), + { + "room_version": self.room_version, + "creation_content": {_TWIN_KEY: f"twin-{twin.name}"}, + }, + ratelimit=False, + ) + ) + + # Persisting. + + @contextmanager + def _persisting(self, twin: _Twin) -> Iterator[None]: + """Persist the twin's events: room A's in Python, room B's through + Rust as far as it goes. Records the cache updates and cache stream + rows.""" + self.hs.config.experimental.rust_event_persistence = ( + "off" if twin.name == "a" else "prepare" + ) + store = self.events_store + python_txn = PersistEventsStore._persist_events_txn.__get__(store) + if twin.name == "a": + + def txn(t: LoggingTransaction, **kwargs: Any) -> None: + start = len(t.after_callbacks or []) + python_txn(t, **kwargs) + twin.effects.extend(self._python_effects(t, start)) + else: + + def txn(t: LoggingTransaction, **kwargs: Any) -> None: + self._rust_txn(twin, t, **kwargs) + + cache_stream_before = self._cache_stream_position() + store._persist_events_txn = txn # type: ignore[method-assign,assignment] + try: + yield + finally: + del store._persist_events_txn + self.hs.config.experimental.rust_event_persistence = "prepare" + twin.cache_rows.extend(self._cache_stream_rows(cache_stream_before)) + + def _rust_txn( + self, + twin: _Twin, + txn: LoggingTransaction, + *, + room_id: str, + events_and_contexts: list[EventPersistencePair], + **kwargs: Any, + ) -> None: + """Room B's transaction: the Rust part, committed first, then the + rest in Python.""" + effects, left = self.persister.persist_events_txn_for_test( + room_id, events_and_contexts + ) + twin.effects.extend(_rust_effects(effects)) + start = len(txn.after_callbacks or []) + _python_steps_after_rust( + self.events_store, + txn, + room_id=room_id, + all_events_and_contexts=events_and_contexts, + events_and_contexts=[events_and_contexts[i] for i in left], + **kwargs, + ) + twin.effects.extend(self._python_effects(txn, start)) + + def _python_effects( + self, txn: LoggingTransaction, start: int + ) -> list[tuple[str, str | None, tuple, dict[str, Any]]]: + effects = [] + for callback, args, kwargs in (txn.after_callbacks or [])[start:]: + name = callback.__name__ + owner = getattr(callback, "__self__", None) + # A single-key cache's `prefill` is a lambda over the cache. + cache = next( + ( + cell.cell_contents + for cell in getattr(callback, "__closure__", None) or () + if isinstance(cell.cell_contents, DeferredCache) + ), + None, + ) + if owner is self.store: + owner_name = None + elif owner is not None: + owner_name = getattr(owner, "name", None) or type(owner).__name__ + elif name == "" and cache is not None: + # The transaction prefills only `_get_state_group_for_event`. + name, owner_name = "prefill", "_get_state_group_for_event" + else: + owner_name = callback.__qualname__ + effects.append((name, owner_name, args, kwargs)) + return effects + + def _cache_stream_position(self) -> int: + rows = self.get_success( + self.store.db_pool.execute( + "test", + "SELECT COALESCE(MAX(stream_id), 0) FROM cache_invalidation_stream_by_instance", + ) + ) + return rows[0][0] + + def _cache_stream_rows(self, after: int) -> list[tuple]: + return self.get_success( + self.store.db_pool.execute( + "test", + """SELECT stream_id, instance_name, cache_func, keys + FROM cache_invalidation_stream_by_instance + WHERE stream_id > ? ORDER BY stream_id""", + after, + ) + ) + + # Building the DAG. + + def _build( + self, + twin: _Twin, + event_type: str, + content: dict[str, Any] | Callable[[_Twin], dict[str, Any]], + *, + state_key: str | None = None, + prev: list[int] | None = None, + redacts: int | None = None, + sender: str | None = None, + ) -> EventPersistencePair: + """An event of the twin's, not yet persisted, added to its events. + `prev` and `redacts` are indices of the twin's events. `content` can + be a function of the twin, for content that names its events.""" + sender = sender or self.alice + if callable(content): + content = content(twin) + event_dict: dict[str, Any] = { + "type": event_type, + "content": dict(content), + "room_id": twin.room_id, + "sender": sender, + } + if state_key is not None: + event_dict["state_key"] = state_key + if redacts is not None: + redacts_id = twin.events[redacts] + event_dict["redacts"] = redacts_id + if KNOWN_ROOM_VERSIONS[self.room_version].updated_redaction_rules: + event_dict["content"]["redacts"] = redacts_id + event, unpersisted_context = self.get_success( + self.hs.get_event_creation_handler().create_event( + create_requester(sender), + event_dict, + prev_event_ids=( + [twin.events[i] for i in prev] if prev is not None else None + ), + ) + ) + context = self.get_success(unpersisted_context.persist(event)) + twin.events.append(event.event_id) + return event, context + + def _step( + self, + build: Callable[[_Twin], list[EventPersistencePair]], + *, + backfilled: bool = False, + ) -> None: + """Build a chunk in each twin with `build`, and persist it.""" + for twin in self.twins: + chunk = build(twin) + with self._persisting(twin): + self.get_success( + self.persistence.persist_events(chunk, backfilled=backfilled) + ) + self.assertEqual(len(self.twins[0].events), len(self.twins[1].events)) + + def _send( + self, + event_type: str, + content: dict[str, Any] | Callable[[_Twin], dict[str, Any]], + **kwargs: Any, + ) -> int: + """Persist one event in each twin. Returns its index.""" + self._step(lambda twin: [self._build(twin, event_type, content, **kwargs)]) + return len(self.twins[0].events) - 1 + + def _message(self, body: str, **kwargs: Any) -> int: + return self._send( + EventTypes.Message, {"msgtype": "m.text", "body": body}, **kwargs + ) + + def _outlier(self, pair: EventPersistencePair) -> EventPersistencePair: + """A copy of the event as an outlier. Outliers have no state, except + in a state DAG room, where `/get_missing_events` persists them with + it.""" + event, context = pair + copy = make_event_from_dict( + event.get_pdu_json(), + event.room_version, + event.internal_metadata.get_dict(), + ) + copy.internal_metadata.outlier = True + if event.room_version.msc4242_state_dags: + return copy, context + return copy, EventContext.for_outlier(self.hs.get_storage_controllers()) + + # Comparing. + + def _dump(self, twin: _Twin) -> dict[str, list[tuple]]: + """The twin's rows of each table, normalised and sorted.""" + raw: dict[str, tuple[list[str], list[tuple]]] = {} + for table in _TABLES: + columns = self._columns(table) + if "room_id" in columns: + where = "room_id = ?" + elif table == "event_auth_chain_links": + where = """origin_chain_id IN ( + SELECT chain_id FROM event_auth_chains + INNER JOIN events USING (event_id) WHERE room_id = ? + )""" + else: + where = "event_id IN (SELECT event_id FROM events WHERE room_id = ?)" + # The driver cannot decode a `tsvector`. + selected = [ + f"{column}::text" if data_type == "tsvector" else column + for column, data_type in zip(columns, self._types[table]) + ] + rows = self.get_success( + self.store.db_pool.execute( + "test", + f"SELECT {', '.join(selected)} FROM {table} WHERE {where}", + twin.room_id, + ) + ) + raw[table] = (columns, rows) + + normaliser = self._normaliser(twin, raw) + dump = {} + for table, (columns, rows) in raw.items(): + dump[table] = sorted( + ( + tuple( + normaliser.value(table, column, value) + for column, value in zip(columns, row) + ) + for row in rows + ), + key=repr, + ) + dump["cache_invalidation_stream_by_instance"] = [ + (instance_name, cache_func, normaliser.value("", "", keys)) + for _, instance_name, cache_func, keys in twin.cache_rows + ] + dump["effects"] = sorted( + ( + ( + name, + owner, + tuple( + normaliser.position(kind, arg) + if kind + else normaliser.value("", "", arg) + for kind, arg in zip( + _EFFECT_ARGS.get(name, ()) + (None,) * len(args), args + ) + ), + tuple(sorted(kwargs.items())), + ) + for name, owner, args, kwargs in twin.effects + ), + key=repr, + ) + return dump + + def _normaliser( + self, twin: _Twin, raw: dict[str, tuple[list[str], list[tuple]]] + ) -> _Normaliser: + # The scenario's events by index, then the rest (the room's + # creation) by stream ordering. + event_ids = {event_id: f"$e{i}" for i, event_id in enumerate(twin.events)} + columns, rows = raw["events"] + by_ordering = sorted( + rows, key=lambda row: row[columns.index("stream_ordering")] + ) + for row in by_ordering: + event_id = row[columns.index("event_id")] + event_ids.setdefault(event_id, f"$s{len(event_ids)}") + + values: dict[str, set[int]] = {} + for table, (columns, rows) in raw.items(): + for i, column in enumerate(columns): + kind = _STREAM_COLUMNS.get((table, column), _STREAM_COLUMNS.get(column)) + if not kind: + continue + values.setdefault(kind, set()).update( + row[i] for row in rows if isinstance(row[i], int) + ) + values.setdefault("cache", set()).update(row[0] for row in twin.cache_rows) + ranks = { + kind: {value: rank for rank, value in enumerate(sorted(found))} + for kind, found in values.items() + } + normaliser = _Normaliser(twin, event_ids, ranks) + + columns, rows = raw["event_auth_chains"] + for row in rows: + chain_id = row[columns.index("chain_id")] + name = normaliser.string(row[columns.index("event_id")]) + if chain_id not in normaliser.chains or name < normaliser.chains[chain_id]: + normaliser.chains[chain_id] = name + return normaliser + + _column_cache: dict[str, list[str]] = {} + _types: dict[str, list[str]] = {} + + def _columns(self, table: str) -> list[str]: + if table not in self._column_cache: + rows = self.get_success( + self.store.db_pool.execute( + "test", + """SELECT column_name, data_type FROM information_schema.columns + WHERE table_name = ? ORDER BY ordinal_position""", + table, + ) + ) + self.assertTrue(rows, f"no table {table}") + self._column_cache[table] = [row[0] for row in rows] + self._types[table] = [row[1] for row in rows] + return self._column_cache[table] + + def _assert_twins_match(self) -> None: + a, b = (self._dump(twin) for twin in self.twins) + differences = [table for table in a if a[table] != b[table]] + for table in differences: + self.assertEqual(a[table], b[table], f"{table} differs") + + +def _rust_effects( + effects: PersistEffects, +) -> list[tuple[str, str | None, tuple, dict[str, Any]]]: + """Rust's effects as the Python transaction's `call_after` callbacks.""" + calls: list[tuple[str, str | None, tuple, dict[str, Any]]] = [] + for *args, backfilled in effects.events: + calls.append( + ( + "_invalidate_caches_for_event", + None, + tuple(args), + {"backfilled": backfilled}, + ) + ) + for event_id in effects.event_cache_invalidations: + calls.append(("_invalidate_local_get_event_cache", None, (event_id,), {})) + return calls + + +def _python_steps_after_rust( + store: PersistEventsStore, + txn: LoggingTransaction, + *, + room_id: str, + all_events_and_contexts: list[EventPersistencePair], + events_and_contexts: list[EventPersistencePair], + inhibit_local_membership_updates: bool, + state_delta_for_room: Any, + new_forward_extremities: set[str] | None, + new_event_links: Any, + sliding_sync_table_changes: Any, + new_state_dag_forward_extremities: set[str] | None = None, + sticky_events_to_un_soft_fail: Any = frozenset(), +) -> None: + """The steps of `_persist_events_txn` that Rust does not do yet, as it + runs them. `events_and_contexts` are the events Rust left. + + A copy, so that the production transaction need not be split where the + Rust port has got to. Each step ported to Rust leaves this copy, and it + goes once Rust runs the whole transaction. A step added to Python's + transaction meanwhile runs in room A only, so its rows differ and the + comparison fails.""" + min_stream_order = all_events_and_contexts[0][0].internal_metadata.stream_ordering + room_version = all_events_and_contexts[0][0].room_version + + store._update_room_depths_txn(txn, room_id, events_and_contexts=events_and_contexts) + events_and_contexts = store._update_outliers_txn( + txn, events_and_contexts=events_and_contexts + ) + events_and_contexts = store._apply_existing_redaction_txn( + txn, room_id, room_version, events_and_contexts=events_and_contexts + ) + redacted = store._compute_newly_redacted_event_ids_txn( + txn, room_id, room_version, events_and_contexts=events_and_contexts + ) + if redacted: + store.store.delete_sticky_events_txn(txn, redacted) + store._store_event_txn(txn, events_and_contexts=events_and_contexts) + if new_forward_extremities: + max_stream_order = all_events_and_contexts[-1][ + 0 + ].internal_metadata.stream_ordering + assert max_stream_order is not None + store._update_forward_extremities_txn( + txn, + room_id, + new_forward_extremities=new_forward_extremities, + max_stream_order=max_stream_order, + ) + if new_state_dag_forward_extremities: + store._set_state_dag_extremities_txn( + txn, room_id, new_state_dag_forward_extremities + ) + store._persist_transaction_ids_txn(txn, events_and_contexts) + store._store_event_state_mappings_txn(txn, events_and_contexts) + store._persist_event_auth_chain_txn( + txn, [e for e, _ in events_and_contexts], new_event_links + ) + events_and_contexts = store._store_rejected_events_txn( + txn, events_and_contexts=events_and_contexts + ) + store._update_metadata_tables_txn( + txn, + events_and_contexts=events_and_contexts, + all_events_and_contexts=all_events_and_contexts, + inhibit_local_membership_updates=inhibit_local_membership_updates, + ) + if state_delta_for_room: + assert sliding_sync_table_changes is not None + assert min_stream_order is not None + store._update_current_state_txn( + txn, + room_id, + state_delta_for_room, + min_stream_order, + sliding_sync_table_changes, + ) + if store._msc4354_enabled: + store.store.insert_sticky_events_txn(txn, [ev for ev, _ in events_and_contexts]) + if sticky_events_to_un_soft_fail: + store.store.un_soft_fail_sticky_events_txn( + txn, sticky_events_to_un_soft_fail + ) + store._update_sliding_sync_tables_with_new_persisted_events_txn( + txn, room_id, events_and_contexts + ) + + +class PersistTxnTestCase(_TwinTestBase): + def test_room_creation(self) -> None: + self._assert_twins_match() + + def test_messages(self) -> None: + for body in ("one", "two", "three"): + self._message(body) + self._assert_twins_match() + + def test_state_changes(self) -> None: + self._send(EventTypes.Topic, {"topic": "Topic"}, state_key="") + self._send(EventTypes.Name, {"name": "Name"}, state_key="") + self._send( + EventTypes.Member, + {"membership": "invite"}, + state_key=self.bob, + ) + self._send( + EventTypes.Member, + {"membership": "join"}, + state_key=self.bob, + sender=self.bob, + ) + self._message("hello", sender=self.bob) + self._assert_twins_match() + + def test_fork(self) -> None: + base = self._message("base") + self._send(EventTypes.Topic, {"topic": "A"}, state_key="", prev=[base]) + self._send(EventTypes.Topic, {"topic": "B"}, state_key="", prev=[base]) + # Joins the two branches, resolving state. + self._message("merge") + self._assert_twins_match() + + def test_redaction_after_target(self) -> None: + target = self._message("secret") + self._send(EventTypes.Redaction, {}, redacts=target) + self._assert_twins_match() + + def test_redaction_before_target(self) -> None: + base = self._message("base") + built: dict[str, EventPersistencePair] = {} + + def target(twin: _Twin) -> list[EventPersistencePair]: + # Built, but persisted after its redaction. + built[twin.name] = self._build( + twin, EventTypes.Message, {"msgtype": "m.text", "body": "secret"} + ) + return [] + + self._step(target) + target_index = len(self.twins[0].events) - 1 + self._send(EventTypes.Redaction, {}, redacts=target_index, prev=[base]) + self._step(lambda twin: [built[twin.name]]) + self._assert_twins_match() + + def test_relation(self) -> None: + root = self._message("root") + self._send( + EventTypes.Message, + lambda twin: { + "msgtype": "m.text", + "body": "reply", + "m.relates_to": {"rel_type": "m.thread", "event_id": twin.events[root]}, + }, + ) + self._assert_twins_match() + + def test_rejected(self) -> None: + def rejected(twin: _Twin) -> list[EventPersistencePair]: + event, context = self._build( + twin, EventTypes.Topic, {"topic": "no"}, state_key="" + ) + context.rejected = "auth_error" + return [(event, context)] + + self._step(rejected) + self._message("after") + self._assert_twins_match() + + def test_outlier_then_de_outlier(self) -> None: + built: dict[str, EventPersistencePair] = {} + + def outlier(twin: _Twin) -> list[EventPersistencePair]: + built[twin.name] = self._build( + twin, EventTypes.Topic, {"topic": "outlier"}, state_key="" + ) + return [self._outlier(built[twin.name])] + + # Federation persists outliers as backfilled. + self._step(outlier, backfilled=True) + self._step(lambda twin: [built[twin.name]]) + self._assert_twins_match() + + def test_duplicates_in_chunk(self) -> None: + def duplicated(twin: _Twin) -> list[EventPersistencePair]: + pair = self._build(twin, EventTypes.Topic, {"topic": "twice"}, state_key="") + return [self._outlier(pair), pair] + + self._step(duplicated) + self._assert_twins_match() + + def test_backfill(self) -> None: + self._message("now") + self._step( + lambda twin: [ + self._build( + twin, EventTypes.Message, {"msgtype": "m.text", "body": "old"} + ) + ], + backfilled=True, + ) + self._assert_twins_match() + + def test_mismatch_is_found(self) -> None: + """The comparison sees a row that differs.""" + self._message("hello") + self.get_success( + self.store.db_pool.simple_update( + "events", + {"event_id": self.twins[1].events[0]}, + {"depth": 1000}, + desc="test", + ) + ) + with self.assertRaises(AssertionError): + self._assert_twins_match() + + def test_missing_room(self) -> None: + with self.assertRaisesRegex(RuntimeError, "Room does not exist !nope:test"): + self.persister.persist_events_txn_for_test("!nope:test", []) + + +class StateDagPersistTxnTestCase(PersistTxnTestCase): + room_version = "org.matrix.msc4242.12" + + def default_config(self) -> dict[str, Any]: + config = super().default_config() + config["experimental_features"]["msc4242_enabled"] = True + return config