diff --git a/rust/src/persist/batch.rs b/rust/src/persist/batch.rs new file mode 100644 index 0000000000..a58804c497 --- /dev/null +++ b/rust/src/persist/batch.rs @@ -0,0 +1,245 @@ +/* + * This file is licensed under the Affero General Public License (AGPL) version 3. + * + * Copyright (C) 2026 Element Creations Ltd + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * See the GNU Affero General Public License for more details: + * . + * + */ + +//! A queued batch of one room's events, as `_persist_event_batch` persists +//! it: the events already persisted under the same transaction ID are +//! dropped, the rest go in chunks of [`CHUNK_SIZE`], and each chunk is +//! written while its state groups are held back from deletion. +//! +//! The loop itself is [`super::python::Persister::persist_batch`], which +//! hands each written chunk to Python. This holds its steps. + +use std::collections::{BTreeSet, HashMap, HashSet}; + +use anyhow::anyhow; +use futures::FutureExt; + +use super::context::PersistContext; +use crate::{ + clock::Clock, + database::postgres::{ + pool::ConnectionPool, + write::{run_write_interaction, Column, WriteError}, + }, + events::Event, + storage::events::Querier, +}; + +/// How many events `_persist_event_batch` writes in one transaction. +pub const CHUNK_SIZE: usize = 100; + +/// `get_already_persisted_events`: for each event sent with a transaction +/// ID from a device, the event already persisted with the same room, +/// sender, device and transaction ID, if there is one. An event that +/// repeats an earlier one of the batch maps to that one. +pub async fn already_persisted_events( + querier: &impl Querier, + events: &[(Event, PersistContext)], +) -> anyhow::Result> { + let mut mapping = HashMap::new(); + let mut by_txn_id: HashMap<(String, String, String, String), String> = HashMap::new(); + for (event, _) in events { + let metadata = event.internal_metadata(); + let (Some(device_id), Some(txn_id)) = (metadata.device_id()?, metadata.txn_id()?) else { + continue; + }; + // Python skips an empty device or transaction ID. + if device_id.is_empty() || txn_id.is_empty() { + continue; + } + let key = ( + event.room_id().to_owned(), + event.sender().to_owned(), + device_id, + txn_id, + ); + if let Some(existing) = by_txn_id.get(&key) { + mapping.insert(event.event_id().to_owned(), existing.clone()); + continue; + } + let rows = querier + .query( + "SELECT event_id FROM event_txn_id_device_id + WHERE room_id = $1 AND user_id = $2 AND device_id = $3 AND txn_id = $4", + &[&key.0, &key.1, &key.2, &key.3], + ) + .await?; + let existing: Option = rows.first().map(|row| row.try_get(0)).transpose()?; + match existing { + Some(existing) => { + mapping.insert(event.event_id().to_owned(), existing.clone()); + by_txn_id.insert(key, existing); + } + None => { + by_txn_id.insert(key, event.event_id().to_owned()); + } + } + } + Ok(mapping) +} + +/// The state groups `persisting_state_group_references` holds back from +/// deletion for `events`: those of each event that is neither rejected nor +/// an outlier, and the groups before them. +pub fn referenced_state_groups( + events: &[(Event, PersistContext)], +) -> anyhow::Result> { + let mut groups = BTreeSet::new(); + for (event, context) in events { + if context.is_rejected() || event.internal_metadata().is_outlier()? { + continue; + } + let group = context + .state_group + .ok_or_else(|| anyhow!("{} has no state group", event.event_id()))?; + groups.insert(group); + // Python skips a falsy group before the event. + if let Some(before) = context.state_group_before_event.filter(|g| *g != 0) { + groups.insert(before); + } + } + Ok(groups) +} + +/// The state groups were deleted before they could be marked. +#[derive(Debug)] +pub struct StateGroupsDeleted(pub Vec); + +impl std::fmt::Display for StateGroupsDeleted { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // As Python's `shortstr`. + let shown: Vec = self.0.iter().take(5).map(i64::to_string).collect(); + let more = if self.0.len() > 5 { ", ..." } else { "" }; + write!( + f, + "state groups have been deleted: [{}{more}]", + shown.join(", ") + ) + } +} + +impl std::error::Error for StateGroupsDeleted {} + +/// `_mark_state_groups_as_persisting_txn`: check the groups still exist, +/// locking them against a concurrent delete, and record that this instance +/// is persisting them. +pub async fn mark_state_groups_as_persisting( + pool: &ConnectionPool, + clock: &Clock, + instance_name: &str, + groups: &BTreeSet, +) -> Result<(), WriteError> { + let groups: Vec = groups.iter().copied().collect(); + run_write_interaction(pool, "mark_state_groups_as_persisting", |txn, _| { + let groups = &groups; + async move { + let rows = txn + .query( + "SELECT id FROM state_groups WHERE id = ANY($1) FOR KEY SHARE", + &[groups], + ) + .await?; + let existing = rows + .iter() + .map(|row| row.try_get(0)) + .collect::, _>>() + .map_err(anyhow::Error::new)?; + let missing: Vec = groups + .iter() + .copied() + .filter(|g| !existing.contains(g)) + .collect(); + if !missing.is_empty() { + return Err(anyhow::Error::new(StateGroupsDeleted(missing)).into()); + } + let now = i64::try_from(clock.now_millis()).map_err(anyhow::Error::new)?; + let instance_names = vec![instance_name; groups.len()]; + let inserted = vec![now; groups.len()]; + txn.upsert_many( + "state_groups_persisting", + &[ + Column::new("state_group", "bigint", groups), + Column::new("instance_name", "text", &instance_names), + ], + &[Column::new("inserted_ts", "bigint", &inserted)], + ) + .await?; + Ok(()) + } + .boxed() + }) + .await +} + +/// `_finish_persisting_txn`: this instance no longer persists the groups. +/// If the persist failed, any pending deletion of them is pushed back, as +/// they may now be referenced. Otherwise it is dropped, as they are. +pub async fn finish_persisting( + pool: &ConnectionPool, + clock: &Clock, + instance_name: &str, + groups: &BTreeSet, + error: bool, +) -> Result<(), WriteError> { + let groups: Vec = groups.iter().copied().collect(); + run_write_interaction(pool, "finish_persisting", |txn, _| { + let groups = &groups; + async move { + txn.execute( + "DELETE FROM state_groups_persisting + WHERE instance_name = $1 AND state_group = ANY($2)", + &[&instance_name, groups], + ) + .await?; + if error { + // `_bump_deletion_txn`. + let now = i64::try_from(clock.now_millis()).map_err(anyhow::Error::new)?; + txn.execute( + "UPDATE state_groups_pending_deletion + SET sequence_number = DEFAULT, insertion_ts = $1 + WHERE state_group = ANY($2)", + &[&now, groups], + ) + .await?; + } else { + txn.execute( + "DELETE FROM state_groups_pending_deletion WHERE state_group = ANY($1)", + &[groups], + ) + .await?; + } + Ok(()) + } + .boxed() + }) + .await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn deleted_groups_are_shown_as_python_shows_them() { + assert_eq!( + StateGroupsDeleted(vec![1, 2]).to_string(), + "state groups have been deleted: [1, 2]" + ); + assert_eq!( + StateGroupsDeleted((1..=6).collect()).to_string(), + "state groups have been deleted: [1, 2, 3, 4, 5, ...]" + ); + } +} diff --git a/rust/src/persist/mod.rs b/rust/src/persist/mod.rs index 3c0207067b..953ee16a3e 100644 --- a/rust/src/persist/mod.rs +++ b/rust/src/persist/mod.rs @@ -18,6 +18,7 @@ //! This is being built up from the inside out (see //! `renders/event-persistence-rust-plan.md`). So far: //! +//! - [`batch`]: the steps around a queued batch's chunks. //! - [`chain_cover`]: the chain cover index of a chunk's state events. //! - [`context`]: what persisting needs from an `EventContext`. //! - [`effects`]: the Python cache updates a Rust transaction asks for once @@ -38,6 +39,7 @@ use pyo3::prelude::*; +pub mod batch; pub mod chain_cover; pub mod context; pub mod effects; @@ -56,6 +58,9 @@ 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_class::()?; + child_module.add_class::()?; child_module.add_class::()?; child_module.add_class::()?; child_module.add_function(wrap_pyfunction!( diff --git a/rust/src/persist/prepare.rs b/rust/src/persist/prepare.rs index f6e374b760..fb22f5d6bc 100644 --- a/rust/src/persist/prepare.rs +++ b/rust/src/persist/prepare.rs @@ -40,8 +40,6 @@ use std::{ sync::Arc, }; -use pyo3::prelude::*; - use super::{ chain_cover::{self, ChainCoverEvent, NewEventChainLinks}, context::PersistContext, @@ -59,8 +57,7 @@ use crate::{ state::EventMap, }; -/// A chunk of one room's events, in the shapes the steps read. Built while -/// the GIL is held, so that the steps run without it. +/// A chunk of one room's events, in the shapes the steps read. pub struct PrepareInput { pub room_id: String, pub backfilled: bool, @@ -74,17 +71,16 @@ pub struct PrepareInput { } impl PrepareInput { + /// The chunk `events_and_contexts`, which share their data with the + /// Python events they were taken from. pub fn new( - py: Python<'_>, room_id: String, - events_and_contexts: &[(Py, PersistContext)], + events_and_contexts: &[(Event, PersistContext)], backfilled: bool, - ) -> PyResult { - let (events, event_map) = reader::chunk(py, events_and_contexts)?; - let chunk_events: Vec<&Event> = events_and_contexts - .iter() - .map(|(event, _)| event.bind(py).get()) - .collect(); + ) -> anyhow::Result { + let (events, event_map) = reader::chunk(events_and_contexts)?; + let chunk_events: Vec<&Event> = + events_and_contexts.iter().map(|(event, _)| event).collect(); Ok(PrepareInput { room_id, backfilled, diff --git a/rust/src/persist/python.rs b/rust/src/persist/python.rs index b565f9d9f3..7e1e59e99b 100644 --- a/rust/src/persist/python.rs +++ b/rust/src/persist/python.rs @@ -17,7 +17,7 @@ //! `RustHandlers.persister` hands out, and helpers for tests. use std::{ - collections::{HashMap, HashSet}, + collections::{BTreeSet, HashMap, HashSet}, sync::{Arc, Mutex}, }; @@ -26,6 +26,7 @@ use http::StatusCode; use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyDict}; use super::{ + batch, context::PersistContext, effects::PersistEffects, extremities::MissingPrevStateEvents, @@ -186,83 +187,25 @@ impl Persister { if !matches!(&*prepared, Prepared::Chunk(_)) { return Err(PyRuntimeError::new_err("the chunk was prepared in Python")); } - let events: Vec<(Event, PersistContext)> = events_and_contexts - .into_iter() - .map(|(event, context)| (event.bind(py).get().shallow_clone(), context)) - .collect(); + let events = shallow_events(py, events_and_contexts); let pool = self.inner.pool.get()?; let inner = Arc::clone(&self.inner); - let future = async move { + Ok(async move { let Prepared::Chunk(chunk) = &*prepared else { unreachable!("checked above"); }; - // Taken as late as possible, as Python takes them: until the - // chunk is written, no later event on the stream can be - // announced. - let event_stream = match orderings { - StreamOrderings::Set => None, - StreamOrderings::Allocate { backfilled } => { - let stream = if backfilled { - TxnStream::backfill(Arc::clone(&inner.backfill_stream)) - } else { - TxnStream::events(Arc::clone(&inner.events_stream)) - }; - let ids = allocate_stream_orderings(&pool, &stream, events.len()) - .await - .map_err(persist_error_to_py)?; - set_persisted_positions(&inner, &events, &ids, backfilled) - .map_err(persist_error_to_py)?; - Some((stream, ids, backfilled)) - } - }; - let events = events - .into_iter() - .map(|(event, context)| TxnEvent::new(event, context)) - .collect::>>() - .map_err(persist_error_to_py)?; - - let cache_stream = TxnStream::caches(Arc::clone(&inner.cache_stream)); - let sticky_stream = TxnStream::sticky_events(Arc::clone(&inner.sticky_stream)); - let profile_stream = TxnStream::profile_updates(Arc::clone(&inner.profile_stream)); - let input = TxnInput { - room_id: &room_id, - events: &events, - prepared: chunk, - instance_name: &inner.instance_name, - server_name: &inner.server_name, - clock: inner.runtime.clock(), + write_chunk( + &inner, + &pool, + &room_id, + events, + chunk, + orderings, inhibit_local_membership_updates, - enable_search: inner.enable_search, - enable_ephemeral_messages: inner.enable_ephemeral_messages, - msc4354_enabled: inner.msc4354_enabled, - include_profile_updates_in_sync: inner.include_profile_updates_in_sync, - cache_stream: &cache_stream, - sticky_stream: &sticky_stream, - profile_stream: &profile_stream, - }; - let mut effects = run_write_interaction(&pool, "persist_events", |txn, _| { - let input = &input; - async move { - let mut effects = PersistEffects::default(); - txn::persist_events_txn(txn, input, &mut effects).await?; - Ok(effects) - } - .boxed() - }) + ) .await - .map_err(write_error_to_py)?; - cache_stream.hand_over(effects.cache_stream_ids()); - sticky_stream.hand_over(effects.sticky_stream_ids()); - profile_stream.hand_over(effects.profile_stream_ids()); - // Finished on drop if the transaction failed, as `get_next_mult` - // finishes them when its block raises. - if let Some((stream, ids, backfilled)) = event_stream { - stream.hand_over(&ids); - effects.add_stream_ids(ids, backfilled); - } - Ok(effects) - }; - Ok(future.boxed()) + } + .boxed()) } /// Convert the chunk while the GIL is held, and return the future that @@ -275,59 +218,39 @@ impl Persister { backfilled: bool, shadow: bool, ) -> PyResult>> { - let input = PrepareInput::new(py, room_id, &events_and_contexts, backfilled)?; + let events = shallow_events(py, events_and_contexts); + let input = PrepareInput::new(room_id, &events, backfilled).map_err(persist_error_to_py)?; let pool = self.inner.pool.get()?; - let handle = self.inner.runtime.tokio_handle()?; let inner = Arc::clone(&self.inner); Ok(async move { - let resolver = NativeStateResolver { - pool: pool.clone(), - handle, - caches: Arc::clone(&inner.store.caches), - conflict_cache: Arc::clone(&inner.conflict_cache), - }; - let reader = PersistReader::new(inner.sync.reader(pool.clone())); - let options = PrepareOptions { - server_name: &inner.server_name, - clock: inner.runtime.clock(), - sticky: inner.msc4354_enabled, - dry_run: shadow, - }; - // 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, - ) + let (prepared, counts) = prepare_chunk(&inner, &pool, &input, shadow) .await - { - Ok(prepared) => { - return Ok(PreparedChunk { - prepared: Arc::new(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)), - } - } + .map_err(persist_error_to_py)?; + Ok(PreparedChunk { + prepared: Arc::new(prepared), + counts: Mutex::new(Some(counts)), + }) } .boxed()) } + /// The future of the batch's next step. See [`Self::batch_next`]. + fn start_batch_next(&self, batch: &PersistBatch) -> BoxFuture<'static, PyResult> { + let state = Arc::clone(&batch.state); + async move { batch_next(&mut *state.lock().await).await }.boxed() + } + + /// The future that releases the batch. See [`Self::batch_abort`]. + fn start_batch_abort(&self, batch: &PersistBatch) -> BoxFuture<'static, PyResult<()>> { + let state = Arc::clone(&batch.state); + async move { + let mut state = state.lock().await; + state.done = true; + state.release_guard(true).await + } + .boxed() + } + /// Run `future` to completion on the runtime, releasing the GIL. For the /// `_blocking` entry points. fn block_on( @@ -481,6 +404,74 @@ impl Persister { self.block_on(py, future) } + /// `_persist_event_batch` for a queued batch of `room_id`'s events, + /// one chunk at a time: Python awaits [`Self::batch_next`] until it + /// resolves with a [`PersistedBatch`], applying the effects of each + /// [`WrittenChunk`] before the next step, and calls + /// [`Self::batch_abort`] if it cannot. Converts the events, without + /// I/O. + #[pyo3(signature = (room_id, events_and_contexts, *, backfilled))] + fn persist_batch( + &self, + py: Python<'_>, + room_id: String, + events_and_contexts: Vec<(Py, PersistContext)>, + backfilled: bool, + ) -> PyResult { + let state = BatchState { + inner: Arc::clone(&self.inner), + pool: self.inner.pool.get()?, + room_id, + backfilled, + events: shallow_events(py, events_and_contexts), + deduplicated: None, + next: 0, + guard_held: false, + done: false, + }; + Ok(PersistBatch { + state: Arc::new(tokio::sync::Mutex::new(state)), + }) + } + + /// The batch's next step: a Deferred resolving to the [`WrittenChunk`] + /// written, or to the [`PersistedBatch`] once the batch is done, or + /// failing as the Python path would. The chunk's stream IDs are handed + /// to Python unfinished, so Python must not cancel the Deferred: if it + /// is cancelled, the step runs on and the IDs are lost. If the step + /// fails, Rust finishes them and releases the state groups. + fn batch_next<'py>( + &self, + py: Python<'py>, + batch: &PersistBatch, + ) -> PyResult> { + let future = self.start_batch_next(batch); + create_deferred_with(py, &self.inner.runtime, future, OnCancel::RunOn) + } + + /// As [`Self::batch_next`], but blocks. + fn batch_next_blocking(&self, py: Python<'_>, batch: &PersistBatch) -> PyResult { + let future = self.start_batch_next(batch); + self.block_on(py, future) + } + + /// Give up the batch, as Python could not apply a chunk's effects: + /// release its state groups as a failed persist does. + fn batch_abort<'py>( + &self, + py: Python<'py>, + batch: &PersistBatch, + ) -> PyResult> { + let future = self.start_batch_abort(batch); + create_deferred_with(py, &self.inner.runtime, future, OnCancel::RunOn) + } + + /// As [`Self::batch_abort`], but blocks. + fn batch_abort_blocking(&self, py: Python<'_>, batch: &PersistBatch) -> PyResult<()> { + let future = self.start_batch_abort(batch); + self.block_on(py, future) + } + /// As [`Self::prepare_chunk`], but blocks. #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))] fn prepare_chunk_blocking( @@ -497,6 +488,416 @@ impl Persister { } } +/// The events of `events_and_contexts`, sharing their data, internal +/// metadata included, with the Python events, so that Rust can read them +/// without the GIL and set their stream orderings. +fn shallow_events( + py: Python<'_>, + events_and_contexts: Vec<(Py, PersistContext)>, +) -> Vec<(Event, PersistContext)> { + events_and_contexts + .into_iter() + .map(|(event, context)| (event.bind(py).get().shallow_clone(), context)) + .collect() +} + +/// Prepare the chunk `input`, with what it counted. +async fn prepare_chunk( + inner: &PersisterInner, + pool: &ConnectionPool, + input: &PrepareInput, + shadow: bool, +) -> anyhow::Result<(Prepared, PersistCounts)> { + let resolver = NativeStateResolver { + pool: pool.clone(), + handle: inner.runtime.tokio_handle()?, + caches: Arc::clone(&inner.store.caches), + conflict_cache: Arc::clone(&inner.conflict_cache), + }; + let reader = PersistReader::new(inner.sync.reader(pool.clone())); + let options = PrepareOptions { + server_name: &inner.server_name, + clock: inner.runtime.clock(), + sticky: inner.msc4354_enabled, + dry_run: shadow, + }; + // 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((prepared, 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(e), + } + } +} + +/// Write the chunk `events`, which Rust prepared as `chunk`, taking its +/// stream orderings first if `orderings` says so. Resolves with the +/// effects, whose stream IDs are unfinished. +async fn write_chunk( + inner: &PersisterInner, + pool: &ConnectionPool, + room_id: &str, + events: Vec<(Event, PersistContext)>, + chunk: &prepare::PreparedChunk, + orderings: StreamOrderings, + inhibit_local_membership_updates: bool, +) -> PyResult { + // Taken as late as possible, as Python takes them: until the chunk is + // written, no later event on the stream can be announced. + let event_stream = match orderings { + StreamOrderings::Set => None, + StreamOrderings::Allocate { backfilled } => { + let stream = if backfilled { + TxnStream::backfill(Arc::clone(&inner.backfill_stream)) + } else { + TxnStream::events(Arc::clone(&inner.events_stream)) + }; + let ids = allocate_stream_orderings(pool, &stream, events.len()) + .await + .map_err(persist_error_to_py)?; + set_persisted_positions(inner, &events, &ids, backfilled) + .map_err(persist_error_to_py)?; + Some((stream, ids, backfilled)) + } + }; + let events = events + .into_iter() + .map(|(event, context)| TxnEvent::new(event, context)) + .collect::>>() + .map_err(persist_error_to_py)?; + + let cache_stream = TxnStream::caches(Arc::clone(&inner.cache_stream)); + let sticky_stream = TxnStream::sticky_events(Arc::clone(&inner.sticky_stream)); + let profile_stream = TxnStream::profile_updates(Arc::clone(&inner.profile_stream)); + let input = TxnInput { + room_id, + events: &events, + prepared: chunk, + instance_name: &inner.instance_name, + server_name: &inner.server_name, + clock: inner.runtime.clock(), + inhibit_local_membership_updates, + enable_search: inner.enable_search, + enable_ephemeral_messages: inner.enable_ephemeral_messages, + msc4354_enabled: inner.msc4354_enabled, + include_profile_updates_in_sync: inner.include_profile_updates_in_sync, + cache_stream: &cache_stream, + sticky_stream: &sticky_stream, + profile_stream: &profile_stream, + }; + let mut effects = run_write_interaction(pool, "persist_events", |txn, _| { + let input = &input; + async move { + let mut effects = PersistEffects::default(); + txn::persist_events_txn(txn, input, &mut effects).await?; + Ok(effects) + } + .boxed() + }) + .await + .map_err(write_error_to_py)?; + cache_stream.hand_over(effects.cache_stream_ids()); + sticky_stream.hand_over(effects.sticky_stream_ids()); + profile_stream.hand_over(effects.profile_stream_ids()); + // Finished on drop if the transaction failed, as `get_next_mult` + // finishes them when its block raises. + if let Some((stream, ids, backfilled)) = event_stream { + stream.hand_over(&ids); + effects.add_stream_ids(ids, backfilled); + } + Ok(effects) +} + +/// A queued batch of one room's events, persisted chunk by chunk by +/// [`batch_next`]. See [`Persister::persist_batch`]. +struct BatchState { + inner: Arc, + pool: ConnectionPool, + room_id: String, + backfilled: bool, + events: Vec<(Event, PersistContext)>, + /// Set by the first step. + deduplicated: Option, + /// The position in `kept` of the next chunk. + next: usize, + /// The state groups are marked as being persisted for the chunk last + /// handed to Python, which the next step or [`batch_abort`] releases. + guard_held: bool, + /// The batch has been persisted, handed back or failed. + done: bool, +} + +/// The batch once the events already persisted are dropped. +struct Deduplicated { + replaced_events: HashMap, + /// The positions in the batch of the events to persist. + kept: Vec, + /// What `persisting_state_group_references` holds back from deletion + /// around each chunk: the groups of the whole batch, as in Python. + state_groups: BTreeSet, +} + +impl BatchState { + fn shallow(&self, i: usize) -> (Event, PersistContext) { + (self.events[i].0.shallow_clone(), self.events[i].1.clone()) + } + + /// `_finish_persisting_txn` for the chunk last written, if its state + /// groups are still marked. + async fn release_guard(&mut self, error: bool) -> PyResult<()> { + if !std::mem::take(&mut self.guard_held) { + return Ok(()); + } + let Some(deduplicated) = &self.deduplicated else { + return Ok(()); + }; + batch::finish_persisting( + &self.pool, + self.inner.runtime.clock(), + &self.inner.instance_name, + &deduplicated.state_groups, + error, + ) + .await + .map_err(write_error_to_py) + } +} + +impl Drop for BatchState { + /// A batch dropped with its guard held, which Python should not let + /// happen, releases it as a failure, as Python's `finally` would. + fn drop(&mut self) { + if !self.guard_held { + return; + } + let (Some(deduplicated), Ok(handle)) = + (self.deduplicated.take(), self.inner.runtime.tokio_handle()) + else { + return; + }; + let inner = Arc::clone(&self.inner); + let pool = self.pool.clone(); + handle.spawn(async move { + let result = batch::finish_persisting( + &pool, + inner.runtime.clock(), + &inner.instance_name, + &deduplicated.state_groups, + true, + ) + .await; + if let Err(e) = result { + log::warn!("Failed to release the state groups of a dropped batch: {e}"); + } + }); + } +} + +/// What one [`batch_next`] step did. +enum BatchStep { + /// A chunk was written, and its effects are Python's to apply. + Written(WrittenChunk), + /// The batch is done. + Done(PersistedBatch), +} + +impl<'py> IntoPyObject<'py> for BatchStep { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> PyResult> { + match self { + BatchStep::Written(chunk) => Ok(Bound::new(py, chunk)?.into_any()), + BatchStep::Done(batch) => Ok(Bound::new(py, batch)?.into_any()), + } + } +} + +/// One step of `_persist_event_batch` for `state`. The first step drops the +/// events already persisted under the same transaction ID +/// (`get_already_persisted_events`). Each step releases the state groups +/// of the chunk before it, then prepares the next chunk, marks the batch's +/// state groups as being persisted, takes the chunk's stream orderings and +/// writes it. +/// +/// A chunk Rust cannot prepare, because its state resolution falls back to +/// Python or preparing it fails, is handed back with the rest of the batch +/// in [`PersistedBatch::remaining`], unwritten. +async fn batch_next(state: &mut BatchState) -> PyResult { + if state.done { + return Err(PyRuntimeError::new_err("the batch is done")); + } + let result = batch_step(state).await; + if !matches!(result, Ok(BatchStep::Written(_))) { + state.done = true; + } + result +} + +async fn batch_step(state: &mut BatchState) -> PyResult { + if state.deduplicated.is_none() { + // Retried as Python's `runInteraction`s are. + let mut attempt = 1; + let replaced_events = loop { + let result = async { + let querier = PoolQuerier::checkout(&state.pool).await?; + batch::already_persisted_events(&querier, &state.events).await + } + .await; + match result { + Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => { + log::warn!( + "[TXN RETRY] {{get_already_persisted_events}} {e:#} {attempt}/{MAX_ATTEMPTS}" + ); + attempt += 1; + } + result => break result.map_err(persist_error_to_py)?, + } + }; + let kept: Vec = (0..state.events.len()) + .filter(|&i| !replaced_events.contains_key(state.events[i].0.event_id())) + .collect(); + let kept_events: Vec<_> = kept.iter().map(|&i| state.shallow(i)).collect(); + let state_groups = + batch::referenced_state_groups(&kept_events).map_err(persist_error_to_py)?; + state.deduplicated = Some(Deduplicated { + replaced_events, + kept, + state_groups, + }); + } + state.release_guard(false).await?; + + let deduplicated = state.deduplicated.as_ref().expect("set above"); + let start = state.next; + if start >= deduplicated.kept.len() { + return Ok(BatchStep::Done(PersistedBatch { + replaced_events: deduplicated.replaced_events.clone(), + remaining: Vec::new(), + fallback_reason: None, + })); + } + let end = (start + batch::CHUNK_SIZE).min(deduplicated.kept.len()); + let indices = deduplicated.kept[start..end].to_vec(); + let chunk: Vec<(Event, PersistContext)> = indices.iter().map(|&i| state.shallow(i)).collect(); + let hand_back = |reason: String| { + BatchStep::Done(PersistedBatch { + replaced_events: deduplicated.replaced_events.clone(), + remaining: deduplicated.kept[start..].to_vec(), + fallback_reason: Some(reason), + }) + }; + + let inner = &state.inner; + let pool = &state.pool; + let prepared = async { + let input = PrepareInput::new(state.room_id.clone(), &chunk, state.backfilled)?; + prepare_chunk(inner, pool, &input, false).await + } + .await; + let (prepared, counts) = match prepared { + Ok((Prepared::StateResFallback(reason), _)) => return Ok(hand_back(reason)), + // The chunk path prepares it again, and raises or falls back to + // Python as the error and the strictness say. + Err(e) => { + log::warn!( + "Failed to prepare events for {} in Rust, handing them back: {e:#}", + state.room_id + ); + return Ok(hand_back(format!("{e:#}"))); + } + Ok(prepared) => prepared, + }; + inner.metrics.record(|total| total.add(counts)); + let prepared = Arc::new(prepared); + let Prepared::Chunk(prepared_chunk) = &*prepared else { + unreachable!("handed back above"); + }; + + let guarded = !deduplicated.state_groups.is_empty(); + if guarded { + batch::mark_state_groups_as_persisting( + pool, + inner.runtime.clock(), + &inner.instance_name, + &deduplicated.state_groups, + ) + .await + .map_err(write_error_to_py)?; + state.guard_held = true; + } + let written = write_chunk( + inner, + pool, + &state.room_id, + chunk, + prepared_chunk, + StreamOrderings::Allocate { + backfilled: state.backfilled, + }, + state.backfilled, + ) + .await; + let effects = match written { + Ok(effects) => effects, + Err(e) => { + // As Python's `finally`, whose error replaces the persist's. + state.release_guard(true).await?; + return Err(e); + } + }; + state.next = end; + Ok(BatchStep::Written(WrittenChunk { + effects, + prepared, + indices, + })) +} + +/// [`batch_next`]'s chunk, written, whose effects Python must apply. +#[pyclass(frozen)] +pub struct WrittenChunk { + effects: PersistEffects, + prepared: Arc, + indices: Vec, +} + +#[pymethods] +impl WrittenChunk { + /// What Python applies with `apply_rust_persist_effects`, with the + /// chunk's stream IDs unfinished. + #[getter] + fn effects(&self) -> PersistEffects { + self.effects.clone() + } + + /// What Rust prepared for the chunk. + #[getter] + fn prepared(&self) -> PreparedChunk { + PreparedChunk { + prepared: Arc::clone(&self.prepared), + // Counted when it was prepared. + counts: Mutex::new(None), + } + } + + /// The chunk's positions in the batch. + #[getter] + fn indices(&self) -> Vec { + self.indices.clone() + } +} + /// Where the stream orderings of a chunk come from. #[derive(Clone, Copy)] enum StreamOrderings { @@ -560,6 +961,44 @@ fn set_persisted_positions( Ok(()) } +/// A batch being persisted. See [`Persister::persist_batch`]. +#[pyclass(frozen)] +pub struct PersistBatch { + state: Arc>, +} + +/// What [`Persister::persist_batch`] did, once the batch is done. +#[pyclass(frozen)] +pub struct PersistedBatch { + replaced_events: HashMap, + remaining: Vec, + fallback_reason: Option, +} + +#[pymethods] +impl PersistedBatch { + /// The events not persisted as they repeat an event already persisted + /// with the same transaction ID, mapped to that event. + #[getter] + fn replaced_events(&self) -> HashMap { + self.replaced_events.clone() + } + + /// The positions in the batch of the events left for Python to + /// persist, in order: the chunk Rust could not prepare and those after + /// it. Empty if Rust persisted them all. + #[getter] + fn remaining(&self) -> Vec { + self.remaining.clone() + } + + /// Why Rust left the `remaining` events. + #[getter] + fn fallback_reason(&self) -> Option { + self.fallback_reason.clone() + } +} + /// What [`Persister::prepare_chunk`] computed, in the shapes Python builds /// its own types from. If `state_res_fallback` is set, the rest is empty /// and Python should compute the chunk itself. diff --git a/rust/src/persist/reader.rs b/rust/src/persist/reader.rs index 5d6b50d9ef..289834b40f 100644 --- a/rust/src/persist/reader.rs +++ b/rust/src/persist/reader.rs @@ -727,7 +727,7 @@ impl PyPersistReaderForTest { now_ms: i64, state_dag: bool, ) -> PyResult<(StateAndExtremities, PersistCounts)> { - let (events, event_map) = chunk(py, &events_and_contexts)?; + let (events, event_map) = py_chunk(py, &events_and_contexts)?; let resolver = self.resolver(); self.run(py, |reader| { async move { @@ -1199,7 +1199,7 @@ impl PyPersistReaderForTest { HashSet, Bound<'py, PyDict>, )> { - let (events, event_map) = chunk(py, &events_and_contexts)?; + let (events, event_map) = py_chunk(py, &events_and_contexts)?; let old: HashSet> = arcs(old_latest_event_ids).into_iter().collect(); let new: HashSet> = arcs(new_latest_event_ids).into_iter().collect(); let resolver = self.resolver(); @@ -1266,7 +1266,7 @@ impl PyPersistReaderForTest { to_delete: Vec<(String, String)>, to_insert: HashMap<(String, String), String>, ) -> PyResult { - let (events, _) = chunk(py, &events_and_contexts)?; + let (events, _) = py_chunk(py, &events_and_contexts)?; let delta = DeltaState { to_delete: to_delete .iter() @@ -1523,14 +1523,15 @@ fn state_map(map: HashMap<(String, String), String>) -> StateMap { /// The chunk's events as the computation reads them, and as state /// resolution's map of events in flight. pub(super) fn chunk( - py: Python<'_>, - events_and_contexts: &[(Py, PersistContext)], -) -> PyResult<(Vec, EventMap)> { - let events = new_events(py, events_and_contexts)?; + events_and_contexts: &[(Event, PersistContext)], +) -> anyhow::Result<(Vec, EventMap)> { + let events = events_and_contexts + .iter() + .map(|(event, context)| NewEvent::new(event, context)) + .collect::>()?; let event_map = events_and_contexts .iter() .map(|(event, _)| { - let event = event.bind(py).get(); ( Arc::from(event.event_id()), StateResEvent::new(event.shallow_clone()), @@ -1540,6 +1541,18 @@ pub(super) fn chunk( Ok((events, event_map)) } +/// [`chunk`] of the Python events, for the test entry points. +fn py_chunk( + py: Python<'_>, + events_and_contexts: &[(Py, PersistContext)], +) -> PyResult<(Vec, EventMap)> { + let events: Vec<(Event, PersistContext)> = events_and_contexts + .iter() + .map(|(event, context)| (event.bind(py).get().shallow_clone(), context.clone())) + .collect(); + chunk(&events).map_err(persist_error_to_py) +} + fn new_events( py: Python<'_>, events_and_contexts: &[(Py, PersistContext)], diff --git a/synapse/config/experimental.py b/synapse/config/experimental.py index e5d97e87fb..f1c18484e7 100644 --- a/synapse/config/experimental.py +++ b/synapse/config/experimental.py @@ -415,8 +415,9 @@ class ExperimentalConfig(Config): # - "prepare": Rust's results for that are used. # - "transaction": Rust also runs the write transaction. # - "full": Rust persists each queued batch. - # The levels are being built up in that order. So far "full" only - # has Rust take the stream orderings of each chunk as well. + # The levels are being built up in that order. So far at "full", + # Python still applies what each chunk Rust writes asks for and runs + # what follows it, and recalculates current state on request. rust_event_persistence = experimental.get("rust_event_persistence", "off") # YAML reads an unquoted `off` as a boolean. if rust_event_persistence is False: diff --git a/synapse/storage/controllers/persist_events.py b/synapse/storage/controllers/persist_events.py index 7d8865bb2e..ce287e8a2b 100644 --- a/synapse/storage/controllers/persist_events.py +++ b/synapse/storage/controllers/persist_events.py @@ -82,7 +82,13 @@ from synapse.storage.databases.main.events import ( persist_event_counter, ) from synapse.storage.databases.main.events_worker import EventRedactBehaviour -from synapse.synapse_rust.persist import Persister, PreparedChunk +from synapse.storage.databases.main.rust_persist import finish_rust_persist_stream_ids +from synapse.synapse_rust.persist import ( + PersistBatch, + Persister, + PreparedChunk, + WrittenChunk, +) from synapse.types import ( PersistedEventPosition, RoomStreamToken, @@ -90,7 +96,11 @@ from synapse.types import ( get_domain_from_id, ) from synapse.types.state import StateFilter -from synapse.util.async_helpers import ObservableDeferred, yieldable_gather_results +from synapse.util.async_helpers import ( + ObservableDeferred, + delay_cancellation, + yieldable_gather_results, +) from synapse.util.metrics import Measure from synapse.util.stringutils import shortstr @@ -871,6 +881,14 @@ class EventsPersistenceStorageController: if not events_and_contexts: return replaced_events + if ( + self._rust_persister is not None + and self.hs.config.experimental.rust_event_persistence_at_least("full") + ): + return await self._persist_event_batch_in_rust( + room_id, events_and_contexts, backfilled + ) + # Check if any of the events have a transaction ID that has already been # persisted, and if so we don't persist it again. # @@ -892,6 +910,16 @@ class EventsPersistenceStorageController: if not events_and_contexts: return replaced_events + await self._persist_chunks(room_id, events_and_contexts, backfilled) + return replaced_events + + async def _persist_chunks( + self, + room_id: str, + events_and_contexts: list[EventPersistencePair], + backfilled: bool, + ) -> None: + """Persist the events, deduplicated, in chunks of 100.""" chunks = [ events_and_contexts[x : x + 100] for x in range(0, len(events_and_contexts), 100) @@ -924,7 +952,98 @@ class EventsPersistenceStorageController: rust_transaction=self._rust_transaction(prepared), ) - return replaced_events + async def _persist_event_batch_in_rust( + self, + room_id: str, + events_and_contexts: list[EventPersistencePair], + backfilled: bool, + ) -> dict[str, str]: + """`_persist_event_batch` at `rust_event_persistence: full`: Rust + deduplicates the batch, then prepares and writes it chunk by chunk, + and Python applies what each chunk Rust wrote asks for before the + next. A chunk Rust cannot prepare is handed back with the rest of + the batch and persisted chunk by chunk, as at lower levels.""" + assert self._rust_persister is not None + persister = self._rust_persister + handle = persister.persist_batch( + room_id, events_and_contexts, backfilled=backfilled + ) + + async def persist() -> dict[str, str]: + while True: + try: + # The Deferred follows the logcontext rules already. + step = await persister.batch_next(handle) + except BaseException: + # Rust finished the stream orderings of a chunk it + # failed to write. Tell replication, as `get_next_mult` + # does when its block raises. + self.hs.get_replication_notifier().notify_replication() + raise + if not isinstance(step, WrittenChunk): + break + await self._apply_rust_batch_chunk( + room_id, events_and_contexts, backfilled, step, handle + ) + + if step.remaining: + logger.debug( + "Persisting %d events of %s chunk by chunk: %s", + len(step.remaining), + room_id, + step.fallback_reason, + ) + await self._persist_chunks( + room_id, + [events_and_contexts[i] for i in step.remaining], + backfilled, + ) + return step.replaced_events + + # Rust carries on with a step if this is cancelled, and its effects + # must be applied. Waiting also keeps the room's queue from starting + # its next batch while this one, or the rest Rust handed back, is + # still being written. + return await delay_cancellation(persist()) + + async def _apply_rust_batch_chunk( + self, + room_id: str, + events_and_contexts: list[EventPersistencePair], + backfilled: bool, + step: WrittenChunk, + handle: PersistBatch, + ) -> None: + """Apply what a chunk of the batch Rust wrote asks for, and finish + its stream orderings. If that fails, release the batch.""" + assert self._rust_persister is not None + effects = step.effects + try: + try: + prepared = step.prepared + chunk = [events_and_contexts[i] for i in step.indices] + # Counted as the chunk path counts, once prepared. + self._count_rust_prepare("full", "rust") + if prepared.state_res_stats is not None: + self._state_resolution_handler.record_rust_state_res_metrics( + room_id, prepared.state_res_stats + ) + except BaseException: + # Only the apply finishes the orderings otherwise. + finish_rust_persist_stream_ids( + self.persist_events_store.store, effects, succeeded=False + ) + raise + await self.persist_events_store._apply_rust_chunk( + room_id, + chunk, + effects, + backfilled=backfilled, + new_forward_extremities=prepared.new_forward_extremities, + ) + except BaseException: + await self._rust_persister.batch_abort(handle) + raise def _rust_transaction(self, prepared: _PreparedChunk) -> RustTransaction | None: """Whether to write the chunk in Rust: at `rust_event_persistence: @@ -1071,10 +1190,9 @@ class EventsPersistenceStorageController: with Measure( self._clock, name="rust_prepare_chunk", server_name=self.server_name ): - rust = await make_deferred_yieldable( - self._rust_persister.prepare_chunk( - room_id, chunk, backfilled, shadow=shadow - ) + # The Deferred follows the logcontext rules already. + rust = await self._rust_persister.prepare_chunk( + room_id, chunk, backfilled, shadow=shadow ) if rust.state_res_fallback is not None: logger.info( diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py index bc7f8e55aa..c0948a5080 100644 --- a/synapse/storage/databases/main/events.py +++ b/synapse/storage/databases/main/events.py @@ -57,7 +57,6 @@ from synapse.events import ( ) from synapse.events.py_protocol import MSC4242Event, supports_msc4242_state_dag from synapse.events.snapshot import EventPersistencePair -from synapse.logging.context import make_deferred_yieldable from synapse.logging.opentracing import trace from synapse.metrics import SERVER_NAME_LABEL from synapse.storage._base import db_to_json, make_in_list_sql_clause @@ -98,7 +97,7 @@ from synapse.util.stringutils import non_null_str_or_none if TYPE_CHECKING: from synapse.server import HomeServer from synapse.storage.databases.main import DataStore - from synapse.synapse_rust.persist import Persister, PreparedChunk + from synapse.synapse_rust.persist import PersistEffects, Persister, PreparedChunk logger = logging.getLogger(__name__) @@ -1195,18 +1194,17 @@ class PersistEventsStore: # caller waits for both, as Python's transaction cannot be cancelled. async def persist_and_apply() -> None: + # The persister's Deferreds follow the logcontext rules already. persister = rust_transaction.persister allocate = rust_transaction.allocate_stream_orderings if allocate: try: - effects = await make_deferred_yieldable( - persister.persist_chunk( - room_id, - events_and_contexts, - rust_transaction.prepared, - backfilled=backfilled, - inhibit_local_membership_updates=inhibit_local_membership_updates, - ) + effects = await persister.persist_chunk( + room_id, + events_and_contexts, + rust_transaction.prepared, + backfilled=backfilled, + inhibit_local_membership_updates=inhibit_local_membership_updates, ) except BaseException: # Rust finished the stream orderings it took. Tell @@ -1215,44 +1213,71 @@ class PersistEventsStore: self.hs.get_replication_notifier().notify_replication() raise else: - effects = await make_deferred_yieldable( - persister.persist_events_txn( - room_id, - events_and_contexts, - rust_transaction.prepared, - inhibit_local_membership_updates=inhibit_local_membership_updates, - ) - ) - # The copy of a duplicated event that Rust wrote, as Python's - # `_add_to_cache` reads the deduplicated chunk. - written = self._filter_events_and_contexts_for_duplicates( - events_and_contexts - ) - events = {event.event_id: event for event, _ in written} - if not allocate: - apply_rust_persist_effects(self.store, effects, events) - return - # What follows the write runs before the stream orderings are - # finished, as it runs inside `get_next_mult` in Python: the - # events must not be announced before, say, the MSC4293 - # redactions they start are recorded. - try: - apply_rust_persist_effects( - self.store, effects, events, finish_stream_ids=False - ) - await self._after_persisting_events( + effects = await persister.persist_events_txn( room_id, events_and_contexts, - use_negative_stream_ordering=backfilled, + rust_transaction.prepared, + inhibit_local_membership_updates=inhibit_local_membership_updates, + ) + if allocate: + await self._apply_rust_chunk( + room_id, + events_and_contexts, + effects, + backfilled=backfilled, new_forward_extremities=new_forward_extremities, ) - except BaseException: - finish_rust_persist_stream_ids(self.store, effects, succeeded=False) - raise - finish_rust_persist_stream_ids(self.store, effects, succeeded=True) + else: + apply_rust_persist_effects( + self.store, effects, self._events_written(events_and_contexts) + ) await delay_cancellation(persist_and_apply()) + def _events_written( + self, events_and_contexts: list[EventPersistencePair] + ) -> dict[str, EventBase]: + """The copy of each event that a write of the chunk stores, as + Python's `_add_to_cache` reads the deduplicated chunk.""" + written = self._filter_events_and_contexts_for_duplicates(events_and_contexts) + return {event.event_id: event for event, _ in written} + + async def _apply_rust_chunk( + self, + room_id: str, + events_and_contexts: list[EventPersistencePair], + effects: "PersistEffects", + *, + backfilled: bool, + new_forward_extremities: set[str] | None, + ) -> None: + """Apply the effects of a chunk Rust wrote, having taken its stream + orderings, then run `_after_persisting_events`, then finish the + orderings. + + What follows the write runs before the orderings are finished, as + it runs inside `get_next_mult` in Python: the events must not be + announced before, say, the MSC4293 redactions they start are + recorded. The orderings are finished whatever happens. + """ + try: + apply_rust_persist_effects( + self.store, + effects, + self._events_written(events_and_contexts), + finish_stream_ids=False, + ) + await self._after_persisting_events( + room_id, + events_and_contexts, + use_negative_stream_ordering=backfilled, + new_forward_extremities=new_forward_extremities, + ) + except BaseException: + finish_rust_persist_stream_ids(self.store, effects, succeeded=False) + raise + finish_rust_persist_stream_ids(self.store, effects, succeeded=True) + def _persist_events_txn( self, txn: LoggingTransaction, diff --git a/synapse/synapse_rust/persist.pyi b/synapse/synapse_rust/persist.pyi index 9fe43e347d..e2d94e38e1 100644 --- a/synapse/synapse_rust/persist.pyi +++ b/synapse/synapse_rust/persist.pyi @@ -106,6 +106,71 @@ class Persister: ) -> PersistEffects: """As `persist_chunk`, but blocks. For tests.""" + def persist_batch( + self, + room_id: str, + events_and_contexts: list[tuple[EventBase, EventContext]], + *, + backfilled: bool, + ) -> PersistBatch: + """`_persist_event_batch` for a queued batch of `room_id`'s events, + one chunk at a time: await `batch_next` until it resolves with a + `PersistedBatch`, applying the effects of each `WrittenChunk` before + the next step, and call `batch_abort` if they cannot be applied. + Converts the events, without I/O.""" + def batch_next( + self, batch: PersistBatch + ) -> Deferred[WrittenChunk | PersistedBatch]: + """The batch's next step. The first drops the events already + persisted under the same transaction ID. Each releases the state + groups of the chunk before, then prepares the next chunk, marks the + batch's state groups as being persisted, takes the chunk's stream + orderings and writes it. + + The chunk's stream IDs are handed over unfinished: must not be + cancelled. If the step fails, Rust finishes them and releases the + state groups.""" + def batch_next_blocking(self, batch: PersistBatch) -> WrittenChunk | PersistedBatch: + """As `batch_next`, but blocks. For tests.""" + def batch_abort(self, batch: PersistBatch) -> Deferred[None]: + """Give up the batch, releasing its state groups as a failed + persist does.""" + def batch_abort_blocking(self, batch: PersistBatch) -> None: + """As `batch_abort`, but blocks. For tests.""" + +class PersistBatch: + """A batch being persisted. See `Persister.persist_batch`.""" + +class WrittenChunk: + """A chunk `Persister.batch_next` wrote.""" + + @property + def effects(self) -> PersistEffects: + """To apply with `apply_rust_persist_effects`, with the chunk's + stream IDs unfinished.""" + @property + def prepared(self) -> PreparedChunk: + """What Rust prepared for the chunk.""" + @property + def indices(self) -> list[int]: + """The chunk's positions in the batch.""" + +class PersistedBatch: + """What `Persister.persist_batch` did, once the batch is done.""" + + @property + def replaced_events(self) -> dict[str, str]: + """The events not persisted as they repeat an event already + persisted with the same transaction ID, mapped to that event.""" + @property + def remaining(self) -> list[int]: + """The positions in the batch of the events left for Python to + persist, in order: the chunk Rust could not prepare and those after + it. Empty if Rust persisted them all.""" + @property + def fallback_reason(self) -> str | None: + """Why Rust left the `remaining` events.""" + class PreparedChunk: """What `Persister.prepare_chunk` computed. If `state_res_fallback` is set, the rest is empty and Python should compute the chunk itself.""" diff --git a/tests/storage/test_rust_persist_prepare.py b/tests/storage/test_rust_persist_prepare.py index ffa11e747c..d8fecf034d 100644 --- a/tests/storage/test_rust_persist_prepare.py +++ b/tests/storage/test_rust_persist_prepare.py @@ -13,7 +13,7 @@ # """How `_persist_event_batch` prepares chunks in Rust -(`rust/src/persist/prepare.rs`) at the `shadow` and `prepare` levels of +(`rust/src/persist/prepare.rs`) at the `shadow`, `prepare` and `full` levels of `rust_event_persistence`. These need the native Rust Postgres driver: SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.storage.test_rust_persist_prepare @@ -21,7 +21,7 @@ from typing import Any from unittest import SkipTest -from unittest.mock import Mock +from unittest.mock import Mock, patch import attr @@ -118,6 +118,30 @@ class _PrepareTestBase(unittest.HomeserverTestCase): if after[result] != self._counts_before[result] } + def _mock_persister(self, **kwargs: Any) -> Mock: + """A mock persister. At `full` it hands each batch back to be + persisted chunk by chunk, so the mocked `prepare_chunk` is used.""" + + def persist_batch( + room_id: str, events_and_contexts: list, **kwargs: Any + ) -> Mock: + return Mock(events=len(events_and_contexts)) + + def batch_next(handle: Mock) -> defer.Deferred: + return defer.succeed( + Mock( + replaced_events={}, + remaining=list(range(handle.events)), + fallback_reason="mocked", + ) + ) + + return Mock( + persist_batch=Mock(side_effect=persist_batch), + batch_next=Mock(side_effect=batch_next), + **kwargs, + ) + def _send(self, body: str, expect_code: int = 200) -> str | None: return self.helper.send( self.room_id, body, tok=self.alice_tok, expect_code=expect_code @@ -308,7 +332,7 @@ class ShadowTestCase(_PrepareTestBase): def test_rust_rejection_python_accepts(self) -> None: self.persistence._rust_persist_strict = False - self.persistence._rust_persister = Mock( + self.persistence._rust_persister = self._mock_persister( prepare_chunk=Mock(return_value=defer.fail(SynapseError(500, "no"))) ) with self.assertLogs("synapse.storage.controllers.persist_events", "WARNING"): @@ -317,7 +341,7 @@ class ShadowTestCase(_PrepareTestBase): def test_rust_error_is_logged(self) -> None: self.persistence._rust_persist_strict = False - self.persistence._rust_persister = Mock( + self.persistence._rust_persister = self._mock_persister( prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom"))) ) with self.assertLogs("synapse.storage.controllers.persist_events", "ERROR"): @@ -326,7 +350,7 @@ class ShadowTestCase(_PrepareTestBase): class PrepareTestCase(_PrepareTestBase): - level = "prepare" + level: RustEventPersistenceLevel = "prepare" def test_room_lifecycle(self) -> None: chunks = self._room_lifecycle() @@ -424,7 +448,7 @@ class PrepareTestCase(_PrepareTestBase): def test_rejection_is_raised(self) -> None: """A chunk Rust rejects is not computed again in Python.""" self.persistence._rust_persist_strict = False - self.persistence._rust_persister = Mock( + self.persistence._rust_persister = self._mock_persister( prepare_chunk=Mock(return_value=defer.fail(SynapseError(500, "no"))) ) in_python = Mock(side_effect=AssertionError("computed in Python")) @@ -434,7 +458,7 @@ class PrepareTestCase(_PrepareTestBase): in_python.assert_not_called() def test_fallback(self) -> None: - self.persistence._rust_persister = Mock( + self.persistence._rust_persister = self._mock_persister( prepare_chunk=Mock( return_value=defer.succeed(Mock(state_res_fallback="no index")) ) @@ -448,7 +472,7 @@ class PrepareTestCase(_PrepareTestBase): def test_error_uses_python(self) -> None: self.persistence._rust_persist_strict = False - self.persistence._rust_persister = Mock( + self.persistence._rust_persister = self._mock_persister( prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom"))) ) with self.assertLogs("synapse.storage.controllers.persist_events", "ERROR"): @@ -460,8 +484,55 @@ class PrepareTestCase(_PrepareTestBase): ) def test_error_fails_strict_mode(self) -> None: - self.persistence._rust_persister = Mock( + self.persistence._rust_persister = self._mock_persister( prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom"))) ) self._send("hello", expect_code=500) self.assertEqual(self._counted(), {"error": 1}) + + +class FullTestCase(PrepareTestCase): + """At `full`, Rust persists each batch, chunk by chunk, and hands back + to the chunk path any chunk it cannot prepare. The prepare counter + counts the same.""" + + level: RustEventPersistenceLevel = "full" + + def test_hand_back_in_the_middle(self) -> None: + """A batch whose second chunk Rust cannot prepare: Rust writes the + first and hands back the second alone, which Python persists.""" + room_id = self.helper.create_room_as( + self.alice, room_version=RoomVersions.V1.identifier, tok=self.alice_tok + ) + # Siblings, so the messages need no state resolution and the topics, + # in the second chunk, need it on state resolution v1. + messages = [ + self._new_event( + room_id, EventTypes.Message, {"msgtype": "m.text", "body": str(i)} + ) + for i in range(100) + ] + topics = [ + self._new_event(room_id, EventTypes.Topic, {"topic": t}, state_key="") + for t in ("A", "B") + ] + before = self._counts() + persist_chunks = self.persistence._persist_chunks + handed_back: list[list[str]] = [] + + async def recording( + room_id: str, events_and_contexts: list, backfilled: bool + ) -> None: + handed_back.append([event.event_id for event, _ in events_and_contexts]) + await persist_chunks(room_id, events_and_contexts, backfilled) + + with patch.object(self.persistence, "_persist_chunks", side_effect=recording): + self.get_success(self.persistence.persist_events([*messages, *topics])) + self.assertEqual(handed_back, [[event.event_id for event, _ in topics]]) + after = self._counts() + counted = {r: after[r] - before[r] for r in _RESULTS if after[r] != before[r]} + self.assertEqual(counted, {"rust": 1, "fallback": 1}) + for event, _ in [*messages, *topics]: + self.assertIsNotNone( + self.get_success(self.store.get_event(event.event_id, allow_none=True)) + ) diff --git a/tests/storage/test_rust_persist_prereqs.py b/tests/storage/test_rust_persist_prereqs.py index bcc3b184a5..8ca9681588 100644 --- a/tests/storage/test_rust_persist_prereqs.py +++ b/tests/storage/test_rust_persist_prereqs.py @@ -40,7 +40,11 @@ from synapse.config.experimental import ExperimentalConfig from synapse.config.homeserver import HomeServerConfig from synapse.events import EventBase, make_event_from_dict from synapse.events.snapshot import EventContext -from synapse.logging.context import make_deferred_yieldable +from synapse.logging.context import ( + LoggingContext, + current_context, + make_deferred_yieldable, +) from synapse.metrics import SERVER_NAME_LABEL, _PreCollectHooks from synapse.rest import admin from synapse.rest.client import login, room @@ -52,6 +56,7 @@ from synapse.storage.engines import PostgresEngine from synapse.storage.rust_pool import RustConnectionPool from synapse.synapse_rust.persist import ( PersistEffects, + WrittenChunk, persist_context_for_test, persist_gil_stats, persistence_json_for_test, @@ -871,6 +876,8 @@ class _CountingPersister: self.transactions = 0 # Chunks written with Rust taking their stream orderings. self.chunks = 0 + # Batches persisted in Rust. + self.batches = 0 self.last_effects: PersistEffects | None = None def prepare_chunk(self, *args: Any, **kwargs: Any) -> Any: @@ -895,6 +902,19 @@ class _CountingPersister: return self._persister.persist_chunk(*args, **kwargs).addCallback(keep) + def persist_batch(self, *args: Any, **kwargs: Any) -> Any: + self.batches += 1 + return self._persister.persist_batch(*args, **kwargs) + + def batch_next(self, *args: Any, **kwargs: Any) -> Any: + def keep(step: Any) -> Any: + if isinstance(step, WrittenChunk): + self.chunks += 1 + self.last_effects = step.effects + return step + + return self._persister.batch_next(*args, **kwargs).addCallback(keep) + def __getattr__(self, name: str) -> Any: return getattr(self._persister, name) @@ -1031,7 +1051,8 @@ class RustTransactionTestCase(unittest.HomeserverTestCase): def persist_events_txn(self, *args: Any, **kwargs: Any) -> Any: d: defer.Deferred = defer.Deferred() gate.append((args, kwargs, d)) - return d + # As the persister's Deferreds do. + return make_deferred_yieldable(d) def __getattr__(self, name: str) -> Any: return getattr(blocking, name) @@ -1125,16 +1146,20 @@ class RustFullTestCase(unittest.HomeserverTestCase): self.token = self.login("alice", "pass") self.room_id = self.helper.create_room_as(self.user, tok=self.token) - def _build(self, body: str) -> tuple[EventBase, EventContext]: + def _build( + self, body: str, *, txn_id: str | None = None + ) -> tuple[EventBase, EventContext]: + """A message, sent with `txn_id` from `DEVICE` if given.""" event, unpersisted_context = self.get_success( self.hs.get_event_creation_handler().create_event( - create_requester(self.user), + create_requester(self.user, device_id="DEVICE" if txn_id else None), { "type": EventTypes.Message, "content": {"msgtype": "m.text", "body": body}, "room_id": self.room_id, "sender": self.user, }, + txn_id=txn_id, ) ) context = self.get_success(unpersisted_context.persist(event)) @@ -1263,6 +1288,207 @@ class RustFullTestCase(unittest.HomeserverTestCase): assert failed is not None self.assertGreaterEqual(id_gen.get_current_token(), failed) + def test_batch_is_one_rust_call(self) -> None: + """A batch of 150 events is one Rust call, written in two chunks, + with none of it in Python's store entry point.""" + # Siblings on the room's extremity: an event cannot be built on + # one not yet persisted. + pairs = [self._build(f"message {i}") for i in range(150)] + batches, chunks = self.persister.batches, self.persister.chunks + with patch.object( + PersistEventsStore, + "_persist_events_and_state_updates", + side_effect=AssertionError("persisted in Python"), + ): + self.get_success(self.persistence.persist_events(pairs)) + self.assertEqual(self.persister.batches, batches + 1) + self.assertEqual(self.persister.chunks, chunks + 2) + effects = self.persister.last_effects + assert effects is not None + self.assertEqual(len(effects.stream_ids), 50) + + orderings = [self._event_row(event.event_id)[0] for event, _ in pairs] + self.assertEqual(orderings, sorted(orderings)) + self.assertGreaterEqual( + self.store._stream_id_gen.get_current_token(), orderings[-1] + ) + + def test_transaction_ids(self) -> None: + """An event repeating the transaction ID of an earlier one, in the + same batch or an earlier one, is not persisted and maps to it.""" + self.login("alice", "pass", device_id="DEVICE") + first = self._build("one", txn_id="txn") + again = self._build("one again", txn_id="txn") + persisted, _ = self.get_success(self.persistence.persist_events([first, again])) + self.assertEqual( + [event.event_id for event in persisted], [first[0].event_id] * 2 + ) + + later = self._build("one later", txn_id="txn") + persisted, _ = self.get_success(self.persistence.persist_events([later])) + self.assertEqual([event.event_id for event in persisted], [first[0].event_id]) + for event, _ in (again, later): + self.assertIsNone( + self.get_success(self.store.get_event(event.event_id, allow_none=True)) + ) + + def test_state_groups_held_while_writing(self) -> None: + """The batch's state groups are marked as being persisted while each + chunk is written, and no longer once it is.""" + held: list[list[tuple[int, str]]] = [] + persist_events = self.hs.get_datastores().persist_events + assert persist_events is not None + apply_chunk = persist_events._apply_rust_chunk + + async def checking(*args: Any, **kwargs: Any) -> None: + held.append( + await self.store.db_pool.simple_select_list( + "state_groups_persisting", + None, + ("state_group", "instance_name"), + ) + ) + await apply_chunk(*args, **kwargs) + + event, context = self._build("hello") + with patch.object(persist_events, "_apply_rust_chunk", side_effect=checking): + self.get_success(self.persistence.persist_events([(event, context)])) + assert context.state_group is not None + assert context.state_group_before_event is not None + groups = {context.state_group, context.state_group_before_event} + self.assertEqual( + [sorted(rows) for rows in held], + [sorted((group, "master") for group in groups)], + ) + rows = self.get_success( + self.store.db_pool.simple_select_list( + "state_groups_persisting", None, ("state_group",) + ) + ) + self.assertEqual(rows, []) + + def test_deleted_state_group_fails(self) -> None: + """A chunk whose state group was deleted meanwhile is not written.""" + event, context = self._build("hello") + self.get_success( + self.store.db_pool.simple_delete_one( + "state_groups", {"id": context.state_group} + ) + ) + failure = self.get_failure( + self.persistence.persist_events([(event, context)]), RuntimeError + ) + self.assertIn("state groups have been deleted", str(failure.value)) + self.assertIsNone( + self.get_success(self.store.get_event(event.event_id, allow_none=True)) + ) + + def _pending_deletion(self, state_group: int) -> tuple[int, int] | None: + """`(sequence_number, insertion_ts)` of the group's pending deletion.""" + row = self.get_success( + self.store.db_pool.simple_select_one( + "state_groups_pending_deletion", + {"state_group": state_group}, + ("sequence_number", "insertion_ts"), + allow_none=True, + ) + ) + return (row[0], row[1]) if row is not None else None + + def test_pending_deletion_after_persisting(self) -> None: + """Once a chunk is written, its state groups are no longer pending + deletion. If the write fails, their pending deletion is pushed + back, as they may now be referenced.""" + event, context = self._build("hello") + assert context.state_group is not None + self.get_success( + self.store.db_pool.simple_insert( + "state_groups_pending_deletion", + {"state_group": context.state_group, "insertion_ts": 0}, + ) + ) + before = self._pending_deletion(context.state_group) + assert before is not None + + failing, failing_context = self._build("partial") + failing_context.partial_state = True + self.reactor.advance(1) + self.get_failure( + self.persistence.persist_events([(failing, failing_context)]), + PartialStateConflictError, + ) + after = self._pending_deletion(context.state_group) + assert after is not None + self.assertGreater(after[0], before[0]) + self.assertEqual(after[1], self.clock.time_msec()) + + self.get_success(self.persistence.persist_events([(event, context)])) + self.assertIsNone(self._pending_deletion(context.state_group)) + + def test_replaced_event_and_chunks(self) -> None: + """A batch whose first event repeats a persisted transaction ID, and + whose other 150 go in two chunks: each chunk hands Python the events + it wrote.""" + self.login("alice", "pass", device_id="DEVICE") + original = self._build("original", txn_id="txn") + self.get_success(self.persistence.persist_events([original])) + + repeat = self._build("repeat", txn_id="txn") + new = [self._build(f"message {i}") for i in range(150)] + chunks = self.persister.chunks + persisted, _ = self.get_success(self.persistence.persist_events([repeat, *new])) + self.assertEqual(self.persister.chunks, chunks + 2) + effects = self.persister.last_effects + assert effects is not None + self.assertEqual(len(effects.stream_ids), 50) + self.assertEqual( + [event.event_id for event in persisted], + [original[0].event_id] + [event.event_id for event, _ in new], + ) + self.assertIsNone( + self.get_success(self.store.get_event(repeat[0].event_id, allow_none=True)) + ) + orderings = [self._event_row(event.event_id)[0] for event, _ in new] + self.assertEqual(orderings, sorted(orderings)) + for i, (event, _) in enumerate(new): + stored = self.get_success(self.store.get_event(event.event_id)) + self.assertEqual(stored.content["body"], f"message {i}") + + def test_same_event_in_two_chunks(self) -> None: + """An event in two chunks of a batch is written once, and every + ordering the batch took is finished.""" + first = self._build("first") + rest = [self._build(f"message {i}") for i in range(100)] + self.get_success(self.persistence.persist_events([first, *rest, first])) + rows = self.get_success( + self.store.db_pool.simple_select_onecol( + "events", {"event_id": first[0].event_id}, "stream_ordering" + ) + ) + self.assertEqual(len(rows), 1) + effects = self.persister.last_effects + assert effects is not None + self.assertEqual(len(effects.stream_ids), 2) + self.assertGreaterEqual( + self.store._stream_id_gen.get_current_token(), max(effects.stream_ids) + ) + + def test_logcontext_kept(self) -> None: + """Awaiting the batch resumes in the caller's logcontext, so what + follows it is accounted to the caller.""" + event, context = self._build("hello") + seen = [] + + async def persist() -> None: + with LoggingContext(name="persist", server_name="test") as ctx: + await self.persistence._persist_event_batch_in_rust( + self.room_id, [(event, context)], False + ) + seen.append(current_context() is ctx) + + self.get_success(defer.ensureDeferred(persist())) + self.assertEqual(seen, [True]) + def test_gil_count(self) -> None: """Taking the stream orderings and setting them on the events takes the GIL at no persist site.""" @@ -1360,64 +1586,70 @@ class RustFullTestCase(unittest.HomeserverTestCase): self.assertGreaterEqual( self.store._stream_id_gen.get_current_token(), stream_ordering ) + # `batch_abort` released the state groups. + rows = self.get_success( + self.store.db_pool.simple_select_list( + "state_groups_persisting", None, ("state_group",) + ) + ) + self.assertEqual(rows, []) def test_cancelled_persist_applies_its_effects(self) -> None: - """Cancelling the persist while Rust takes the orderings and writes - the chunk waits for both, and the orderings are finished.""" + """Cancelling the batch while Rust writes a chunk waits for Rust, + so the room's queue cannot start another batch alongside, and the + chunk's effects are applied and its orderings finished.""" gate: list[tuple[Any, Any, defer.Deferred]] = [] blocking = self.persister._persister class Gated: - def persist_chunk(self, *args: Any, **kwargs: Any) -> Any: + def batch_next(self, *args: Any, **kwargs: Any) -> Any: + # Only the first step, which writes the chunk. + if gate: + return blocking.batch_next(*args, **kwargs) d: defer.Deferred = defer.Deferred() gate.append((args, kwargs, d)) - return d + # As the persister's Deferreds do. + return make_deferred_yieldable(d) def __getattr__(self, name: str) -> Any: return getattr(blocking, name) self.persister._persister = Gated() - persist_events = self.hs.get_datastores().persist_events - assert persist_events is not None - original = persist_events._persist_events_in_rust + original = self.persistence._persist_event_batch_in_rust cancelled = [] - async def cancel_then_wait(*args: Any, **kwargs: Any) -> None: - d = defer.ensureDeferred(original(*args, **kwargs)) + async def cancel_then_wait(*args: Any) -> dict[str, str]: + d = defer.ensureDeferred(original(*args)) d.cancel() try: - await make_deferred_yieldable(d) + return await make_deferred_yieldable(d) except defer.CancelledError: cancelled.append(True) + return {} + event, context = self._build("hello") with patch.object( - persist_events, "_persist_events_in_rust", side_effect=cancel_then_wait + self.persistence, + "_persist_event_batch_in_rust", + side_effect=cancel_then_wait, ): - channel = make_request( - self.reactor, - self.site, - "PUT", - f"/rooms/{self.room_id}/send/m.room.message/txn1", - {"msgtype": "m.text", "body": "hello"}, - access_token=self.token, - await_result=False, + persisted = defer.ensureDeferred( + self.persistence.persist_events([(event, context)]) ) self.pump() self.assertEqual(len(gate), 1) - # Cancelled, but waiting for the write. + # Cancelled, but waiting for Rust. self.assertEqual(cancelled, []) args, kwargs, d = gate[0] - blocking.persist_chunk(*args, **kwargs).chainDeferred(d) - channel.await_result() - self.assertEqual(channel.code, 200, channel.json_body) - event_id = channel.json_body["event_id"] + blocking.batch_next(*args, **kwargs).chainDeferred(d) + self.get_success(persisted) self.assertEqual(cancelled, [True]) effects = self.persister.last_effects assert effects is not None (stream_ordering,) = effects.stream_ids - self.assertEqual(self._event_row(event_id)[0], stream_ordering) + self.assertEqual(self._event_row(event.event_id)[0], stream_ordering) self.assertGreaterEqual( self.store._stream_id_gen.get_current_token(), stream_ordering )