Start the Rust persist transaction, with a twin-room test

The first commit of phase 2 of the event persistence plan
(`transaction` level). `rust/src/persist/txn/` will hold the port of
`_persist_events_txn`. This commit adds its driver: the share lock on
the room's `rooms` row, the cache updates Python queues for every event
up front, and dropping duplicate events. It returns the events left for
the steps still in Python.

`tests/storage/test_rust_persist_txn.py` builds the same DAG in two
rooms. Room A persists in Python. Room B prepares in Rust, runs the
Rust part of the transaction in its own committed transaction through
`Persister.persist_events_txn_for_test`, and then runs the Python steps
after it. The test compares the rows of every table the transaction
writes, after normalising room and event IDs, stream positions, state
groups and chain IDs, and compares the cache updates both rooms asked
for. Each later commit moves more steps from the Python half to the
Rust half, in transaction order, until the facade runs the whole
transaction. Mutating the driver (dropping the redaction
invalidation, the relation, the state key, or the outlier rule for
duplicates) fails the comparison.

Two open items from phase 1:

- Reads before the transaction now retry. Python runs each read in a
  `runInteraction` that retries a lost connection or a cancelled
  statement. `Persister.prepare_chunk` now prepares the whole chunk
  again on such an error, up to five attempts. That is safe because its
  only write, bumping the state groups it uses, can be repeated. The
  reader and the state resolution store keep the pool's error as the
  cause, so `is_retryable_error` can classify it. A `WriteError` does
  not count, as its write interaction has already retried, and
  `WriteError::from_pg` shares the classification.
- The Rust persister needs the state tables in the main database, as
  it reads and writes both in one transaction. `Persister::for_homeserver`
  now checks for that, and the persister is off when they are apart.

Also fixed in the dispatch: at `off`, `_prepare_chunk` ran Rust when a
persister existed. That only happens when a test changes the level at
runtime. The twin test does exactly that.

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

rust/src/persist/python.rs

@@ -27,13 +27,16 @@ use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyDict};
 
 use super::{
     context::PersistContext,
+    effects::PersistEffects,
     extremities::MissingPrevStateEvents,
     metrics::{PersistCounts, PersistMetrics},
     prepare::{self, PrepareInput, PrepareOptions, Prepared},
     reader::{self, PersistReader, PyDelta, PySlidingSync},
     state_delta::NativeStateResolver,
+    txn::{self, TxnEvent, TxnInput},
 };
 use crate::{
+    database::postgres::write::{is_retryable_error, run_write_interaction, MAX_ATTEMPTS},
     deferred::{create_deferred_with, OnCancel},
     errors::SynapseError,
     event_auth::AuthFailure,
@@ -86,7 +89,9 @@ pub(super) struct PersisterInner {
 impl Persister {
     /// Build the persister for `homeserver`, whose main database uses the
     /// native driver through `connection_pool` (a `RustConnectionPool`).
-    /// `None` if this instance does not write events.
+    /// `None` if this instance does not write events, or if the state tables
+    /// are in another database: persisting reads and writes them in the
+    /// main database's transactions.
     pub fn for_homeserver(
         py: Python<'_>,
         homeserver: &HomeServer,
@@ -100,6 +105,16 @@ impl Persister {
             return Ok(None);
         }
 
+        if !homeserver
+            .state_database_pool(py)?
+            .is(&homeserver.main_database_pool(py)?)
+        {
+            log::warn!(
+                "The state tables are not in the main database: persisting events in Python"
+            );
+            return Ok(None);
+        }
+
         let config = homeserver.config(py)?;
         let main = homeserver.main_datastore(py)?;
         let tracker = |attr: &Bound<'_, pyo3::types::PyString>| -> PyResult<Arc<StreamIdTracker>> {
@@ -165,15 +180,37 @@ impl Persister {
                 sticky: inner.msc4354_enabled,
                 dry_run: shadow,
             };
-            let mut counts = PersistCounts::default();
-            let prepared =
-                prepare::prepare_chunk(&reader, &resolver, &pool, &input, &options, &mut counts)
-                    .await
-                    .map_err(persist_error_to_py)?;
-            Ok(PreparedChunk {
-                prepared,
-                counts: Mutex::new(Some(counts)),
-            })
+            // Python reads in `runInteraction`s, which retry on a lost
+            // connection or a cancelled statement. Here the whole chunk is
+            // prepared again. That is safe as its only write, bumping the
+            // state groups it uses, can be repeated, and runs in a write
+            // interaction with its own retries.
+            let mut attempt = 1;
+            loop {
+                let mut counts = PersistCounts::default();
+                match prepare::prepare_chunk(
+                    &reader,
+                    &resolver,
+                    &pool,
+                    &input,
+                    &options,
+                    &mut counts,
+                )
+                .await
+                {
+                    Ok(prepared) => {
+                        return Ok(PreparedChunk {
+                            prepared,
+                            counts: Mutex::new(Some(counts)),
+                        })
+                    }
+                    Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => {
+                        log::warn!("[TXN RETRY] {{prepare_chunk}} {e:#} {attempt}/{MAX_ATTEMPTS}");
+                        attempt += 1;
+                    }
+                    Err(e) => return Err(persist_error_to_py(e)),
+                }
+            }
         }
         .boxed())
     }
@@ -231,6 +268,47 @@ impl Persister {
         }
     }
 
+    /// Run the part of the write transaction ported so far over the chunk,
+    /// whose events must have their stream orderings set, and commit it.
+    /// Returns what it collected for Python to apply and the positions in
+    /// the chunk of the events left for the Python steps after it. For
+    /// `tests/storage/test_rust_persist_txn.py`.
+    fn persist_events_txn_for_test(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+    ) -> PyResult<(PersistEffects, Vec<usize>)> {
+        let events = events_and_contexts
+            .into_iter()
+            .map(|(event, context)| TxnEvent::new(event.bind(py).get().shallow_clone(), context))
+            .collect::<anyhow::Result<Vec<_>>>()
+            .map_err(persist_error_to_py)?;
+        let pool = self.inner.pool.get()?;
+        let future = async move {
+            let input = TxnInput {
+                room_id: &room_id,
+                events: &events,
+            };
+            run_write_interaction(&pool, "persist_events", |txn, _| {
+                let input = &input;
+                async move {
+                    let mut effects = PersistEffects::default();
+                    let left = txn::persist_events_txn(txn, input, &mut effects).await?;
+                    let left = left
+                        .iter()
+                        .filter_map(|e| input.events.iter().position(|x| std::ptr::eq(x, *e)))
+                        .collect();
+                    Ok((effects, left))
+                }
+                .boxed()
+            })
+            .await
+            .map_err(|e| persist_error_to_py(anyhow::Error::new(e)))
+        };
+        self.block_on(py, future.boxed())
+    }
+
     /// As [`Self::prepare_chunk`], but blocks.
     #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))]
     fn prepare_chunk_blocking(