Prepare chunks in Rust at the shadow and prepare levels

`Persister.prepare_chunk` runs the steps of the previous commits for a
chunk: the extremities and state delta, then the sliding sync table
changes and sticky events, then the chain cover index. The chain cover
step comes last as the only one that writes, so a failure in an earlier
step leaves nothing for Python to find. Sliding sync and sticky events
now come before the stream orderings are allocated, as nothing there
reads them. Each step reads the clock when it needs the time.

`_persist_event_batch` dispatches on `rust_event_persistence`:

- `shadow`: Rust runs first with its chain cover write rolled back, then
  Python. The results are compared and Python's are used. The comparison
  ignores orders Python leaves to chance and the IDs of new chains. Both
  rejecting the chunk with the same error is a match.
- `prepare` and above: Rust's results are used. A `SynapseError` is
  raised as Python would raise it. Python computes the chunk when Rust
  cannot resolve state in the room, or on any other Rust error.

Trial runs strict by default (`SYNAPSE_TEST_RUST_PERSIST_STRICT`), so a
mismatch or Rust error fails the test. Rust's counts and state res stats
are recorded by Python once it uses the result.

At `prepare`, the `_calculate_state_and_extrem` and
`calculate_chain_cover_index_for_events` timings give way to
`rust_prepare_chunk`. Python's state resolution and state group caches
are not filled, so creating a room costs one more transaction
(`RoomsCreateTestCase` counts it) and the first local event after a
fork resolves state again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LXMkcZVkcZhi21augBff4a

changelog.d/rust-persist-prepare.misc

@@ -1 +1 @@
-Add the Rust reads for persisting events before the write transaction, behind `experimental_features.rust_event_persistence`.
+Compute what persisting events works out before the write transaction in Rust, behind `experimental_features.rust_event_persistence`.

rust/src/database/postgres/write.rs

@@ -407,6 +407,31 @@ impl Drop for WriteTxn {
 pub async fn run_write_interaction<'env, T, F>(
     pool: &ConnectionPool,
     name: &str,
+    f: F,
+) -> Result<T, WriteError>
+where
+    F: for<'t> FnMut(&'t WriteTxn, &'t &'env ()) -> BoxFuture<'t, Result<T, WriteError>>,
+{
+    run_interaction(pool, name, true, f).await
+}
+
+/// As [`run_write_interaction`], but rolls the transaction back even when
+/// `f` succeeds. For a dry run whose writes must not be seen.
+pub async fn run_rolled_back_interaction<'env, T, F>(
+    pool: &ConnectionPool,
+    name: &str,
+    f: F,
+) -> Result<T, WriteError>
+where
+    F: for<'t> FnMut(&'t WriteTxn, &'t &'env ()) -> BoxFuture<'t, Result<T, WriteError>>,
+{
+    run_interaction(pool, name, false, f).await
+}
+
+async fn run_interaction<'env, T, F>(
+    pool: &ConnectionPool,
+    name: &str,
+    commit: bool,
     mut f: F,
 ) -> Result<T, WriteError>
 where
@@ -414,7 +439,7 @@ where
 {
     let mut attempt = 1;
     loop {
-        let result = run_once(pool, &mut f).await;
+        let result = run_once(pool, commit, &mut f).await;
         match result {
             Err(error) if error.is_retryable() && attempt < MAX_ATTEMPTS => {
                 warn!("[TXN RETRY] {{{name}}} {error} {attempt}/{MAX_ATTEMPTS}");
@@ -425,7 +450,11 @@ where
     }
 }
 
-async fn run_once<'env, T, F>(pool: &ConnectionPool, f: &mut F) -> Result<T, WriteError>
+async fn run_once<'env, T, F>(
+    pool: &ConnectionPool,
+    commit: bool,
+    f: &mut F,
+) -> Result<T, WriteError>
 where
     F: for<'t> FnMut(&'t WriteTxn, &'t &'env ()) -> BoxFuture<'t, Result<T, WriteError>>,
 {
@@ -437,10 +466,14 @@ where
     })?;
     let txn = WriteTxn::begin(conn).await?;
     match f(&txn, &&()).await {
-        Ok(value) => {
+        Ok(value) if commit => {
             txn.commit().await?;
             Ok(value)
         }
+        Ok(value) => {
+            txn.rollback().await;
+            Ok(value)
+        }
         Err(error) => {
             txn.rollback().await;
             Err(error)

rust/src/persist/chain_cover.rs

@@ -51,7 +51,7 @@ use super::reader::{self, AuthEventChain, ChainLink, ChainToCalculate};
 use crate::{
     database::postgres::{
         pool::ConnectionPool,
-        write::{run_write_interaction, WriteTxn},
+        write::{run_rolled_back_interaction, run_write_interaction, WriteTxn},
     },
     events::Event,
 };
@@ -223,10 +223,15 @@ impl ChainCoverStore for WriteTxn {
 /// are left out. The result can also hold events that are not in `events`:
 /// those put off by an earlier chunk, and auth events persisted out of
 /// band, once their auth events are known.
+///
+/// A `dry_run` rolls the transaction back, so that nothing it puts off is
+/// written. Its chain IDs are still used up, as the sequence is not
+/// transactional.
 pub async fn calculate_chain_cover_index_for_events(
     pool: &ConnectionPool,
     room_id: &str,
     events: &[ChainCoverEvent],
+    dry_run: bool,
 ) -> anyhow::Result<HashMap<Arc<str>, NewEventChainLinks>> {
     // Ensure there are no duplicates.
     let mut seen = HashSet::new();
@@ -238,15 +243,26 @@ pub async fn calculate_chain_cover_index_for_events(
         return Ok(HashMap::new());
     }
     let state_events = &state_events;
-    Ok(
-        run_write_interaction(pool, "_calculate_chain_cover_index_for_events", |txn, _| {
+    let name = "_calculate_chain_cover_index_for_events";
+    // The closure is written out twice: bound to a variable, it would lose
+    // the lifetimes the interactions need.
+    Ok(if dry_run {
+        run_rolled_back_interaction(pool, name, |txn, _| {
             async move {
                 Ok(calculate_chain_cover_index_for_events_txn(txn, room_id, state_events).await?)
             }
             .boxed()
         })
-        .await?,
-    )
+        .await?
+    } else {
+        run_write_interaction(pool, name, |txn, _| {
+            async move {
+                Ok(calculate_chain_cover_index_for_events_txn(txn, room_id, state_events).await?)
+            }
+            .boxed()
+        })
+        .await?
+    })
 }
 
 /// `calculate_chain_cover_index_for_events_txn`, over deduplicated state

rust/src/persist/metrics.rs

@@ -79,6 +79,39 @@ impl PersistMetrics {
 }
 
 impl PersistCounts {
+    /// Add `other` to these counts. Its persisted position, if it has one,
+    /// is the newer.
+    pub fn add(&mut self, other: PersistCounts) {
+        self.persisted_events += other.persisted_events;
+        for (key, count) in other.persisted_events_by_type {
+            *self.persisted_events_by_type.entry(key).or_default() += count;
+        }
+        self.state_delta += other.state_delta;
+        self.state_delta_single_event += other.state_delta_single_event;
+        self.state_delta_reuse_delta += other.state_delta_reuse_delta;
+        self.state_resolutions += other.state_resolutions;
+        self.potential_times_prune_extremities += other.potential_times_prune_extremities;
+        self.times_pruned_extremities += other.times_pruned_extremities;
+        for (mine, theirs) in [
+            (&mut self.forward_extremities, other.forward_extremities),
+            (
+                &mut self.stale_forward_extremities,
+                other.stale_forward_extremities,
+            ),
+            (
+                &mut self.state_dag_forward_extremities,
+                other.state_dag_forward_extremities,
+            ),
+        ] {
+            for (value, count) in theirs {
+                *mine.entry(value).or_default() += count;
+            }
+        }
+        if other.persisted_position.is_some() {
+            self.persisted_position = other.persisted_position;
+        }
+    }
+
     /// The counts as the dict `flush_rust_persist_metrics` reads.
     pub fn to_py<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
         let dict = PyDict::new(py);
@@ -112,6 +145,34 @@ impl PersistCounts {
 mod tests {
     use super::*;
 
+    #[test]
+    fn add_sums_counts_and_keeps_the_newer_position() {
+        let mut total = PersistCounts {
+            state_delta: 1,
+            persisted_position: Some(5),
+            ..Default::default()
+        };
+        total.forward_extremities.insert(1, 1);
+        let mut more = PersistCounts {
+            state_delta: 2,
+            times_pruned_extremities: 1,
+            ..Default::default()
+        };
+        more.forward_extremities.insert(1, 2);
+        more.forward_extremities.insert(3, 1);
+
+        total.add(more.clone());
+        assert_eq!(total.state_delta, 3);
+        assert_eq!(total.times_pruned_extremities, 1);
+        assert_eq!(total.forward_extremities, BTreeMap::from([(1, 3), (3, 1)]));
+        // `more` had no position, so the old one stays.
+        assert_eq!(total.persisted_position, Some(5));
+
+        more.persisted_position = Some(7);
+        total.add(more);
+        assert_eq!(total.persisted_position, Some(7));
+    }
+
     #[test]
     fn take_starts_again() {
         let metrics = PersistMetrics::default();

rust/src/persist/mod.rs

@@ -24,6 +24,8 @@
 //!   it commits.
 //! - [`extremities`]: the room's forward extremities after a chunk.
 //! - [`metrics`]: counts for the Python metrics, flushed when scraped.
+//! - [`prepare`]: everything a chunk needs before its write transaction,
+//!   built from the steps above.
 //! - [`python`]: the [`python::Persister`] and test helpers.
 //! - [`reader`]: what persisting reads before its write transaction.
 //! - [`sliding_sync`]: what a chunk changes in the sliding sync tables.
@@ -40,6 +42,7 @@ pub mod context;
 pub mod effects;
 pub mod extremities;
 pub mod metrics;
+pub mod prepare;
 pub mod python;
 pub mod reader;
 pub mod sliding_sync;
@@ -50,6 +53,7 @@ pub mod sticky;
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let child_module = PyModule::new(py, "persist")?;
     child_module.add_class::<python::Persister>()?;
+    child_module.add_class::<python::PreparedChunk>()?;
     child_module.add_class::<effects::PersistEffects>()?;
     child_module.add_class::<reader::PyPersistReaderForTest>()?;
     child_module.add_function(wrap_pyfunction!(

rust/src/persist/prepare.rs

@@ -0,0 +1,234 @@
+/*
+ * 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:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! What `_persist_event_batch` computes for a chunk before its write
+//! transaction: the forward extremities and state delta, then the sliding
+//! sync table changes and the sticky events to un-soft-fail, then the chain
+//! cover index.
+//!
+//! Python computes the sliding sync changes and the sticky events inside
+//! `_persist_events_and_state_updates`, once the chunk's stream orderings
+//! are allocated. Neither reads them, so here they come before the
+//! allocation, where they cannot hold up other rooms' events.
+//!
+//! The chain cover step comes last because it is the only one that writes:
+//! it puts off events whose auth events are missing. A failure in an
+//! earlier step leaves nothing behind for Python to find when it computes
+//! the chunk again. A dry run, for shadow mode, rolls that write back.
+//!
+//! Python's computation fills two Python caches on the way: the state
+//! resolution handler's `_state_cache` of resolved state group sets, and
+//! the state group cache. This does not, so when Rust's result is used the
+//! next Python read of the same groups misses. Creating a room costs one
+//! more `_get_state_groups_from_groups` transaction, and the first local
+//! event after a fork resolves state again.
+
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+};
+
+use pyo3::prelude::*;
+
+use super::{
+    chain_cover::{self, ChainCoverEvent, NewEventChainLinks},
+    context::PersistContext,
+    extremities::NewEvent,
+    metrics::PersistCounts,
+    reader::{self, PersistReader},
+    sliding_sync::{self, SlidingSyncEvent, SlidingSyncTableChanges},
+    state_delta::{self, NativeStateResolver, StateAndExtremities, StateResFallback},
+    sticky,
+};
+use crate::{
+    clock::Clock,
+    database::postgres::pool::ConnectionPool,
+    events::Event,
+    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.
+pub struct PrepareInput {
+    pub room_id: String,
+    pub backfilled: bool,
+    /// Whether every event is in a room version with MSC4242 state DAGs.
+    pub state_dag: bool,
+    pub events: Vec<NewEvent>,
+    pub event_map: EventMap,
+    pub chain_cover_events: Vec<ChainCoverEvent>,
+    pub sliding_sync_events: Vec<SlidingSyncEvent>,
+    pub sticky_events: Vec<Event>,
+}
+
+impl PrepareInput {
+    pub fn new(
+        py: Python<'_>,
+        room_id: String,
+        events_and_contexts: &[(Py<Event>, PersistContext)],
+        backfilled: bool,
+    ) -> PyResult<Self> {
+        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();
+        Ok(PrepareInput {
+            room_id,
+            backfilled,
+            state_dag: chunk_events
+                .iter()
+                .all(|event| event.room_version().msc4242_state_dags),
+            events,
+            event_map,
+            chain_cover_events: chunk_events
+                .iter()
+                .filter_map(|event| ChainCoverEvent::new(event))
+                .collect(),
+            sliding_sync_events: chunk_events
+                .iter()
+                .filter_map(|event| SlidingSyncEvent::new(event))
+                .collect(),
+            sticky_events: chunk_events
+                .iter()
+                .map(|event| event.shallow_clone())
+                .collect(),
+        })
+    }
+}
+
+/// What the chunk's write transaction takes.
+#[derive(Debug)]
+pub struct PreparedChunk {
+    /// The extremities and state delta. All `None` for a backfilled chunk.
+    pub state: StateAndExtremities,
+    pub new_event_links: HashMap<Arc<str>, NewEventChainLinks>,
+    /// `None` when the state does not change.
+    pub sliding_sync: Option<SlidingSyncTableChanges>,
+    pub sticky_events_to_un_soft_fail: HashSet<Arc<str>>,
+}
+
+/// The outcome of [`prepare_chunk`].
+#[derive(Debug)]
+pub enum Prepared {
+    Chunk(Box<PreparedChunk>),
+    /// The Rust state resolution cannot resolve the room: Python should
+    /// compute the chunk. Nothing was written.
+    StateResFallback(String),
+}
+
+/// What [`prepare_chunk`] needs besides the chunk.
+pub struct PrepareOptions<'a> {
+    pub server_name: &'a str,
+    /// Read when a step needs the time, as Python reads its clock.
+    pub clock: &'a Clock,
+    /// `experimental_features.msc4354_enabled`: whether to look for sticky
+    /// events to un-soft-fail.
+    pub sticky: bool,
+    /// Roll back the chain cover step's write.
+    pub dry_run: bool,
+}
+
+/// Compute what the chunk's write transaction takes, as
+/// `_persist_event_batch` does. Adds to `counts` what Python counts on the
+/// way. The caller should drop them on a [`Prepared::StateResFallback`],
+/// since Python counts the chunk again.
+pub async fn prepare_chunk(
+    reader: &PersistReader<'_>,
+    resolver: &NativeStateResolver,
+    pool: &ConnectionPool,
+    input: &PrepareInput,
+    options: &PrepareOptions<'_>,
+    counts: &mut PersistCounts,
+) -> anyhow::Result<Prepared> {
+    let room_id = input.room_id.as_str();
+    let now_ms = || i64::try_from(options.clock.now_millis());
+    let state = if input.backfilled {
+        StateAndExtremities::default()
+    } else {
+        let result = if input.state_dag {
+            state_delta::process_state_dag_forward_extremities_and_state_delta(
+                reader,
+                resolver,
+                room_id,
+                options.server_name,
+                &input.events,
+                &input.event_map,
+                now_ms()?,
+                counts,
+            )
+            .await
+        } else {
+            state_delta::calculate_new_forward_extremities_and_state_delta(
+                reader,
+                resolver,
+                room_id,
+                options.server_name,
+                &input.events,
+                &input.event_map,
+                now_ms()?,
+                counts,
+            )
+            .await
+        };
+        match result {
+            Ok(state) => state,
+            Err(e) => match e.downcast::<StateResFallback>() {
+                Ok(fallback) => return Ok(Prepared::StateResFallback(fallback.0)),
+                Err(e) => return Err(e),
+            },
+        }
+    };
+
+    let mut sliding_sync = None;
+    let mut sticky_events_to_un_soft_fail = HashSet::new();
+    if let Some(delta) = &state.state_delta {
+        sliding_sync = Some(
+            sliding_sync::calculate_sliding_sync_table_changes(
+                reader,
+                room_id,
+                options.server_name,
+                &input.sliding_sync_events,
+                delta,
+            )
+            .await?,
+        );
+        if options.sticky {
+            sticky_events_to_un_soft_fail = sticky::compute_sticky_events_to_un_soft_fail(
+                reader,
+                room_id,
+                &input.sticky_events,
+                delta,
+                now_ms()?,
+            )
+            .await?;
+        }
+    }
+
+    let new_event_links = chain_cover::calculate_chain_cover_index_for_events(
+        pool,
+        room_id,
+        &input.chain_cover_events,
+        options.dry_run,
+    )
+    .await?;
+
+    Ok(Prepared::Chunk(Box::new(PreparedChunk {
+        state,
+        new_event_links,
+        sliding_sync,
+        sticky_events_to_un_soft_fail,
+    })))
+}

rust/src/persist/python.rs

@@ -16,21 +16,31 @@
 //! The Python face of event persistence: the [`Persister`] that
 //! `RustHandlers.persister` hands out, and helpers for tests.
 
-use std::sync::Arc;
+use std::{
+    collections::{HashMap, HashSet},
+    sync::{Arc, Mutex},
+};
 
+use futures::{future::BoxFuture, FutureExt};
 use http::StatusCode;
 use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyDict};
 
 use super::{
-    context::PersistContext, extremities::MissingPrevStateEvents, metrics::PersistMetrics,
+    context::PersistContext,
+    extremities::MissingPrevStateEvents,
+    metrics::{PersistCounts, PersistMetrics},
+    prepare::{self, PrepareInput, PrepareOptions, Prepared},
+    reader::{self, PersistReader, PyDelta, PySlidingSync},
+    state_delta::NativeStateResolver,
 };
 use crate::{
+    deferred::{create_deferred_with, OnCancel},
     errors::SynapseError,
     event_auth::AuthFailure,
     events::Event,
     homeserver::HomeServer,
     runtime::RustRuntime,
-    state::ConflictCache,
+    state::{ConflictCache, StateResStats},
     storage::{
         native_pool::NativePool,
         store::Store,
@@ -44,11 +54,12 @@ use crate::{
 ///
 /// Holds what every level needs: the pool, the shared caches, the stream ID
 /// trackers of the events, backfill and cache streams (shared with their
-/// Python `MultiWriterIdGenerator`s) and the metrics. The entry points come
-/// with the levels that use them.
+/// Python `MultiWriterIdGenerator`s) and the metrics.
+///
+/// Each entry point returns a Deferred and has a `_blocking` twin for
+/// tests (see `_BlockingRustCalls`).
 #[pyclass(frozen)]
 pub struct Persister {
-    #[allow(dead_code)] // Read by the levels to come.
     pub(super) inner: Arc<PersisterInner>,
 }
 
@@ -68,6 +79,8 @@ pub(super) struct PersisterInner {
     pub metrics: PersistMetrics,
     /// The Rust state resolution handler's conflict cache.
     pub conflict_cache: Arc<ConflictCache>,
+    /// `experimental_features.msc4354_enabled`.
+    pub msc4354_enabled: bool,
 }
 
 impl Persister {
@@ -87,6 +100,7 @@ impl Persister {
             return Ok(None);
         }
 
+        let config = homeserver.config(py)?;
         let main = homeserver.main_datastore(py)?;
         let tracker = |attr: &Bound<'_, pyo3::types::PyString>| -> PyResult<Arc<StreamIdTracker>> {
             let id_gen = main.getattr(attr)?;
@@ -108,7 +122,7 @@ impl Persister {
                 runtime,
                 sync: Arc::new(SyncStore::new(
                     Arc::clone(&store.caches),
-                    SyncStoreConfig::from_config(&homeserver.config(py)?),
+                    SyncStoreConfig::from_config(&config),
                 )),
                 store,
                 events_stream: tracker(intern!(py, "_stream_id_gen"))?,
@@ -118,9 +132,62 @@ impl Persister {
                 instance_name,
                 metrics: PersistMetrics::default(),
                 conflict_cache,
+                msc4354_enabled: config.experimental.msc4354_enabled,
             }),
         }))
     }
+
+    /// Convert the chunk while the GIL is held, and return the future that
+    /// prepares it without the GIL.
+    fn start_prepare_chunk(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        backfilled: bool,
+        shadow: bool,
+    ) -> PyResult<BoxFuture<'static, PyResult<PreparedChunk>>> {
+        let input = PrepareInput::new(py, room_id, &events_and_contexts, backfilled)?;
+        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,
+            };
+            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)),
+            })
+        }
+        .boxed())
+    }
+
+    /// Run `future` to completion on the runtime, releasing the GIL. For the
+    /// `_blocking` entry points.
+    fn block_on<T: Send>(
+        &self,
+        py: Python<'_>,
+        future: BoxFuture<'static, PyResult<T>>,
+    ) -> PyResult<T> {
+        let handle = self.inner.runtime.tokio_handle()?;
+        py.detach(|| handle.block_on(future))
+    }
 }
 
 #[pymethods]
@@ -130,6 +197,151 @@ impl Persister {
     fn take_metrics<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
         self.inner.metrics.take().to_py(py)
     }
+
+    /// What `_persist_event_batch` computes for a chunk of one room's events
+    /// before the write transaction: a Deferred resolving to a
+    /// [`PreparedChunk`]. In `shadow` mode the chain cover step's write is
+    /// rolled back. Nothing is counted until [`Self::record_prepare_counts`].
+    /// Cancelling the Deferred stops the computation.
+    #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))]
+    fn prepare_chunk<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        backfilled: bool,
+        shadow: bool,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let future =
+            self.start_prepare_chunk(py, room_id, events_and_contexts, backfilled, shadow)?;
+        create_deferred_with(py, &self.inner.runtime, future, OnCancel::Abort)
+    }
+
+    /// Add what preparing `chunk` counted to the metrics, once Python has
+    /// decided to use it. Python counts the chunks it computes itself. Does
+    /// nothing the second time.
+    fn record_prepare_counts(&self, chunk: &PreparedChunk) {
+        let counts = chunk
+            .counts
+            .lock()
+            .unwrap_or_else(|e| e.into_inner())
+            .take();
+        if let Some(counts) = counts {
+            self.inner.metrics.record(|total| total.add(counts));
+        }
+    }
+
+    /// As [`Self::prepare_chunk`], but blocks.
+    #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))]
+    fn prepare_chunk_blocking(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        backfilled: bool,
+        shadow: bool,
+    ) -> PyResult<PreparedChunk> {
+        let future =
+            self.start_prepare_chunk(py, room_id, events_and_contexts, backfilled, shadow)?;
+        self.block_on(py, future)
+    }
+}
+
+/// 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.
+#[pyclass(frozen)]
+pub struct PreparedChunk {
+    prepared: Prepared,
+    /// What preparing counted, until it is recorded.
+    counts: Mutex<Option<PersistCounts>>,
+}
+
+impl PreparedChunk {
+    fn chunk(&self) -> Option<&prepare::PreparedChunk> {
+        match &self.prepared {
+            Prepared::Chunk(chunk) => Some(chunk),
+            Prepared::StateResFallback(_) => None,
+        }
+    }
+}
+
+#[pymethods]
+impl PreparedChunk {
+    /// Why the Rust state resolution could not resolve the room.
+    #[getter]
+    fn state_res_fallback(&self) -> Option<&str> {
+        match &self.prepared {
+            Prepared::Chunk(_) => None,
+            Prepared::StateResFallback(reason) => Some(reason),
+        }
+    }
+
+    #[getter]
+    fn new_forward_extremities(&self) -> Option<HashSet<String>> {
+        let ids = self.chunk()?.state.new_forward_extremities.clone()?;
+        Some(reader::strings(ids))
+    }
+
+    /// `(to_delete, to_insert, no_longer_in_room)`.
+    #[getter]
+    fn state_delta<'py>(&self, py: Python<'py>) -> PyResult<Option<PyDelta<'py>>> {
+        self.chunk()
+            .and_then(|chunk| chunk.state.state_delta.as_ref())
+            .map(|delta| reader::delta_to_py(py, delta))
+            .transpose()
+    }
+
+    #[getter]
+    fn new_state_dag_extremities(&self) -> Option<HashSet<String>> {
+        let ids = self.chunk()?.state.new_state_dag_extremities.clone()?;
+        Some(reader::strings(ids))
+    }
+
+    #[getter]
+    fn state_res_stats(&self) -> Option<StateResStats> {
+        self.chunk()?.state.state_res_stats.clone()
+    }
+
+    /// Event ID to `(chain_id, sequence_number, links)`.
+    #[getter]
+    #[allow(clippy::type_complexity)]
+    fn new_event_links(&self) -> HashMap<String, (i64, i64, Vec<(i64, i64)>)> {
+        self.chunk()
+            .map(|chunk| {
+                chunk
+                    .new_event_links
+                    .iter()
+                    .map(|(event_id, links)| {
+                        (
+                            event_id.to_string(),
+                            (links.chain_id, links.sequence_number, links.links.clone()),
+                        )
+                    })
+                    .collect()
+            })
+            .unwrap_or_default()
+    }
+
+    /// As `PersistReaderForTest.calculate_sliding_sync_table_changes`
+    /// returns it. `None` when the state does not change.
+    #[getter]
+    fn sliding_sync_table_changes<'py>(
+        &self,
+        py: Python<'py>,
+    ) -> PyResult<Option<PySlidingSync<'py>>> {
+        self.chunk()
+            .and_then(|chunk| chunk.sliding_sync.as_ref())
+            .map(|changes| reader::sliding_sync_to_py(py, changes))
+            .transpose()
+    }
+
+    #[getter]
+    fn sticky_events_to_un_soft_fail(&self) -> HashSet<String> {
+        self.chunk()
+            .map(|chunk| reader::strings(chunk.sticky_events_to_un_soft_fail.iter().cloned()))
+            .unwrap_or_default()
+    }
 }
 
 /// Convert an error from persisting to the Python exception the Python

rust/src/persist/reader.rs

@@ -687,7 +687,7 @@ fn arcs(ids: Vec<String>) -> Vec<Arc<str>> {
     ids.into_iter().map(Arc::from).collect()
 }
 
-fn strings(ids: impl IntoIterator<Item = Arc<str>>) -> HashSet<String> {
+pub(super) fn strings(ids: impl IntoIterator<Item = Arc<str>>) -> HashSet<String> {
     ids.into_iter().map(|id| id.to_string()).collect()
 }
 
@@ -1053,14 +1053,16 @@ impl PyPersistReaderForTest {
         })
     }
 
-    /// The chain cover step, in a write transaction of its own. Event ID to
-    /// `(chain_id, sequence_number, links)`.
+    /// The chain cover step, in a write transaction of its own, rolled back
+    /// for a `dry_run`. Event ID to `(chain_id, sequence_number, links)`.
     #[allow(clippy::type_complexity)]
+    #[pyo3(signature = (room_id, events, *, dry_run = false))]
     fn calculate_chain_cover_index_for_events(
         &self,
         py: Python<'_>,
         room_id: String,
         events: Vec<Py<Event>>,
+        dry_run: bool,
     ) -> PyResult<HashMap<String, (i64, i64, Vec<(i64, i64)>)>> {
         let events: Vec<chain_cover::ChainCoverEvent> = events
             .iter()
@@ -1070,7 +1072,7 @@ impl PyPersistReaderForTest {
             .detach(|| {
                 self.handle
                     .block_on(chain_cover::calculate_chain_cover_index_for_events(
-                        &self.pool, &room_id, &events,
+                        &self.pool, &room_id, &events, dry_run,
                     ))
             })
             .map_err(persist_error_to_py)?;
@@ -1360,13 +1362,7 @@ impl PyPersistReaderForTest {
         to_delete: Vec<(String, String)>,
         to_insert: HashMap<(String, String), String>,
         no_longer_in_room: bool,
-    ) -> PyResult<(
-        Option<i64>,
-        Bound<'py, PyDict>,
-        Bound<'py, PyDict>,
-        Vec<(String, String, String, String)>,
-        Vec<String>,
-    )> {
+    ) -> PyResult<PySlidingSync<'py>> {
         let events: Vec<sliding_sync::SlidingSyncEvent> = events
             .iter()
             .filter_map(|event| sliding_sync::SlidingSyncEvent::new(event.bind(py).get()))
@@ -1392,36 +1388,7 @@ impl PyPersistReaderForTest {
             }
             .boxed()
         })?;
-        let shared =
-            state_values_to_py(py, &changes.membership_snapshot_shared_insert_values.state)?;
-        if let Some(has_known_state) = changes
-            .membership_snapshot_shared_insert_values
-            .has_known_state
-        {
-            shared.set_item("has_known_state", has_known_state)?;
-        }
-        Ok((
-            changes.joined_room_bump_stamp_to_fully_insert,
-            state_values_to_py(py, &changes.joined_room_updates)?,
-            shared,
-            changes
-                .to_insert_membership_snapshots
-                .iter()
-                .map(|info| {
-                    (
-                        info.user_id.to_string(),
-                        info.sender.to_string(),
-                        info.membership_event_id.to_string(),
-                        info.membership.to_string(),
-                    )
-                })
-                .collect(),
-            changes
-                .to_delete_membership_snapshots
-                .iter()
-                .map(|user_id| user_id.to_string())
-                .collect(),
-        ))
+        sliding_sync_to_py(py, &changes)
     }
 
     /// The passing event IDs.
@@ -1482,11 +1449,58 @@ fn state_values_to_py<'py>(
     Ok(dict)
 }
 
+/// Sliding sync table changes as `(bump_stamp_to_fully_insert,
+/// joined_room_updates, membership_snapshot_shared_insert_values,
+/// to_insert_membership_snapshots, to_delete_membership_snapshots)`, each
+/// snapshot as `(user_id, sender, membership_event_id, membership)`.
+pub(super) type PySlidingSync<'py> = (
+    Option<i64>,
+    Bound<'py, PyDict>,
+    Bound<'py, PyDict>,
+    Vec<(String, String, String, String)>,
+    Vec<String>,
+);
+
+pub(super) fn sliding_sync_to_py<'py>(
+    py: Python<'py>,
+    changes: &sliding_sync::SlidingSyncTableChanges,
+) -> PyResult<PySlidingSync<'py>> {
+    let shared = state_values_to_py(py, &changes.membership_snapshot_shared_insert_values.state)?;
+    if let Some(has_known_state) = changes
+        .membership_snapshot_shared_insert_values
+        .has_known_state
+    {
+        shared.set_item("has_known_state", has_known_state)?;
+    }
+    Ok((
+        changes.joined_room_bump_stamp_to_fully_insert,
+        state_values_to_py(py, &changes.joined_room_updates)?,
+        shared,
+        changes
+            .to_insert_membership_snapshots
+            .iter()
+            .map(|info| {
+                (
+                    info.user_id.to_string(),
+                    info.sender.to_string(),
+                    info.membership_event_id.to_string(),
+                    info.membership.to_string(),
+                )
+            })
+            .collect(),
+        changes
+            .to_delete_membership_snapshots
+            .iter()
+            .map(|user_id| user_id.to_string())
+            .collect(),
+    ))
+}
+
 /// A delta as Python's test reads it: `(to_delete, to_insert,
 /// no_longer_in_room)`.
-type PyDelta<'py> = (Vec<(String, String)>, Bound<'py, PyDict>, bool);
+pub(super) type PyDelta<'py> = (Vec<(String, String)>, Bound<'py, PyDict>, bool);
 
-fn delta_to_py<'py>(py: Python<'py>, delta: &DeltaState) -> PyResult<PyDelta<'py>> {
+pub(super) fn delta_to_py<'py>(py: Python<'py>, delta: &DeltaState) -> PyResult<PyDelta<'py>> {
     Ok((
         delta
             .to_delete
@@ -1508,7 +1522,7 @@ 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.
-fn chunk(
+pub(super) fn chunk(
     py: Python<'_>,
     events_and_contexts: &[(Py<Event>, PersistContext)],
 ) -> PyResult<(Vec<NewEvent>, EventMap)> {

synapse/config/experimental.py

@@ -410,12 +410,13 @@ class ExperimentalConfig(Config):
         #  - "off": all Python.
         #  - "shadow": Rust also computes what is worked out before the write
         #    transaction (forward extremities, the current state delta, the
-        #    chain cover index), and differences from Python are logged.
+        #    chain cover index, the sliding sync table changes and the sticky
+        #    events to un-soft-fail), and differences from Python are logged.
         #  - "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 every level
-        # other than "off" only builds the Rust persister.
+        # The levels are being built up in that order. So far "transaction"
+        # and "full" act as "prepare".
         rust_event_persistence = experimental.get("rust_event_persistence", "off")
         # YAML reads an unquoted `off` as a boolean.
         if rust_event_persistence is False:

synapse/state/__init__.py

@@ -1013,6 +1013,12 @@ class StateResolutionHandler:
         ).inc()
         return state
 
+    def record_rust_state_res_metrics(
+        self, room_id: str, rust_stats: StateResStats
+    ) -> None:
+        """Record the work of a resolution the Rust event persister ran."""
+        self._record_state_res_metrics(room_id, ContextResourceUsage(), rust_stats)
+
     def _record_state_res_metrics(
         self,
         room_id: str,

synapse/storage/controllers/persist_events.py

@@ -36,6 +36,7 @@ from typing import (
     Iterable,
     Sequence,
     TypeVar,
+    cast,
 )
 
 import attr
@@ -70,11 +71,17 @@ from synapse.storage.controllers.state import StateStorageController
 from synapse.storage.databases import Databases
 from synapse.storage.databases.main.events import (
     DeltaState,
+    NewEventChainLinks,
+    PreparedTableChanges,
+    SlidingSyncMembershipInfo,
+    SlidingSyncMembershipSnapshotSharedInsertValues,
+    SlidingSyncStateInsertValues,
+    SlidingSyncTableChanges,
     event_counter,
     persist_event_counter,
 )
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
-from synapse.synapse_rust.persist import Persister
+from synapse.synapse_rust.persist import Persister, PreparedChunk
 from synapse.types import (
     PersistedEventPosition,
     RoomStreamToken,
@@ -154,6 +161,27 @@ times_pruned_extremities = Counter(
     labelnames=[SERVER_NAME_LABEL],
 )
 
+rust_prepare_counter = Counter(
+    "synapse_storage_events_rust_prepare",
+    "Chunks of events whose pre-transaction computation ran in Rust, by "
+    "`rust_event_persistence` level and result: `rust` (Rust's result was "
+    "used), `matched` or `mismatch` (shadow mode), `fallback` (Rust cannot "
+    "resolve state in the room) or `error`",
+    labelnames=["level", "result", SERVER_NAME_LABEL],
+)
+
+
+@attr.s(auto_attribs=True, slots=True, frozen=True)
+class _PreparedChunk:
+    """What `_persist_event_batch` computes for a chunk before persisting it."""
+
+    new_forward_extremities: set[str] | None
+    state_delta: DeltaState | None
+    new_state_dag_extremities: set[str] | None
+    new_event_links: dict[str, NewEventChainLinks]
+    # `None` to leave them to `_persist_events_and_state_updates`.
+    table_changes: PreparedTableChanges | None
+
 
 @attr.s(auto_attribs=True, slots=True)
 class _PersistEventsTask:
@@ -388,6 +416,153 @@ def flush_rust_persist_metrics(persister: Persister, server_name: str) -> None:
         )
 
 
+def _prepared_chunk_from_rust(room_id: str, rust: PreparedChunk) -> _PreparedChunk:
+    """Build the Python types from what `Persister.prepare_chunk` returned."""
+    state_delta = None
+    if rust.state_delta is not None:
+        to_delete, to_insert, no_longer_in_room = rust.state_delta
+        state_delta = DeltaState(
+            to_delete=to_delete,
+            to_insert=to_insert,
+            no_longer_in_room=no_longer_in_room,
+        )
+
+    sliding_sync_table_changes = None
+    if rust.sliding_sync_table_changes is not None:
+        bump_stamp, joined, shared, inserts, deletes = rust.sliding_sync_table_changes
+        sliding_sync_table_changes = SlidingSyncTableChanges(
+            room_id=room_id,
+            joined_room_bump_stamp_to_fully_insert=bump_stamp,
+            joined_room_updates=cast(SlidingSyncStateInsertValues, joined),
+            membership_snapshot_shared_insert_values=cast(
+                SlidingSyncMembershipSnapshotSharedInsertValues, shared
+            ),
+            to_insert_membership_snapshots=[
+                SlidingSyncMembershipInfo(
+                    user_id=user_id,
+                    sender=sender,
+                    membership_event_id=membership_event_id,
+                    membership=membership,
+                )
+                for user_id, sender, membership_event_id, membership in inserts
+            ],
+            to_delete_membership_snapshots=deletes,
+        )
+
+    return _PreparedChunk(
+        new_forward_extremities=rust.new_forward_extremities,
+        state_delta=state_delta,
+        new_state_dag_extremities=rust.new_state_dag_extremities,
+        new_event_links={
+            event_id: NewEventChainLinks(
+                chain_id=chain_id, sequence_number=sequence_number, links=links
+            )
+            for event_id, (
+                chain_id,
+                sequence_number,
+                links,
+            ) in rust.new_event_links.items()
+        },
+        table_changes=PreparedTableChanges(
+            sliding_sync_table_changes=sliding_sync_table_changes,
+            sticky_events_to_un_soft_fail=rust.sticky_events_to_un_soft_fail,
+        ),
+    )
+
+
+def _prepared_chunk_differences(
+    python: _PreparedChunk, rust: _PreparedChunk
+) -> list[str]:
+    """How Rust's preparation of a chunk differs from Python's, ignoring the
+    orders Python leaves to chance and the IDs of new chains."""
+    differences = []
+
+    def compare(name: str, expected: object, actual: object) -> None:
+        if expected != actual:
+            differences.append(f"{name}: Python {expected!r}, Rust {actual!r}")
+
+    compare(
+        "new_forward_extremities",
+        python.new_forward_extremities,
+        rust.new_forward_extremities,
+    )
+    compare(
+        "new_state_dag_extremities",
+        python.new_state_dag_extremities,
+        rust.new_state_dag_extremities,
+    )
+
+    def delta(
+        delta: DeltaState | None,
+    ) -> tuple[set[tuple[str, str]], dict[tuple[str, str], str], bool] | None:
+        if delta is None:
+            return None
+        return set(delta.to_delete), dict(delta.to_insert), delta.no_longer_in_room
+
+    compare("state_delta", delta(python.state_delta), delta(rust.state_delta))
+    compare(
+        "new_event_links",
+        _normalise_chain_links(python.new_event_links),
+        _normalise_chain_links(rust.new_event_links),
+    )
+
+    assert python.table_changes is not None
+    assert rust.table_changes is not None
+
+    def sliding_sync(changes: SlidingSyncTableChanges | None) -> object:
+        if changes is None:
+            return None
+        return (
+            changes.joined_room_bump_stamp_to_fully_insert,
+            dict(changes.joined_room_updates),
+            dict(changes.membership_snapshot_shared_insert_values),
+            sorted(
+                (i.user_id, i.sender, i.membership_event_id, i.membership)
+                for i in changes.to_insert_membership_snapshots
+            ),
+            sorted(changes.to_delete_membership_snapshots),
+        )
+
+    compare(
+        "sliding_sync_table_changes",
+        sliding_sync(python.table_changes.sliding_sync_table_changes),
+        sliding_sync(rust.table_changes.sliding_sync_table_changes),
+    )
+    compare(
+        "sticky_events_to_un_soft_fail",
+        python.table_changes.sticky_events_to_un_soft_fail,
+        rust.table_changes.sticky_events_to_un_soft_fail,
+    )
+    return differences
+
+
+def _normalise_chain_links(
+    links: dict[str, NewEventChainLinks],
+) -> dict[str, tuple[str, int, list[tuple[str, int]]]]:
+    """The chain links with each new chain named after the lowest event ID
+    on it. Python and Rust take new chain IDs from the sequence separately,
+    so only the shape of the new chains can match. A new chain is one with
+    an event at sequence number 1. Events that extend an existing chain
+    come after its first event, so the IDs of existing chains must match."""
+    names: dict[int, str] = {}
+    new_chains = {link.chain_id for link in links.values() if link.sequence_number == 1}
+    for event_id, link in sorted(links.items()):
+        if link.chain_id in new_chains:
+            names.setdefault(link.chain_id, f"new {event_id}")
+
+    def name(chain_id: int) -> str:
+        return names.get(chain_id, str(chain_id))
+
+    return {
+        event_id: (
+            name(link.chain_id),
+            link.sequence_number,
+            sorted((name(chain_id), seq) for chain_id, seq in link.links),
+        )
+        for event_id, link in links.items()
+    }
+
+
 class EventsPersistenceStorageController:
     """High level interface for handling persisting newly received events.
 
@@ -426,6 +601,9 @@ class EventsPersistenceStorageController:
         # "off". It needs the native Rust Postgres driver.
         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
+        # implementations, raises rather than being logged. Tests set it.
+        self._rust_persist_strict = False
         if rust_level != "off":
             self._rust_persister = hs.get_rust_handlers().persister
             if self._rust_persister is None:
@@ -720,55 +898,7 @@ class EventsPersistenceStorageController:
         for chunk in chunks:
             # We can't easily parallelize these since different chunks
             # might contain the same event. :(
-
-            new_forward_extremities = None
-            state_delta_for_room = None
-            new_state_dag_extrems = None
-
-            if not backfilled:
-                if all_supports_msc4242_state_dag(chunk):
-                    with Measure(
-                        self._clock,
-                        name="_process_state_dag_forward_extremities_and_state_delta",
-                        server_name=self.server_name,
-                    ):
-                        (
-                            new_forward_extremities,  # for prev_events
-                            state_delta_for_room,  # for state groups
-                            new_state_dag_extrems,  # for prev_state_events
-                        ) = await self._process_state_dag_forward_extremities_and_state_delta(
-                            room_id, chunk
-                        )
-                else:
-                    with Measure(
-                        self._clock,
-                        name="_calculate_state_and_extrem",
-                        server_name=self.server_name,
-                    ):
-                        # Work out the new "current state" for the room.
-                        # We do this by working out what the new extremities are and then
-                        # calculating the state from that.
-                        (
-                            new_forward_extremities,
-                            state_delta_for_room,
-                        ) = await self._calculate_new_forward_extremities_and_state_delta(
-                            room_id, chunk
-                        )
-
-            with Measure(
-                self._clock,
-                name="calculate_chain_cover_index_for_events",
-                server_name=self.server_name,
-            ):
-                # We now calculate chain ID/sequence numbers for any state events we're
-                # persisting. We ignore out of band memberships as we're not in the room
-                # and won't have their auth chain (we'll fix it up later if we join the
-                # room).
-                #
-                # See: docs/auth_chain_difference_algorithm.md
-                new_event_links = await self.persist_events_store.calculate_chain_cover_index_for_events(
-                    room_id, [e for e, _ in chunk]
-                )
+            prepared = await self._prepare_chunk(room_id, chunk, backfilled)
 
             # Stop the state groups from being deleted while we're persisting
             # them.
@@ -778,16 +908,244 @@ class EventsPersistenceStorageController:
                 await self.persist_events_store._persist_events_and_state_updates(
                     room_id,
                     chunk,
-                    state_delta_for_room=state_delta_for_room,
-                    new_forward_extremities=new_forward_extremities,
+                    state_delta_for_room=prepared.state_delta,
+                    new_forward_extremities=prepared.new_forward_extremities,
                     use_negative_stream_ordering=backfilled,
                     inhibit_local_membership_updates=backfilled,
-                    new_event_links=new_event_links,
-                    new_state_dag_forward_extremities=new_state_dag_extrems,
+                    new_event_links=prepared.new_event_links,
+                    new_state_dag_forward_extremities=prepared.new_state_dag_extremities,
+                    table_changes=prepared.table_changes,
                 )
 
         return replaced_events
 
+    async def _prepare_chunk(
+        self, room_id: str, chunk: list[EventPersistencePair], backfilled: bool
+    ) -> _PreparedChunk:
+        """Compute what persisting the chunk needs, in Python or Rust as
+        `rust_event_persistence` says."""
+        if self._rust_persister is None:
+            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)
+
+        try:
+            prepared = await self._prepare_chunk_in_rust(
+                room_id, chunk, backfilled, shadow=False
+            )
+        except SynapseError:
+            # The chunk is rejected, as Python would reject it.
+            self._count_rust_prepare(level, "rust")
+            raise
+        except Exception:
+            self._count_rust_prepare(level, "error")
+            if self._rust_persist_strict:
+                raise
+            logger.exception(
+                "Failed to prepare events for %s in Rust: using Python", room_id
+            )
+            prepared = None
+        else:
+            self._count_rust_prepare(level, "rust" if prepared else "fallback")
+        if prepared is None:
+            return await self._prepare_chunk_in_python(
+                room_id, chunk, backfilled, table_changes=False
+            )
+        return prepared
+
+    async def _prepare_chunk_in_shadow_mode(
+        self, room_id: str, chunk: list[EventPersistencePair], backfilled: bool
+    ) -> _PreparedChunk:
+        """Compute the chunk in both Rust and Python, compare the two, and
+        return Python's, or raise Python's error.
+
+        Rust goes first: its chain cover step is a dry run, while Python's
+        writes the events it puts off, which Rust would then read as left
+        from an earlier chunk.
+        """
+        rust_error: Exception | None = None
+        try:
+            rust = await self._prepare_chunk_in_rust(
+                room_id, chunk, backfilled, shadow=True
+            )
+        except Exception as e:
+            rust = None
+            rust_error = e
+
+        try:
+            python = await self._prepare_chunk_in_python(
+                room_id, chunk, backfilled, table_changes=True
+            )
+        except Exception as python_error:
+            # Both rejecting the chunk the same way is a match.
+            if rust is None and rust_error is None:
+                self._count_rust_prepare("shadow", "fallback")
+            elif type(rust_error) is type(python_error):
+                self._count_rust_prepare("shadow", "matched")
+            else:
+                self._report_shadow_mismatch(
+                    room_id, [f"Python raised {python_error!r}, Rust {rust_error!r}"]
+                )
+            raise
+
+        if isinstance(rust_error, SynapseError):
+            self._report_shadow_mismatch(
+                room_id, [f"Rust raised {rust_error!r}, Python succeeded"]
+            )
+        elif rust_error is not None:
+            self._count_rust_prepare("shadow", "error")
+            if self._rust_persist_strict:
+                raise rust_error
+            logger.error(
+                "Failed to prepare events for %s in Rust in shadow mode",
+                room_id,
+                exc_info=(type(rust_error), rust_error, rust_error.__traceback__),
+            )
+        elif rust is None:
+            self._count_rust_prepare("shadow", "fallback")
+        else:
+            differences = _prepared_chunk_differences(python, rust)
+            if differences:
+                self._report_shadow_mismatch(room_id, differences)
+            else:
+                self._count_rust_prepare("shadow", "matched")
+        return python
+
+    def _report_shadow_mismatch(self, room_id: str, differences: list[str]) -> None:
+        """Count a shadow mode mismatch, and log it or, in strict mode, raise."""
+        self._count_rust_prepare("shadow", "mismatch")
+        message = "Rust prepared events for %s differently: %s" % (
+            room_id,
+            "; ".join(differences),
+        )
+        if self._rust_persist_strict:
+            raise AssertionError(message)
+        logger.warning("%s", message)
+
+    async def _prepare_chunk_in_rust(
+        self,
+        room_id: str,
+        chunk: list[EventPersistencePair],
+        backfilled: bool,
+        *,
+        shadow: bool,
+    ) -> _PreparedChunk | None:
+        """Compute the chunk in Rust. `None` if the Rust state resolution
+        cannot resolve state in the room, so Python should compute it.
+        Unless `shadow` is set, what it counted is recorded, as the caller
+        uses the result."""
+        assert self._rust_persister is not None
+        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
+                )
+            )
+        if rust.state_res_fallback is not None:
+            logger.info(
+                "Preparing events for %s in Python: %s",
+                room_id,
+                rust.state_res_fallback,
+            )
+            return None
+        prepared = _prepared_chunk_from_rust(room_id, rust)
+        if not shadow:
+            self._rust_persister.record_prepare_counts(rust)
+            if rust.state_res_stats is not None:
+                self._state_resolution_handler.record_rust_state_res_metrics(
+                    room_id, rust.state_res_stats
+                )
+        return prepared
+
+    async def _prepare_chunk_in_python(
+        self,
+        room_id: str,
+        chunk: list[EventPersistencePair],
+        backfilled: bool,
+        *,
+        table_changes: bool,
+    ) -> _PreparedChunk:
+        """Compute the chunk in Python. The sliding sync changes and sticky
+        events are left to `_persist_events_and_state_updates` unless
+        `table_changes` is set."""
+        new_forward_extremities = None
+        state_delta_for_room = None
+        new_state_dag_extrems = None
+
+        if not backfilled:
+            if all_supports_msc4242_state_dag(chunk):
+                with Measure(
+                    self._clock,
+                    name="_process_state_dag_forward_extremities_and_state_delta",
+                    server_name=self.server_name,
+                ):
+                    (
+                        new_forward_extremities,  # for prev_events
+                        state_delta_for_room,  # for state groups
+                        new_state_dag_extrems,  # for prev_state_events
+                    ) = await self._process_state_dag_forward_extremities_and_state_delta(
+                        room_id, chunk
+                    )
+            else:
+                with Measure(
+                    self._clock,
+                    name="_calculate_state_and_extrem",
+                    server_name=self.server_name,
+                ):
+                    # Work out the new "current state" for the room.
+                    # We do this by working out what the new extremities are and then
+                    # calculating the state from that.
+                    (
+                        new_forward_extremities,
+                        state_delta_for_room,
+                    ) = await self._calculate_new_forward_extremities_and_state_delta(
+                        room_id, chunk
+                    )
+
+        with Measure(
+            self._clock,
+            name="calculate_chain_cover_index_for_events",
+            server_name=self.server_name,
+        ):
+            # We now calculate chain ID/sequence numbers for any state events we're
+            # persisting. We ignore out of band memberships as we're not in the room
+            # and won't have their auth chain (we'll fix it up later if we join the
+            # room).
+            #
+            # See: docs/auth_chain_difference_algorithm.md
+            new_event_links = (
+                await self.persist_events_store.calculate_chain_cover_index_for_events(
+                    room_id, [e for e, _ in chunk]
+                )
+            )
+
+        prepared_table_changes = None
+        if table_changes:
+            prepared_table_changes = (
+                await self.persist_events_store.calculate_table_changes(
+                    room_id, chunk, state_delta_for_room
+                )
+            )
+
+        return _PreparedChunk(
+            new_forward_extremities=new_forward_extremities,
+            state_delta=state_delta_for_room,
+            new_state_dag_extremities=new_state_dag_extrems,
+            new_event_links=new_event_links,
+            table_changes=prepared_table_changes,
+        )
+
+    def _count_rust_prepare(self, level: str, result: str) -> None:
+        rust_prepare_counter.labels(
+            level=level, result=result, **{SERVER_NAME_LABEL: self.server_name}
+        ).inc()
+
     async def _calculate_new_forward_extremities_and_state_delta(
         self, room_id: str, ev_ctx_rm: list[EventPersistencePair]
     ) -> tuple[set[str] | None, DeltaState | None]:

synapse/storage/databases/main/events.py

@@ -232,6 +232,16 @@ class SlidingSyncTableChanges:
     to_delete_membership_snapshots: list[str]
 
 
+@attr.s(slots=True, frozen=True, auto_attribs=True)
+class PreparedTableChanges:
+    """The sliding sync table changes and sticky events to un-soft-fail of a
+    chunk, computed before `_persist_events_and_state_updates` is called.
+    """
+
+    sliding_sync_table_changes: "SlidingSyncTableChanges | None"
+    sticky_events_to_un_soft_fail: set[str]
+
+
 @attr.s(slots=True, auto_attribs=True)
 class NewEventChainLinks:
     """Information about new auth chain links that need to be added to the DB.
@@ -305,6 +315,45 @@ class PersistEventsStore:
         self._stream_id_gen: AbstractStreamIdGenerator = self.store._stream_id_gen
 
     @trace
+    async def calculate_table_changes(
+        self,
+        room_id: str,
+        events_and_contexts: list[EventPersistencePair],
+        state_delta_for_room: DeltaState | None,
+    ) -> PreparedTableChanges:
+        """The sliding sync table changes and the sticky events to
+        un-soft-fail for persisting the events with `state_delta_for_room`.
+        """
+        sliding_sync_table_changes = None
+        if state_delta_for_room is not None:
+            sliding_sync_table_changes = (
+                await self._calculate_sliding_sync_table_changes(
+                    room_id, events_and_contexts, state_delta_for_room
+                )
+            )
+
+        sticky_events_to_un_soft_fail: set[str] = set()
+        if self._msc4354_enabled and state_delta_for_room is not None:
+            # When we change the room's current state with `state_delta_for_room`,
+            # that might cause some previously soft-failed sticky events to now pass
+            # the state-dependent auth checks.
+            # In other words, the sticky events could have been valid if they had
+            # waited for these state changes.
+            # For that reason, we give sticky events a second chance.
+            # We compute them here and then un-soft-fail them atomically with the
+            # persistence of the events.
+            sticky_events_to_un_soft_fail = (
+                await self.store.compute_sticky_events_to_un_soft_fail(
+                    room_id,
+                    events_and_contexts,
+                    state_delta_for_room,
+                )
+            )
+        return PreparedTableChanges(
+            sliding_sync_table_changes=sliding_sync_table_changes,
+            sticky_events_to_un_soft_fail=sticky_events_to_un_soft_fail,
+        )
+
     async def _persist_events_and_state_updates(
         self,
         room_id: str,
@@ -316,6 +365,7 @@ class PersistEventsStore:
         use_negative_stream_ordering: bool = False,
         inhibit_local_membership_updates: bool = False,
         new_state_dag_forward_extremities: set[str] | None = None,
+        table_changes: PreparedTableChanges | None = None,
     ) -> None:
         """Persist a set of events alongside updates to the current state and
                 forward extremities tables.
@@ -338,6 +388,9 @@ class PersistEventsStore:
                 not affect the current local state.
             new_state_dag_forward_extremities: A set of event IDs that are the new forward
                 extremities for the state DAG for this room. MSC4242 only.
+            table_changes: The sliding sync table changes and sticky events
+                to un-soft-fail, if the caller computed them. Otherwise they
+                are computed here.
 
         Returns:
             Resolves when the events have been persisted
@@ -385,31 +438,12 @@ class PersistEventsStore:
                 event.internal_metadata.received_ts = self._clock.time_msec()
                 event.internal_metadata.instance_name = self._instance_name
 
-            sliding_sync_table_changes = None
-            if state_delta_for_room is not None:
-                sliding_sync_table_changes = (
-                    await self._calculate_sliding_sync_table_changes(
-                        room_id, events_and_contexts, state_delta_for_room
-                    )
-                )
-
-            sticky_events_to_un_soft_fail: set[str] = set()
-            if self._msc4354_enabled and state_delta_for_room is not None:
-                # When we change the room's current state with `state_delta_for_room`,
-                # that might cause some previously soft-failed sticky events to now pass
-                # the state-dependent auth checks.
-                # In other words, the sticky events could have been valid if they had
-                # waited for these state changes.
-                # For that reason, we give sticky events a second chance.
-                # We compute them here and then un-soft-fail them atomically with the
-                # persistence of the events.
-                sticky_events_to_un_soft_fail = (
-                    await self.store.compute_sticky_events_to_un_soft_fail(
-                        room_id,
-                        events_and_contexts,
-                        state_delta_for_room,
-                    )
+            if table_changes is None:
+                table_changes = await self.calculate_table_changes(
+                    room_id, events_and_contexts, state_delta_for_room
                 )
+            sliding_sync_table_changes = table_changes.sliding_sync_table_changes
+            sticky_events_to_un_soft_fail = table_changes.sticky_events_to_un_soft_fail
 
             await self.db_pool.runInteraction(
                 "persist_events",

synapse/synapse_rust/persist.pyi

@@ -13,9 +13,12 @@
 
 from typing import Any
 
+from twisted.internet.defer import Deferred
+
 from synapse.events import EventBase
 from synapse.events.snapshot import EventContext
 from synapse.synapse_rust.database.postgres import ConnectionPool
+from synapse.synapse_rust.state import StateResStats
 
 class Persister:
     """Persists events in Rust, on the events writer, over the native driver.
@@ -25,6 +28,68 @@ class Persister:
     def take_metrics(self) -> dict[str, Any]:
         """What persisting has counted since the last call. See
         `flush_rust_persist_metrics`."""
+    def prepare_chunk(
+        self,
+        room_id: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        backfilled: bool,
+        *,
+        shadow: bool,
+    ) -> Deferred[PreparedChunk]:
+        """What `_persist_event_batch` computes for a chunk of one room's
+        events before the write transaction. In shadow mode the chain cover
+        step's write is rolled back. Nothing is counted until
+        `record_prepare_counts`."""
+    def record_prepare_counts(self, chunk: PreparedChunk) -> None:
+        """Add what preparing `chunk` counted to the metrics, once."""
+    def prepare_chunk_blocking(
+        self,
+        room_id: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        backfilled: bool,
+        *,
+        shadow: bool,
+    ) -> PreparedChunk:
+        """As `prepare_chunk`, but blocks. For tests."""
+
+class PreparedChunk:
+    """What `Persister.prepare_chunk` computed. If `state_res_fallback` is
+    set, the rest is empty and Python should compute the chunk itself."""
+
+    @property
+    def state_res_fallback(self) -> str | None:
+        """Why the Rust state resolution could not resolve the room."""
+    @property
+    def new_forward_extremities(self) -> set[str] | None: ...
+    @property
+    def state_delta(
+        self,
+    ) -> tuple[list[tuple[str, str]], dict[tuple[str, str], str], bool] | None:
+        """`(to_delete, to_insert, no_longer_in_room)`."""
+    @property
+    def new_state_dag_extremities(self) -> set[str] | None: ...
+    @property
+    def state_res_stats(self) -> StateResStats | None: ...
+    @property
+    def new_event_links(self) -> dict[str, tuple[int, int, list[tuple[int, int]]]]:
+        """Event ID to `(chain_id, sequence_number, links)`."""
+    @property
+    def sliding_sync_table_changes(
+        self,
+    ) -> (
+        tuple[
+            int | None,
+            dict[str, Any],
+            dict[str, Any],
+            list[tuple[str, str, str, str]],
+            list[str],
+        ]
+        | None
+    ):
+        """As `PersistReaderForTest.calculate_sliding_sync_table_changes`
+        returns it. `None` when the state does not change."""
+    @property
+    def sticky_events_to_un_soft_fail(self) -> set[str]: ...
 
 class PersistEffects:
     """The Python cache updates a Rust persist transaction asks for once it
@@ -144,10 +209,10 @@ class PersistReaderForTest:
         target_sequence_number)` rows."""
     def max_sequence_numbers(self, chain_ids: list[int]) -> dict[int, int]: ...
     # The chain cover step (`rust/src/persist/chain_cover.rs`), in a write
-    # transaction of its own. Event ID to `(chain_id, sequence_number,
-    # links)`.
+    # transaction of its own, rolled back for a dry run. Event ID to
+    # `(chain_id, sequence_number, links)`.
     def calculate_chain_cover_index_for_events(
-        self, room_id: str, events: list[EventBase]
+        self, room_id: str, events: list[EventBase], *, dry_run: bool = False
     ) -> dict[str, tuple[int, int, list[tuple[int, int]]]]: ...
     # The extremity calculations (`rust/src/persist/extremities.rs`), each
     # returning its result and what it counted, in the shape

tests/server.py

@@ -130,6 +130,7 @@ from tests.utils import (
     POSTGRES_PORT,
     POSTGRES_USER,
     RUST_PERSIST_FOR_TESTS,
+    RUST_PERSIST_STRICT_FOR_TESTS,
     RUST_STATE_RES_FOR_TESTS,
     RUST_SYNC_FOR_TESTS,
     RUST_SYNC_STRICT_FOR_TESTS,
@@ -1569,6 +1570,7 @@ def setup_test_homeserver(
             persistence._rust_persister = cast(
                 Any, _BlockingRustCalls(persistence._rust_persister)
             )
+            persistence._rust_persist_strict = RUST_PERSIST_STRICT_FOR_TESTS
 
     return hs
 

tests/state/test_rust_handler.py

@@ -57,6 +57,12 @@ class RustStateResHandlerTestCase(unittest.HomeserverTestCase):
         self.user_id = self.register_user("user", "pass")
         self.tok = self.login("user", "pass")
         self.handler = hs.get_state_resolution_handler()
+        # These tests resolve state through the handler as persisting does in
+        # Python. At `rust_event_persistence: prepare` and above the Rust
+        # persister resolves it itself, so turn that off.
+        persistence = hs.get_storage_controllers().persistence
+        assert persistence is not None
+        persistence._rust_persister = None
 
     def _fork(self, room_id: str) -> tuple[str, str]:
         """Send two topics with the same prev events, so persisting the second

tests/storage/test_persist_reader.py

@@ -872,6 +872,29 @@ class ChainCoverTestCase(PersistReaderTestBase):
             {row[0] for row in self._to_calculate()}, {oob.event_id, after.event_id}
         )
 
+    def test_dry_run(self) -> None:
+        """A dry run returns what a real one does and writes nothing."""
+        carol = "@carol:test"
+        oob, _ = self._new(
+            EventTypes.Member, carol, [self.create, self.power_levels, "$unknown"]
+        )
+        name, _ = self._new(
+            EventTypes.Name, "", [self.create, self.power_levels, self.alice_member]
+        )
+        before = self._to_calculate()
+
+        dry = self.rust.calculate_chain_cover_index_for_events(
+            self.room_id, [oob, name], dry_run=True
+        )
+        self.assertEqual(self._to_calculate(), before)
+
+        actual = self.rust.calculate_chain_cover_index_for_events(
+            self.room_id, [oob, name]
+        )
+        self.assertEqual({row[0] for row in self._to_calculate()}, {oob.event_id})
+        self.assertEqual(set(dry), {name.event_id})
+        self.assertEqual(self._normalise(dry), self._normalise(actual))
+
     def test_auth_event_without_a_chain(self) -> None:
         self._put_off(self.bob_member)
         member, _ = pair = self._new(

tests/storage/test_rust_persist_prepare.py

@@ -0,0 +1,458 @@
+#
+# 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:
+# <https://www.gnu.org/licenses/agpl-3.0.html>.
+#
+
+"""How `_persist_event_batch` prepares chunks in Rust
+(`rust/src/persist/prepare.rs`) at the `shadow` and `prepare` 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
+"""
+
+from typing import Any
+from unittest import SkipTest
+from unittest.mock import Mock
+
+import attr
+
+from twisted.internet import defer
+from twisted.internet.testing import MemoryReactor
+
+from synapse.api.constants import EventTypes, StickyEventField
+from synapse.api.errors import SynapseError
+from synapse.api.room_versions import RoomVersions
+from synapse.config.experimental import RustEventPersistenceLevel
+from synapse.events.snapshot import EventPersistencePair
+from synapse.rest import admin
+from synapse.rest.client import login, room
+from synapse.server import HomeServer
+from synapse.storage.controllers.persist_events import (
+    _PreparedChunk,
+    rust_prepare_counter,
+)
+from synapse.storage.rust_pool import RustConnectionPool
+from synapse.types import create_requester
+from synapse.util.clock import Clock
+from synapse.util.duration import Duration
+
+from tests import unittest
+from tests.test_utils.event_injection import inject_event
+
+_RESULTS = ("rust", "matched", "mismatch", "fallback", "error")
+
+
+class _PrepareTestBase(unittest.HomeserverTestCase):
+    """A homeserver at `level`, whatever `SYNAPSE_TEST_RUST_PERSIST` says,
+    with sticky events on."""
+
+    servlets = [
+        admin.register_servlets,
+        room.register_servlets,
+        login.register_servlets,
+    ]
+
+    level: RustEventPersistenceLevel
+
+    def default_config(self) -> dict[str, Any]:
+        config = super().default_config()
+        experimental = config.setdefault("experimental_features", {})
+        experimental["rust_event_persistence"] = self.level
+        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")
+        # The level is read per chunk, so this overrides the environment.
+        hs.config.experimental.rust_event_persistence = self.level
+        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
+
+        self.alice = self.register_user("alice", "pass")
+        self.alice_tok = self.login("alice", "pass")
+        self.bob = self.register_user("bob", "pass")
+        self.bob_tok = self.login("bob", "pass")
+        self.carol = self.register_user("carol", "pass")
+        self.carol_tok = self.login("carol", "pass")
+        self.room_id = self.helper.create_room_as(self.alice, tok=self.alice_tok)
+        self.helper.join(self.room_id, self.bob, tok=self.bob_tok)
+        self._counts_before = self._counts()
+
+    def _counts(self) -> dict[str, float]:
+        return {
+            result: rust_prepare_counter.labels(
+                level=self.level, result=result, server_name=self.hs.hostname
+            )._value.get()
+            for result in _RESULTS
+        }
+
+    def _counted(self) -> dict[str, float]:
+        """What was counted since `prepare`, leaving out zeros."""
+        after = self._counts()
+        return {
+            result: after[result] - self._counts_before[result]
+            for result in _RESULTS
+            if after[result] != self._counts_before[result]
+        }
+
+    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
+        ).get("event_id")
+
+    def _power_levels(self, users: dict[str, int]) -> None:
+        # The creator has unlimited power in room version 12.
+        self.helper.send_state(
+            self.room_id,
+            EventTypes.PowerLevels,
+            {"users": users, "events_default": 0},
+            tok=self.alice_tok,
+        )
+
+    def _sticky(self, sender: str) -> str:
+        """A soft-failed sticky event, as federation would persist one."""
+        return self.get_success(
+            inject_event(
+                self.hs,
+                room_id=self.room_id,
+                sender=sender,
+                type=EventTypes.Message,
+                content={"msgtype": "m.text", "body": "sticky"},
+                internal_metadata={"soft_failed": True},
+                msc4354_sticky=StickyEventField(
+                    duration_ms=Duration(minutes=10).as_millis()
+                ),
+            )
+        ).event_id
+
+    def _soft_failed_sticky_events(self) -> set[str]:
+        return set(
+            self.get_success(
+                self.store.db_pool.runInteraction(
+                    "test",
+                    self.store._get_soft_failed_sticky_events_txn,
+                    self.room_id,
+                    senders=None,
+                )
+            )
+        )
+
+    def _new_event(
+        self, room_id: str, event_type: str, content: dict[str, Any], **kwargs: Any
+    ) -> EventPersistencePair:
+        """An event of Alice's, not yet persisted."""
+        event, unpersisted_context = self.get_success(
+            self.hs.get_event_creation_handler().create_event(
+                create_requester(self.alice),
+                {
+                    "type": event_type,
+                    "content": content,
+                    "room_id": room_id,
+                    "sender": self.alice,
+                    **kwargs,
+                },
+                prev_event_ids=self.get_success(
+                    self.store.get_prev_events_for_room(room_id)
+                ),
+            )
+        )
+        return event, self.get_success(unpersisted_context.persist(event))
+
+    def _fork(self, room_id: str) -> None:
+        """Send two topics with the same prev events, so persisting the second
+        resolves state across the fork."""
+        pairs = [
+            self._new_event(room_id, EventTypes.Topic, {"topic": topic}, state_key="")
+            for topic in ("A", "B")
+        ]
+        for pair in pairs:
+            self.get_success(
+                self.hs.get_event_creation_handler().handle_new_client_event(
+                    create_requester(self.alice), events_and_context=[pair]
+                )
+            )
+
+    def _room_lifecycle(self) -> int:
+        """Persist a room's worth of changes, one chunk each. Returns the
+        number of chunks."""
+        self._send("hello")
+        self.helper.send_state(
+            self.room_id, EventTypes.Name, {"name": "The room"}, tok=self.alice_tok
+        )
+        self.helper.join(self.room_id, self.carol, tok=self.carol_tok)
+        self._power_levels({self.bob: -1})
+        sticky = self._sticky(self.bob)
+        self.assertEqual(self._soft_failed_sticky_events(), {sticky})
+        # Bob may speak again, so his sticky event passes.
+        self._power_levels({self.bob: 0})
+        self.assertEqual(self._soft_failed_sticky_events(), set())
+        self.helper.leave(self.room_id, self.carol, tok=self.carol_tok)
+        self._send("bye")
+        return 8
+
+
+class ShadowTestCase(_PrepareTestBase):
+    level = "shadow"
+
+    def test_room_lifecycle(self) -> None:
+        chunks = self._room_lifecycle()
+        self.assertEqual(self._counted(), {"matched": chunks})
+
+    def _tamper(self, change: dict[str, Any]) -> None:
+        """Make Rust's results differ from Python's by `change`."""
+        prepare_in_rust = self.persistence._prepare_chunk_in_rust
+
+        async def tampered(
+            room_id: str,
+            chunk: list[EventPersistencePair],
+            backfilled: bool,
+            *,
+            shadow: bool,
+        ) -> _PreparedChunk | None:
+            prepared = await prepare_in_rust(room_id, chunk, backfilled, shadow=shadow)
+            assert prepared is not None
+            return attr.evolve(prepared, **change)
+
+        self.persistence._prepare_chunk_in_rust = tampered  # type: ignore[method-assign]
+
+    def test_mismatch_fails_strict_mode(self) -> None:
+        self._tamper({"new_forward_extremities": {"$bogus"}})
+        self._send("hello", expect_code=500)
+        self.assertEqual(self._counted(), {"mismatch": 1})
+
+    def test_mismatch_uses_python(self) -> None:
+        self.persistence._rust_persist_strict = False
+        self._tamper({"new_forward_extremities": {"$bogus"}})
+        with self.assertLogs(
+            "synapse.storage.controllers.persist_events", "WARNING"
+        ) as logs:
+            event_id = self._send("hello")
+        self.assertIn("new_forward_extremities", logs.output[0])
+        self.assertEqual(self._counted(), {"mismatch": 1})
+        self.assertEqual(
+            self.get_success(self.store.get_latest_event_ids_in_room(self.room_id)),
+            frozenset({event_id}),
+        )
+
+    def _shift_chains(self, new: bool) -> None:
+        """Move the new chains (`new`) or the extended existing ones in
+        Rust's results to other chain IDs."""
+        prepare_in_rust = self.persistence._prepare_chunk_in_rust
+
+        async def shifted(
+            room_id: str,
+            chunk: list[EventPersistencePair],
+            backfilled: bool,
+            *,
+            shadow: bool,
+        ) -> _PreparedChunk | None:
+            prepared = await prepare_in_rust(room_id, chunk, backfilled, shadow=shadow)
+            assert prepared is not None
+            for links in prepared.new_event_links.values():
+                if (links.sequence_number == 1) == new:
+                    links.chain_id += 1_000_000
+            return prepared
+
+        self.persistence._prepare_chunk_in_rust = shifted  # type: ignore[method-assign]
+
+    def test_new_chains_are_compared_by_shape(self) -> None:
+        # Rust's dry run and Python take new chain IDs from the sequence
+        # separately, so the comparison names new chains after their events.
+        self._shift_chains(new=True)
+        # Carol's join starts a chain.
+        self.helper.join(self.room_id, self.carol, tok=self.carol_tok)
+        self.assertEqual(self._counted(), {"matched": 1})
+
+    def test_existing_chains_are_compared_by_id(self) -> None:
+        self.helper.join(self.room_id, self.carol, tok=self.carol_tok)
+        self.persistence._rust_persist_strict = False
+        self._shift_chains(new=False)
+        # Carol's leave extends the chain of her join.
+        with self.assertLogs("synapse.storage.controllers.persist_events", "WARNING"):
+            self.helper.leave(self.room_id, self.carol, tok=self.carol_tok)
+        self.assertEqual(self._counted(), {"matched": 1, "mismatch": 1})
+
+    def test_fork(self) -> None:
+        self._fork(self.room_id)
+        self.assertEqual(self._counted(), {"matched": 2})
+
+    def test_backfill(self) -> None:
+        pair = self._new_event(
+            self.room_id, EventTypes.Message, {"msgtype": "m.text", "body": "old"}
+        )
+        self.get_success(self.persistence.persist_events([pair], backfilled=True))
+        self.assertEqual(self._counted(), {"matched": 1})
+
+    def test_rust_rejection_python_accepts(self) -> None:
+        self.persistence._rust_persist_strict = False
+        self.persistence._rust_persister = Mock(
+            prepare_chunk=Mock(return_value=defer.fail(SynapseError(500, "no")))
+        )
+        with self.assertLogs("synapse.storage.controllers.persist_events", "WARNING"):
+            self._send("hello")
+        self.assertEqual(self._counted(), {"mismatch": 1})
+
+    def test_rust_error_is_logged(self) -> None:
+        self.persistence._rust_persist_strict = False
+        self.persistence._rust_persister = Mock(
+            prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom")))
+        )
+        with self.assertLogs("synapse.storage.controllers.persist_events", "ERROR"):
+            self._send("hello")
+        self.assertEqual(self._counted(), {"error": 1})
+
+
+class PrepareTestCase(_PrepareTestBase):
+    level = "prepare"
+
+    def test_room_lifecycle(self) -> None:
+        chunks = self._room_lifecycle()
+        self.assertEqual(self._counted(), {"rust": chunks})
+
+        state = self.get_success(
+            self.store.get_partial_filtered_current_state_ids(self.room_id)
+        )
+        self.assertEqual(
+            {key for key in state if key[0] == EventTypes.Member},
+            {
+                (EventTypes.Member, self.alice),
+                (EventTypes.Member, self.bob),
+                (EventTypes.Member, self.carol),
+            },
+        )
+        membership = self.get_success(
+            self.store.get_event(state[(EventTypes.Member, self.carol)])
+        ).membership
+        self.assertEqual(membership, "leave")
+
+        # The sliding sync tables have the room's name and Carol's leave.
+        (name,) = self.get_success(
+            self.store.db_pool.simple_select_one(
+                "sliding_sync_joined_rooms",
+                {"room_id": self.room_id},
+                ["room_name"],
+            )
+        )
+        self.assertEqual(name, "The room")
+        (carol_membership,) = self.get_success(
+            self.store.db_pool.simple_select_one(
+                "sliding_sync_membership_snapshots",
+                {"room_id": self.room_id, "user_id": self.carol},
+                ["membership"],
+            )
+        )
+        self.assertEqual(carol_membership, "leave")
+
+        # Every state event is in the chain cover index.
+        missing = self.get_success(
+            self.store.db_pool.execute(
+                "test",
+                """SELECT e.event_id FROM events AS e
+                LEFT JOIN event_auth_chains AS c USING (event_id)
+                WHERE e.room_id = ? AND e.state_key IS NOT NULL
+                    AND c.event_id IS NULL""",
+                self.room_id,
+            )
+        )
+        self.assertEqual(missing, [])
+
+    def test_fork(self) -> None:
+        self._fork(self.room_id)
+        self.assertEqual(self._counted(), {"rust": 2})
+        # Rust resolved state, and its work was recorded.
+        handler = self.hs.get_state_resolution_handler()
+        self.assertIn(self.room_id, handler._state_res_metrics)
+        state = self.get_success(
+            self.store.get_partial_filtered_current_state_ids(self.room_id)
+        )
+        topic = self.get_success(self.store.get_event(state[(EventTypes.Topic, "")]))
+        self.assertIn(topic.content["topic"], ("A", "B"))
+
+    def test_state_res_v1_falls_back(self) -> None:
+        room_id = self.helper.create_room_as(
+            self.alice, room_version=RoomVersions.V1.identifier, tok=self.alice_tok
+        )
+        before = self._counted()
+        self._fork(room_id)
+        after = self._counted()
+        # The first topic needs no resolution. The second does.
+        self.assertEqual(after.get("rust", 0) - before.get("rust", 0), 1)
+        self.assertEqual(after.get("fallback"), 1)
+
+    def test_backfill(self) -> None:
+        pair = self._new_event(
+            self.room_id, EventTypes.Message, {"msgtype": "m.text", "body": "old"}
+        )
+        self.get_success(self.persistence.persist_events([pair], backfilled=True))
+        self.assertEqual(self._counted(), {"rust": 1})
+        stream_ordering = pair[0].internal_metadata.stream_ordering
+        assert stream_ordering is not None
+        self.assertLess(stream_ordering, 0)
+
+    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(
+            prepare_chunk=Mock(return_value=defer.fail(SynapseError(500, "no")))
+        )
+        in_python = Mock(side_effect=AssertionError("computed in Python"))
+        self.persistence._prepare_chunk_in_python = in_python  # type: ignore[method-assign]
+        self._send("hello", expect_code=500)
+        self.assertEqual(self._counted(), {"rust": 1})
+        in_python.assert_not_called()
+
+    def test_fallback(self) -> None:
+        self.persistence._rust_persister = Mock(
+            prepare_chunk=Mock(
+                return_value=defer.succeed(Mock(state_res_fallback="no index"))
+            )
+        )
+        event_id = self._send("hello")
+        self.assertEqual(self._counted(), {"fallback": 1})
+        self.assertEqual(
+            self.get_success(self.store.get_latest_event_ids_in_room(self.room_id)),
+            frozenset({event_id}),
+        )
+
+    def test_error_uses_python(self) -> None:
+        self.persistence._rust_persist_strict = False
+        self.persistence._rust_persister = Mock(
+            prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom")))
+        )
+        with self.assertLogs("synapse.storage.controllers.persist_events", "ERROR"):
+            event_id = self._send("hello")
+        self.assertEqual(self._counted(), {"error": 1})
+        self.assertEqual(
+            self.get_success(self.store.get_latest_event_ids_in_room(self.room_id)),
+            frozenset({event_id}),
+        )
+
+    def test_error_fails_strict_mode(self) -> None:
+        self.persistence._rust_persister = Mock(
+            prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom")))
+        )
+        self._send("hello", expect_code=500)
+        self.assertEqual(self._counted(), {"error": 1})

tests/utils.py

@@ -66,6 +66,10 @@ RUST_SYNC_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_SYNC")
 RUST_SYNC_STRICT_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_SYNC_STRICT", "1") != "0"
 # As above, for `experimental_features.rust_event_persistence`.
 RUST_PERSIST_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_PERSIST")
+# As `SYNAPSE_TEST_RUST_SYNC_STRICT`, for the persister.
+RUST_PERSIST_STRICT_FOR_TESTS = (
+    os.environ.get("SYNAPSE_TEST_RUST_PERSIST_STRICT", "1") != "0"
+)
 # At `rust_federation_sending: transactions`, check every transaction body Rust
 # builds against the one Python would build, and fail the test if they differ.
 RUST_FED_STRICT_FOR_TESTS = os.environ.get("SYNAPSE_TEST_RUST_FED_STRICT", "1") != "0"