Persist each queued batch in Rust at rust_event_persistence: full

At `full`, `_persist_event_batch` hands the batch to
`Persister.persist_batch` and steps through it with `batch_next`:

- The first step drops the events already persisted under the same
  transaction ID (`get_already_persisted_events`).
- Each step prepares the next chunk of up to 100 events, marks the
  batch's state groups as being persisted
  (`persisting_state_group_references`, two short native
  transactions), takes the chunk's stream orderings and writes it.
- Python applies what the chunk asks for, runs
  `_after_persisting_events` and finishes the orderings before the next
  step, which first releases the state groups. So each chunk is
  prepared against current caches.
- If Python cannot apply a chunk, `batch_abort` releases the state
  groups as a failed persist does.

A chunk Rust cannot prepare (a state resolution fallback, or an error)
is handed back with the rest of the batch to the chunk-by-chunk path,
which prepares it again and falls back or raises as at lower levels.

Python drives the steps, rather than Rust calling back into Python, so
every step has a `_blocking` twin for tests and nothing waits on the
reactor from Rust. That costs the same hops as a callback: one in and
one out per chunk.

The prepare inputs are now built from `Event`s that share their data
with the Python events, so a chunk's input can be built without the
GIL once the batch has been deduplicated.

The persister's Deferreds already follow the logcontext rules, so they
are awaited directly. Wrapping them in `make_deferred_yieldable` again,
as the `prepare_chunk`, `persist_events_txn` and `persist_chunk` call
sites did, resumed the caller in the sentinel context.

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

rust/src/persist/batch.rs

@@ -0,0 +1,245 @@
+/*
+ * This file is licensed under the Affero General Public License (AGPL) version 3.
+ *
+ * Copyright (C) 2026 Element Creations Ltd
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * See the GNU Affero General Public License for more details:
+ * <https://www.gnu.org/licenses/agpl-3.0.html>.
+ *
+ */
+
+//! A queued batch of one room's events, as `_persist_event_batch` persists
+//! it: the events already persisted under the same transaction ID are
+//! dropped, the rest go in chunks of [`CHUNK_SIZE`], and each chunk is
+//! written while its state groups are held back from deletion.
+//!
+//! The loop itself is [`super::python::Persister::persist_batch`], which
+//! hands each written chunk to Python. This holds its steps.
+
+use std::collections::{BTreeSet, HashMap, HashSet};
+
+use anyhow::anyhow;
+use futures::FutureExt;
+
+use super::context::PersistContext;
+use crate::{
+    clock::Clock,
+    database::postgres::{
+        pool::ConnectionPool,
+        write::{run_write_interaction, Column, WriteError},
+    },
+    events::Event,
+    storage::events::Querier,
+};
+
+/// How many events `_persist_event_batch` writes in one transaction.
+pub const CHUNK_SIZE: usize = 100;
+
+/// `get_already_persisted_events`: for each event sent with a transaction
+/// ID from a device, the event already persisted with the same room,
+/// sender, device and transaction ID, if there is one. An event that
+/// repeats an earlier one of the batch maps to that one.
+pub async fn already_persisted_events(
+    querier: &impl Querier,
+    events: &[(Event, PersistContext)],
+) -> anyhow::Result<HashMap<String, String>> {
+    let mut mapping = HashMap::new();
+    let mut by_txn_id: HashMap<(String, String, String, String), String> = HashMap::new();
+    for (event, _) in events {
+        let metadata = event.internal_metadata();
+        let (Some(device_id), Some(txn_id)) = (metadata.device_id()?, metadata.txn_id()?) else {
+            continue;
+        };
+        // Python skips an empty device or transaction ID.
+        if device_id.is_empty() || txn_id.is_empty() {
+            continue;
+        }
+        let key = (
+            event.room_id().to_owned(),
+            event.sender().to_owned(),
+            device_id,
+            txn_id,
+        );
+        if let Some(existing) = by_txn_id.get(&key) {
+            mapping.insert(event.event_id().to_owned(), existing.clone());
+            continue;
+        }
+        let rows = querier
+            .query(
+                "SELECT event_id FROM event_txn_id_device_id
+                WHERE room_id = $1 AND user_id = $2 AND device_id = $3 AND txn_id = $4",
+                &[&key.0, &key.1, &key.2, &key.3],
+            )
+            .await?;
+        let existing: Option<String> = rows.first().map(|row| row.try_get(0)).transpose()?;
+        match existing {
+            Some(existing) => {
+                mapping.insert(event.event_id().to_owned(), existing.clone());
+                by_txn_id.insert(key, existing);
+            }
+            None => {
+                by_txn_id.insert(key, event.event_id().to_owned());
+            }
+        }
+    }
+    Ok(mapping)
+}
+
+/// The state groups `persisting_state_group_references` holds back from
+/// deletion for `events`: those of each event that is neither rejected nor
+/// an outlier, and the groups before them.
+pub fn referenced_state_groups(
+    events: &[(Event, PersistContext)],
+) -> anyhow::Result<BTreeSet<i64>> {
+    let mut groups = BTreeSet::new();
+    for (event, context) in events {
+        if context.is_rejected() || event.internal_metadata().is_outlier()? {
+            continue;
+        }
+        let group = context
+            .state_group
+            .ok_or_else(|| anyhow!("{} has no state group", event.event_id()))?;
+        groups.insert(group);
+        // Python skips a falsy group before the event.
+        if let Some(before) = context.state_group_before_event.filter(|g| *g != 0) {
+            groups.insert(before);
+        }
+    }
+    Ok(groups)
+}
+
+/// The state groups were deleted before they could be marked.
+#[derive(Debug)]
+pub struct StateGroupsDeleted(pub Vec<i64>);
+
+impl std::fmt::Display for StateGroupsDeleted {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        // As Python's `shortstr`.
+        let shown: Vec<String> = self.0.iter().take(5).map(i64::to_string).collect();
+        let more = if self.0.len() > 5 { ", ..." } else { "" };
+        write!(
+            f,
+            "state groups have been deleted: [{}{more}]",
+            shown.join(", ")
+        )
+    }
+}
+
+impl std::error::Error for StateGroupsDeleted {}
+
+/// `_mark_state_groups_as_persisting_txn`: check the groups still exist,
+/// locking them against a concurrent delete, and record that this instance
+/// is persisting them.
+pub async fn mark_state_groups_as_persisting(
+    pool: &ConnectionPool,
+    clock: &Clock,
+    instance_name: &str,
+    groups: &BTreeSet<i64>,
+) -> Result<(), WriteError> {
+    let groups: Vec<i64> = groups.iter().copied().collect();
+    run_write_interaction(pool, "mark_state_groups_as_persisting", |txn, _| {
+        let groups = &groups;
+        async move {
+            let rows = txn
+                .query(
+                    "SELECT id FROM state_groups WHERE id = ANY($1) FOR KEY SHARE",
+                    &[groups],
+                )
+                .await?;
+            let existing = rows
+                .iter()
+                .map(|row| row.try_get(0))
+                .collect::<Result<HashSet<i64>, _>>()
+                .map_err(anyhow::Error::new)?;
+            let missing: Vec<i64> = groups
+                .iter()
+                .copied()
+                .filter(|g| !existing.contains(g))
+                .collect();
+            if !missing.is_empty() {
+                return Err(anyhow::Error::new(StateGroupsDeleted(missing)).into());
+            }
+            let now = i64::try_from(clock.now_millis()).map_err(anyhow::Error::new)?;
+            let instance_names = vec![instance_name; groups.len()];
+            let inserted = vec![now; groups.len()];
+            txn.upsert_many(
+                "state_groups_persisting",
+                &[
+                    Column::new("state_group", "bigint", groups),
+                    Column::new("instance_name", "text", &instance_names),
+                ],
+                &[Column::new("inserted_ts", "bigint", &inserted)],
+            )
+            .await?;
+            Ok(())
+        }
+        .boxed()
+    })
+    .await
+}
+
+/// `_finish_persisting_txn`: this instance no longer persists the groups.
+/// If the persist failed, any pending deletion of them is pushed back, as
+/// they may now be referenced. Otherwise it is dropped, as they are.
+pub async fn finish_persisting(
+    pool: &ConnectionPool,
+    clock: &Clock,
+    instance_name: &str,
+    groups: &BTreeSet<i64>,
+    error: bool,
+) -> Result<(), WriteError> {
+    let groups: Vec<i64> = groups.iter().copied().collect();
+    run_write_interaction(pool, "finish_persisting", |txn, _| {
+        let groups = &groups;
+        async move {
+            txn.execute(
+                "DELETE FROM state_groups_persisting
+                WHERE instance_name = $1 AND state_group = ANY($2)",
+                &[&instance_name, groups],
+            )
+            .await?;
+            if error {
+                // `_bump_deletion_txn`.
+                let now = i64::try_from(clock.now_millis()).map_err(anyhow::Error::new)?;
+                txn.execute(
+                    "UPDATE state_groups_pending_deletion
+                    SET sequence_number = DEFAULT, insertion_ts = $1
+                    WHERE state_group = ANY($2)",
+                    &[&now, groups],
+                )
+                .await?;
+            } else {
+                txn.execute(
+                    "DELETE FROM state_groups_pending_deletion WHERE state_group = ANY($1)",
+                    &[groups],
+                )
+                .await?;
+            }
+            Ok(())
+        }
+        .boxed()
+    })
+    .await
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn deleted_groups_are_shown_as_python_shows_them() {
+        assert_eq!(
+            StateGroupsDeleted(vec![1, 2]).to_string(),
+            "state groups have been deleted: [1, 2]"
+        );
+        assert_eq!(
+            StateGroupsDeleted((1..=6).collect()).to_string(),
+            "state groups have been deleted: [1, 2, 3, 4, 5, ...]"
+        );
+    }
+}

rust/src/persist/mod.rs

@@ -18,6 +18,7 @@
 //! This is being built up from the inside out (see
 //! `renders/event-persistence-rust-plan.md`). So far:
 //!
+//! - [`batch`]: the steps around a queued batch's chunks.
 //! - [`chain_cover`]: the chain cover index of a chunk's state events.
 //! - [`context`]: what persisting needs from an `EventContext`.
 //! - [`effects`]: the Python cache updates a Rust transaction asks for once
@@ -38,6 +39,7 @@
 
 use pyo3::prelude::*;
 
+pub mod batch;
 pub mod chain_cover;
 pub mod context;
 pub mod effects;
@@ -56,6 +58,9 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     let child_module = PyModule::new(py, "persist")?;
     child_module.add_class::<python::Persister>()?;
     child_module.add_class::<python::PreparedChunk>()?;
+    child_module.add_class::<python::PersistBatch>()?;
+    child_module.add_class::<python::PersistedBatch>()?;
+    child_module.add_class::<python::WrittenChunk>()?;
     child_module.add_class::<effects::PersistEffects>()?;
     child_module.add_class::<reader::PyPersistReaderForTest>()?;
     child_module.add_function(wrap_pyfunction!(

rust/src/persist/prepare.rs

@@ -40,8 +40,6 @@ use std::{
     sync::Arc,
 };
 
-use pyo3::prelude::*;
-
 use super::{
     chain_cover::{self, ChainCoverEvent, NewEventChainLinks},
     context::PersistContext,
@@ -59,8 +57,7 @@ use crate::{
     state::EventMap,
 };
 
-/// A chunk of one room's events, in the shapes the steps read. Built while
-/// the GIL is held, so that the steps run without it.
+/// A chunk of one room's events, in the shapes the steps read.
 pub struct PrepareInput {
     pub room_id: String,
     pub backfilled: bool,
@@ -74,17 +71,16 @@ pub struct PrepareInput {
 }
 
 impl PrepareInput {
+    /// The chunk `events_and_contexts`, which share their data with the
+    /// Python events they were taken from.
     pub fn new(
-        py: Python<'_>,
         room_id: String,
-        events_and_contexts: &[(Py<Event>, PersistContext)],
+        events_and_contexts: &[(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();
+    ) -> anyhow::Result<Self> {
+        let (events, event_map) = reader::chunk(events_and_contexts)?;
+        let chunk_events: Vec<&Event> =
+            events_and_contexts.iter().map(|(event, _)| event).collect();
         Ok(PrepareInput {
             room_id,
             backfilled,

rust/src/persist/python.rs

@@ -17,7 +17,7 @@
 //! `RustHandlers.persister` hands out, and helpers for tests.
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{BTreeSet, HashMap, HashSet},
     sync::{Arc, Mutex},
 };
 
@@ -26,6 +26,7 @@ use http::StatusCode;
 use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyDict};
 
 use super::{
+    batch,
     context::PersistContext,
     effects::PersistEffects,
     extremities::MissingPrevStateEvents,
@@ -186,83 +187,25 @@ impl Persister {
         if !matches!(&*prepared, Prepared::Chunk(_)) {
             return Err(PyRuntimeError::new_err("the chunk was prepared in Python"));
         }
-        let events: Vec<(Event, PersistContext)> = events_and_contexts
-            .into_iter()
-            .map(|(event, context)| (event.bind(py).get().shallow_clone(), context))
-            .collect();
+        let events = shallow_events(py, events_and_contexts);
         let pool = self.inner.pool.get()?;
         let inner = Arc::clone(&self.inner);
-        let future = async move {
+        Ok(async move {
             let Prepared::Chunk(chunk) = &*prepared else {
                 unreachable!("checked above");
             };
-            // Taken as late as possible, as Python takes them: until the
-            // chunk is written, no later event on the stream can be
-            // announced.
-            let event_stream = match orderings {
-                StreamOrderings::Set => None,
-                StreamOrderings::Allocate { backfilled } => {
-                    let stream = if backfilled {
-                        TxnStream::backfill(Arc::clone(&inner.backfill_stream))
-                    } else {
-                        TxnStream::events(Arc::clone(&inner.events_stream))
-                    };
-                    let ids = allocate_stream_orderings(&pool, &stream, events.len())
-                        .await
-                        .map_err(persist_error_to_py)?;
-                    set_persisted_positions(&inner, &events, &ids, backfilled)
-                        .map_err(persist_error_to_py)?;
-                    Some((stream, ids, backfilled))
-                }
-            };
-            let events = events
-                .into_iter()
-                .map(|(event, context)| TxnEvent::new(event, context))
-                .collect::<anyhow::Result<Vec<_>>>()
-                .map_err(persist_error_to_py)?;
-
-            let cache_stream = TxnStream::caches(Arc::clone(&inner.cache_stream));
-            let sticky_stream = TxnStream::sticky_events(Arc::clone(&inner.sticky_stream));
-            let profile_stream = TxnStream::profile_updates(Arc::clone(&inner.profile_stream));
-            let input = TxnInput {
-                room_id: &room_id,
-                events: &events,
-                prepared: chunk,
-                instance_name: &inner.instance_name,
-                server_name: &inner.server_name,
-                clock: inner.runtime.clock(),
+            write_chunk(
+                &inner,
+                &pool,
+                &room_id,
+                events,
+                chunk,
+                orderings,
                 inhibit_local_membership_updates,
-                enable_search: inner.enable_search,
-                enable_ephemeral_messages: inner.enable_ephemeral_messages,
-                msc4354_enabled: inner.msc4354_enabled,
-                include_profile_updates_in_sync: inner.include_profile_updates_in_sync,
-                cache_stream: &cache_stream,
-                sticky_stream: &sticky_stream,
-                profile_stream: &profile_stream,
-            };
-            let mut effects = run_write_interaction(&pool, "persist_events", |txn, _| {
-                let input = &input;
-                async move {
-                    let mut effects = PersistEffects::default();
-                    txn::persist_events_txn(txn, input, &mut effects).await?;
-                    Ok(effects)
-                }
-                .boxed()
-            })
+            )
             .await
-            .map_err(write_error_to_py)?;
-            cache_stream.hand_over(effects.cache_stream_ids());
-            sticky_stream.hand_over(effects.sticky_stream_ids());
-            profile_stream.hand_over(effects.profile_stream_ids());
-            // Finished on drop if the transaction failed, as `get_next_mult`
-            // finishes them when its block raises.
-            if let Some((stream, ids, backfilled)) = event_stream {
-                stream.hand_over(&ids);
-                effects.add_stream_ids(ids, backfilled);
-            }
-            Ok(effects)
-        };
-        Ok(future.boxed())
+        }
+        .boxed())
     }
 
     /// Convert the chunk while the GIL is held, and return the future that
@@ -275,59 +218,39 @@ impl Persister {
         backfilled: bool,
         shadow: bool,
     ) -> PyResult<BoxFuture<'static, PyResult<PreparedChunk>>> {
-        let input = PrepareInput::new(py, room_id, &events_and_contexts, backfilled)?;
+        let events = shallow_events(py, events_and_contexts);
+        let input = PrepareInput::new(room_id, &events, backfilled).map_err(persist_error_to_py)?;
         let pool = self.inner.pool.get()?;
-        let handle = self.inner.runtime.tokio_handle()?;
         let inner = Arc::clone(&self.inner);
         Ok(async move {
-            let resolver = NativeStateResolver {
-                pool: pool.clone(),
-                handle,
-                caches: Arc::clone(&inner.store.caches),
-                conflict_cache: Arc::clone(&inner.conflict_cache),
-            };
-            let reader = PersistReader::new(inner.sync.reader(pool.clone()));
-            let options = PrepareOptions {
-                server_name: &inner.server_name,
-                clock: inner.runtime.clock(),
-                sticky: inner.msc4354_enabled,
-                dry_run: shadow,
-            };
-            // Python reads in `runInteraction`s, which retry on a lost
-            // connection or a cancelled statement. Here the whole chunk is
-            // prepared again. That is safe as its only write, bumping the
-            // state groups it uses, can be repeated, and runs in a write
-            // interaction with its own retries.
-            let mut attempt = 1;
-            loop {
-                let mut counts = PersistCounts::default();
-                match prepare::prepare_chunk(
-                    &reader,
-                    &resolver,
-                    &pool,
-                    &input,
-                    &options,
-                    &mut counts,
-                )
+            let (prepared, counts) = prepare_chunk(&inner, &pool, &input, shadow)
                 .await
-                {
-                    Ok(prepared) => {
-                        return Ok(PreparedChunk {
-                            prepared: Arc::new(prepared),
-                            counts: Mutex::new(Some(counts)),
-                        })
-                    }
-                    Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => {
-                        log::warn!("[TXN RETRY] {{prepare_chunk}} {e:#} {attempt}/{MAX_ATTEMPTS}");
-                        attempt += 1;
-                    }
-                    Err(e) => return Err(persist_error_to_py(e)),
-                }
-            }
+                .map_err(persist_error_to_py)?;
+            Ok(PreparedChunk {
+                prepared: Arc::new(prepared),
+                counts: Mutex::new(Some(counts)),
+            })
         }
         .boxed())
     }
 
+    /// The future of the batch's next step. See [`Self::batch_next`].
+    fn start_batch_next(&self, batch: &PersistBatch) -> BoxFuture<'static, PyResult<BatchStep>> {
+        let state = Arc::clone(&batch.state);
+        async move { batch_next(&mut *state.lock().await).await }.boxed()
+    }
+
+    /// The future that releases the batch. See [`Self::batch_abort`].
+    fn start_batch_abort(&self, batch: &PersistBatch) -> BoxFuture<'static, PyResult<()>> {
+        let state = Arc::clone(&batch.state);
+        async move {
+            let mut state = state.lock().await;
+            state.done = true;
+            state.release_guard(true).await
+        }
+        .boxed()
+    }
+
     /// Run `future` to completion on the runtime, releasing the GIL. For the
     /// `_blocking` entry points.
     fn block_on<T: Send>(
@@ -481,6 +404,74 @@ impl Persister {
         self.block_on(py, future)
     }
 
+    /// `_persist_event_batch` for a queued batch of `room_id`'s events,
+    /// one chunk at a time: Python awaits [`Self::batch_next`] until it
+    /// resolves with a [`PersistedBatch`], applying the effects of each
+    /// [`WrittenChunk`] before the next step, and calls
+    /// [`Self::batch_abort`] if it cannot. Converts the events, without
+    /// I/O.
+    #[pyo3(signature = (room_id, events_and_contexts, *, backfilled))]
+    fn persist_batch(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+        events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+        backfilled: bool,
+    ) -> PyResult<PersistBatch> {
+        let state = BatchState {
+            inner: Arc::clone(&self.inner),
+            pool: self.inner.pool.get()?,
+            room_id,
+            backfilled,
+            events: shallow_events(py, events_and_contexts),
+            deduplicated: None,
+            next: 0,
+            guard_held: false,
+            done: false,
+        };
+        Ok(PersistBatch {
+            state: Arc::new(tokio::sync::Mutex::new(state)),
+        })
+    }
+
+    /// The batch's next step: a Deferred resolving to the [`WrittenChunk`]
+    /// written, or to the [`PersistedBatch`] once the batch is done, or
+    /// failing as the Python path would. The chunk's stream IDs are handed
+    /// to Python unfinished, so Python must not cancel the Deferred: if it
+    /// is cancelled, the step runs on and the IDs are lost. If the step
+    /// fails, Rust finishes them and releases the state groups.
+    fn batch_next<'py>(
+        &self,
+        py: Python<'py>,
+        batch: &PersistBatch,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let future = self.start_batch_next(batch);
+        create_deferred_with(py, &self.inner.runtime, future, OnCancel::RunOn)
+    }
+
+    /// As [`Self::batch_next`], but blocks.
+    fn batch_next_blocking(&self, py: Python<'_>, batch: &PersistBatch) -> PyResult<BatchStep> {
+        let future = self.start_batch_next(batch);
+        self.block_on(py, future)
+    }
+
+    /// Give up the batch, as Python could not apply a chunk's effects:
+    /// release its state groups as a failed persist does.
+    fn batch_abort<'py>(
+        &self,
+        py: Python<'py>,
+        batch: &PersistBatch,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let future = self.start_batch_abort(batch);
+        create_deferred_with(py, &self.inner.runtime, future, OnCancel::RunOn)
+    }
+
+    /// As [`Self::batch_abort`], but blocks.
+    fn batch_abort_blocking(&self, py: Python<'_>, batch: &PersistBatch) -> PyResult<()> {
+        let future = self.start_batch_abort(batch);
+        self.block_on(py, future)
+    }
+
     /// As [`Self::prepare_chunk`], but blocks.
     #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))]
     fn prepare_chunk_blocking(
@@ -497,6 +488,416 @@ impl Persister {
     }
 }
 
+/// The events of `events_and_contexts`, sharing their data, internal
+/// metadata included, with the Python events, so that Rust can read them
+/// without the GIL and set their stream orderings.
+fn shallow_events(
+    py: Python<'_>,
+    events_and_contexts: Vec<(Py<Event>, PersistContext)>,
+) -> Vec<(Event, PersistContext)> {
+    events_and_contexts
+        .into_iter()
+        .map(|(event, context)| (event.bind(py).get().shallow_clone(), context))
+        .collect()
+}
+
+/// Prepare the chunk `input`, with what it counted.
+async fn prepare_chunk(
+    inner: &PersisterInner,
+    pool: &ConnectionPool,
+    input: &PrepareInput,
+    shadow: bool,
+) -> anyhow::Result<(Prepared, PersistCounts)> {
+    let resolver = NativeStateResolver {
+        pool: pool.clone(),
+        handle: inner.runtime.tokio_handle()?,
+        caches: Arc::clone(&inner.store.caches),
+        conflict_cache: Arc::clone(&inner.conflict_cache),
+    };
+    let reader = PersistReader::new(inner.sync.reader(pool.clone()));
+    let options = PrepareOptions {
+        server_name: &inner.server_name,
+        clock: inner.runtime.clock(),
+        sticky: inner.msc4354_enabled,
+        dry_run: shadow,
+    };
+    // Python reads in `runInteraction`s, which retry on a lost connection
+    // or a cancelled statement. Here the whole chunk is prepared again.
+    // That is safe as its only write, bumping the state groups it uses, can
+    // be repeated, and runs in a write interaction with its own retries.
+    let mut attempt = 1;
+    loop {
+        let mut counts = PersistCounts::default();
+        match prepare::prepare_chunk(&reader, &resolver, pool, input, &options, &mut counts).await {
+            Ok(prepared) => return Ok((prepared, counts)),
+            Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => {
+                log::warn!("[TXN RETRY] {{prepare_chunk}} {e:#} {attempt}/{MAX_ATTEMPTS}");
+                attempt += 1;
+            }
+            Err(e) => return Err(e),
+        }
+    }
+}
+
+/// Write the chunk `events`, which Rust prepared as `chunk`, taking its
+/// stream orderings first if `orderings` says so. Resolves with the
+/// effects, whose stream IDs are unfinished.
+async fn write_chunk(
+    inner: &PersisterInner,
+    pool: &ConnectionPool,
+    room_id: &str,
+    events: Vec<(Event, PersistContext)>,
+    chunk: &prepare::PreparedChunk,
+    orderings: StreamOrderings,
+    inhibit_local_membership_updates: bool,
+) -> PyResult<PersistEffects> {
+    // Taken as late as possible, as Python takes them: until the chunk is
+    // written, no later event on the stream can be announced.
+    let event_stream = match orderings {
+        StreamOrderings::Set => None,
+        StreamOrderings::Allocate { backfilled } => {
+            let stream = if backfilled {
+                TxnStream::backfill(Arc::clone(&inner.backfill_stream))
+            } else {
+                TxnStream::events(Arc::clone(&inner.events_stream))
+            };
+            let ids = allocate_stream_orderings(pool, &stream, events.len())
+                .await
+                .map_err(persist_error_to_py)?;
+            set_persisted_positions(inner, &events, &ids, backfilled)
+                .map_err(persist_error_to_py)?;
+            Some((stream, ids, backfilled))
+        }
+    };
+    let events = events
+        .into_iter()
+        .map(|(event, context)| TxnEvent::new(event, context))
+        .collect::<anyhow::Result<Vec<_>>>()
+        .map_err(persist_error_to_py)?;
+
+    let cache_stream = TxnStream::caches(Arc::clone(&inner.cache_stream));
+    let sticky_stream = TxnStream::sticky_events(Arc::clone(&inner.sticky_stream));
+    let profile_stream = TxnStream::profile_updates(Arc::clone(&inner.profile_stream));
+    let input = TxnInput {
+        room_id,
+        events: &events,
+        prepared: chunk,
+        instance_name: &inner.instance_name,
+        server_name: &inner.server_name,
+        clock: inner.runtime.clock(),
+        inhibit_local_membership_updates,
+        enable_search: inner.enable_search,
+        enable_ephemeral_messages: inner.enable_ephemeral_messages,
+        msc4354_enabled: inner.msc4354_enabled,
+        include_profile_updates_in_sync: inner.include_profile_updates_in_sync,
+        cache_stream: &cache_stream,
+        sticky_stream: &sticky_stream,
+        profile_stream: &profile_stream,
+    };
+    let mut effects = run_write_interaction(pool, "persist_events", |txn, _| {
+        let input = &input;
+        async move {
+            let mut effects = PersistEffects::default();
+            txn::persist_events_txn(txn, input, &mut effects).await?;
+            Ok(effects)
+        }
+        .boxed()
+    })
+    .await
+    .map_err(write_error_to_py)?;
+    cache_stream.hand_over(effects.cache_stream_ids());
+    sticky_stream.hand_over(effects.sticky_stream_ids());
+    profile_stream.hand_over(effects.profile_stream_ids());
+    // Finished on drop if the transaction failed, as `get_next_mult`
+    // finishes them when its block raises.
+    if let Some((stream, ids, backfilled)) = event_stream {
+        stream.hand_over(&ids);
+        effects.add_stream_ids(ids, backfilled);
+    }
+    Ok(effects)
+}
+
+/// A queued batch of one room's events, persisted chunk by chunk by
+/// [`batch_next`]. See [`Persister::persist_batch`].
+struct BatchState {
+    inner: Arc<PersisterInner>,
+    pool: ConnectionPool,
+    room_id: String,
+    backfilled: bool,
+    events: Vec<(Event, PersistContext)>,
+    /// Set by the first step.
+    deduplicated: Option<Deduplicated>,
+    /// The position in `kept` of the next chunk.
+    next: usize,
+    /// The state groups are marked as being persisted for the chunk last
+    /// handed to Python, which the next step or [`batch_abort`] releases.
+    guard_held: bool,
+    /// The batch has been persisted, handed back or failed.
+    done: bool,
+}
+
+/// The batch once the events already persisted are dropped.
+struct Deduplicated {
+    replaced_events: HashMap<String, String>,
+    /// The positions in the batch of the events to persist.
+    kept: Vec<usize>,
+    /// What `persisting_state_group_references` holds back from deletion
+    /// around each chunk: the groups of the whole batch, as in Python.
+    state_groups: BTreeSet<i64>,
+}
+
+impl BatchState {
+    fn shallow(&self, i: usize) -> (Event, PersistContext) {
+        (self.events[i].0.shallow_clone(), self.events[i].1.clone())
+    }
+
+    /// `_finish_persisting_txn` for the chunk last written, if its state
+    /// groups are still marked.
+    async fn release_guard(&mut self, error: bool) -> PyResult<()> {
+        if !std::mem::take(&mut self.guard_held) {
+            return Ok(());
+        }
+        let Some(deduplicated) = &self.deduplicated else {
+            return Ok(());
+        };
+        batch::finish_persisting(
+            &self.pool,
+            self.inner.runtime.clock(),
+            &self.inner.instance_name,
+            &deduplicated.state_groups,
+            error,
+        )
+        .await
+        .map_err(write_error_to_py)
+    }
+}
+
+impl Drop for BatchState {
+    /// A batch dropped with its guard held, which Python should not let
+    /// happen, releases it as a failure, as Python's `finally` would.
+    fn drop(&mut self) {
+        if !self.guard_held {
+            return;
+        }
+        let (Some(deduplicated), Ok(handle)) =
+            (self.deduplicated.take(), self.inner.runtime.tokio_handle())
+        else {
+            return;
+        };
+        let inner = Arc::clone(&self.inner);
+        let pool = self.pool.clone();
+        handle.spawn(async move {
+            let result = batch::finish_persisting(
+                &pool,
+                inner.runtime.clock(),
+                &inner.instance_name,
+                &deduplicated.state_groups,
+                true,
+            )
+            .await;
+            if let Err(e) = result {
+                log::warn!("Failed to release the state groups of a dropped batch: {e}");
+            }
+        });
+    }
+}
+
+/// What one [`batch_next`] step did.
+enum BatchStep {
+    /// A chunk was written, and its effects are Python's to apply.
+    Written(WrittenChunk),
+    /// The batch is done.
+    Done(PersistedBatch),
+}
+
+impl<'py> IntoPyObject<'py> for BatchStep {
+    type Target = PyAny;
+    type Output = Bound<'py, PyAny>;
+    type Error = PyErr;
+
+    fn into_pyobject(self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
+        match self {
+            BatchStep::Written(chunk) => Ok(Bound::new(py, chunk)?.into_any()),
+            BatchStep::Done(batch) => Ok(Bound::new(py, batch)?.into_any()),
+        }
+    }
+}
+
+/// One step of `_persist_event_batch` for `state`. The first step drops the
+/// events already persisted under the same transaction ID
+/// (`get_already_persisted_events`). Each step releases the state groups
+/// of the chunk before it, then prepares the next chunk, marks the batch's
+/// state groups as being persisted, takes the chunk's stream orderings and
+/// writes it.
+///
+/// A chunk Rust cannot prepare, because its state resolution falls back to
+/// Python or preparing it fails, is handed back with the rest of the batch
+/// in [`PersistedBatch::remaining`], unwritten.
+async fn batch_next(state: &mut BatchState) -> PyResult<BatchStep> {
+    if state.done {
+        return Err(PyRuntimeError::new_err("the batch is done"));
+    }
+    let result = batch_step(state).await;
+    if !matches!(result, Ok(BatchStep::Written(_))) {
+        state.done = true;
+    }
+    result
+}
+
+async fn batch_step(state: &mut BatchState) -> PyResult<BatchStep> {
+    if state.deduplicated.is_none() {
+        // Retried as Python's `runInteraction`s are.
+        let mut attempt = 1;
+        let replaced_events = loop {
+            let result = async {
+                let querier = PoolQuerier::checkout(&state.pool).await?;
+                batch::already_persisted_events(&querier, &state.events).await
+            }
+            .await;
+            match result {
+                Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => {
+                    log::warn!(
+                        "[TXN RETRY] {{get_already_persisted_events}} {e:#} {attempt}/{MAX_ATTEMPTS}"
+                    );
+                    attempt += 1;
+                }
+                result => break result.map_err(persist_error_to_py)?,
+            }
+        };
+        let kept: Vec<usize> = (0..state.events.len())
+            .filter(|&i| !replaced_events.contains_key(state.events[i].0.event_id()))
+            .collect();
+        let kept_events: Vec<_> = kept.iter().map(|&i| state.shallow(i)).collect();
+        let state_groups =
+            batch::referenced_state_groups(&kept_events).map_err(persist_error_to_py)?;
+        state.deduplicated = Some(Deduplicated {
+            replaced_events,
+            kept,
+            state_groups,
+        });
+    }
+    state.release_guard(false).await?;
+
+    let deduplicated = state.deduplicated.as_ref().expect("set above");
+    let start = state.next;
+    if start >= deduplicated.kept.len() {
+        return Ok(BatchStep::Done(PersistedBatch {
+            replaced_events: deduplicated.replaced_events.clone(),
+            remaining: Vec::new(),
+            fallback_reason: None,
+        }));
+    }
+    let end = (start + batch::CHUNK_SIZE).min(deduplicated.kept.len());
+    let indices = deduplicated.kept[start..end].to_vec();
+    let chunk: Vec<(Event, PersistContext)> = indices.iter().map(|&i| state.shallow(i)).collect();
+    let hand_back = |reason: String| {
+        BatchStep::Done(PersistedBatch {
+            replaced_events: deduplicated.replaced_events.clone(),
+            remaining: deduplicated.kept[start..].to_vec(),
+            fallback_reason: Some(reason),
+        })
+    };
+
+    let inner = &state.inner;
+    let pool = &state.pool;
+    let prepared = async {
+        let input = PrepareInput::new(state.room_id.clone(), &chunk, state.backfilled)?;
+        prepare_chunk(inner, pool, &input, false).await
+    }
+    .await;
+    let (prepared, counts) = match prepared {
+        Ok((Prepared::StateResFallback(reason), _)) => return Ok(hand_back(reason)),
+        // The chunk path prepares it again, and raises or falls back to
+        // Python as the error and the strictness say.
+        Err(e) => {
+            log::warn!(
+                "Failed to prepare events for {} in Rust, handing them back: {e:#}",
+                state.room_id
+            );
+            return Ok(hand_back(format!("{e:#}")));
+        }
+        Ok(prepared) => prepared,
+    };
+    inner.metrics.record(|total| total.add(counts));
+    let prepared = Arc::new(prepared);
+    let Prepared::Chunk(prepared_chunk) = &*prepared else {
+        unreachable!("handed back above");
+    };
+
+    let guarded = !deduplicated.state_groups.is_empty();
+    if guarded {
+        batch::mark_state_groups_as_persisting(
+            pool,
+            inner.runtime.clock(),
+            &inner.instance_name,
+            &deduplicated.state_groups,
+        )
+        .await
+        .map_err(write_error_to_py)?;
+        state.guard_held = true;
+    }
+    let written = write_chunk(
+        inner,
+        pool,
+        &state.room_id,
+        chunk,
+        prepared_chunk,
+        StreamOrderings::Allocate {
+            backfilled: state.backfilled,
+        },
+        state.backfilled,
+    )
+    .await;
+    let effects = match written {
+        Ok(effects) => effects,
+        Err(e) => {
+            // As Python's `finally`, whose error replaces the persist's.
+            state.release_guard(true).await?;
+            return Err(e);
+        }
+    };
+    state.next = end;
+    Ok(BatchStep::Written(WrittenChunk {
+        effects,
+        prepared,
+        indices,
+    }))
+}
+
+/// [`batch_next`]'s chunk, written, whose effects Python must apply.
+#[pyclass(frozen)]
+pub struct WrittenChunk {
+    effects: PersistEffects,
+    prepared: Arc<Prepared>,
+    indices: Vec<usize>,
+}
+
+#[pymethods]
+impl WrittenChunk {
+    /// What Python applies with `apply_rust_persist_effects`, with the
+    /// chunk's stream IDs unfinished.
+    #[getter]
+    fn effects(&self) -> PersistEffects {
+        self.effects.clone()
+    }
+
+    /// What Rust prepared for the chunk.
+    #[getter]
+    fn prepared(&self) -> PreparedChunk {
+        PreparedChunk {
+            prepared: Arc::clone(&self.prepared),
+            // Counted when it was prepared.
+            counts: Mutex::new(None),
+        }
+    }
+
+    /// The chunk's positions in the batch.
+    #[getter]
+    fn indices(&self) -> Vec<usize> {
+        self.indices.clone()
+    }
+}
+
 /// Where the stream orderings of a chunk come from.
 #[derive(Clone, Copy)]
 enum StreamOrderings {
@@ -560,6 +961,44 @@ fn set_persisted_positions(
     Ok(())
 }
 
+/// A batch being persisted. See [`Persister::persist_batch`].
+#[pyclass(frozen)]
+pub struct PersistBatch {
+    state: Arc<tokio::sync::Mutex<BatchState>>,
+}
+
+/// What [`Persister::persist_batch`] did, once the batch is done.
+#[pyclass(frozen)]
+pub struct PersistedBatch {
+    replaced_events: HashMap<String, String>,
+    remaining: Vec<usize>,
+    fallback_reason: Option<String>,
+}
+
+#[pymethods]
+impl PersistedBatch {
+    /// The events not persisted as they repeat an event already persisted
+    /// with the same transaction ID, mapped to that event.
+    #[getter]
+    fn replaced_events(&self) -> HashMap<String, String> {
+        self.replaced_events.clone()
+    }
+
+    /// The positions in the batch of the events left for Python to
+    /// persist, in order: the chunk Rust could not prepare and those after
+    /// it. Empty if Rust persisted them all.
+    #[getter]
+    fn remaining(&self) -> Vec<usize> {
+        self.remaining.clone()
+    }
+
+    /// Why Rust left the `remaining` events.
+    #[getter]
+    fn fallback_reason(&self) -> Option<String> {
+        self.fallback_reason.clone()
+    }
+}
+
 /// What [`Persister::prepare_chunk`] computed, in the shapes Python builds
 /// its own types from. If `state_res_fallback` is set, the rest is empty
 /// and Python should compute the chunk itself.

rust/src/persist/reader.rs

@@ -727,7 +727,7 @@ impl PyPersistReaderForTest {
         now_ms: i64,
         state_dag: bool,
     ) -> PyResult<(StateAndExtremities, PersistCounts)> {
-        let (events, event_map) = chunk(py, &events_and_contexts)?;
+        let (events, event_map) = py_chunk(py, &events_and_contexts)?;
         let resolver = self.resolver();
         self.run(py, |reader| {
             async move {
@@ -1199,7 +1199,7 @@ impl PyPersistReaderForTest {
         HashSet<String>,
         Bound<'py, PyDict>,
     )> {
-        let (events, event_map) = chunk(py, &events_and_contexts)?;
+        let (events, event_map) = py_chunk(py, &events_and_contexts)?;
         let old: HashSet<Arc<str>> = arcs(old_latest_event_ids).into_iter().collect();
         let new: HashSet<Arc<str>> = arcs(new_latest_event_ids).into_iter().collect();
         let resolver = self.resolver();
@@ -1266,7 +1266,7 @@ impl PyPersistReaderForTest {
         to_delete: Vec<(String, String)>,
         to_insert: HashMap<(String, String), String>,
     ) -> PyResult<bool> {
-        let (events, _) = chunk(py, &events_and_contexts)?;
+        let (events, _) = py_chunk(py, &events_and_contexts)?;
         let delta = DeltaState {
             to_delete: to_delete
                 .iter()
@@ -1523,14 +1523,15 @@ fn state_map(map: HashMap<(String, String), String>) -> StateMap {
 /// The chunk's events as the computation reads them, and as state
 /// resolution's map of events in flight.
 pub(super) fn chunk(
-    py: Python<'_>,
-    events_and_contexts: &[(Py<Event>, PersistContext)],
-) -> PyResult<(Vec<NewEvent>, EventMap)> {
-    let events = new_events(py, events_and_contexts)?;
+    events_and_contexts: &[(Event, PersistContext)],
+) -> anyhow::Result<(Vec<NewEvent>, EventMap)> {
+    let events = events_and_contexts
+        .iter()
+        .map(|(event, context)| NewEvent::new(event, context))
+        .collect::<anyhow::Result<_>>()?;
     let event_map = events_and_contexts
         .iter()
         .map(|(event, _)| {
-            let event = event.bind(py).get();
             (
                 Arc::from(event.event_id()),
                 StateResEvent::new(event.shallow_clone()),
@@ -1540,6 +1541,18 @@ pub(super) fn chunk(
     Ok((events, event_map))
 }
 
+/// [`chunk`] of the Python events, for the test entry points.
+fn py_chunk(
+    py: Python<'_>,
+    events_and_contexts: &[(Py<Event>, PersistContext)],
+) -> PyResult<(Vec<NewEvent>, EventMap)> {
+    let events: Vec<(Event, PersistContext)> = events_and_contexts
+        .iter()
+        .map(|(event, context)| (event.bind(py).get().shallow_clone(), context.clone()))
+        .collect();
+    chunk(&events).map_err(persist_error_to_py)
+}
+
 fn new_events(
     py: Python<'_>,
     events_and_contexts: &[(Py<Event>, PersistContext)],

synapse/config/experimental.py

@@ -415,8 +415,9 @@ class ExperimentalConfig(Config):
         #  - "prepare": Rust's results for that are used.
         #  - "transaction": Rust also runs the write transaction.
         #  - "full": Rust persists each queued batch.
-        # The levels are being built up in that order. So far "full" only
-        # has Rust take the stream orderings of each chunk as well.
+        # The levels are being built up in that order. So far at "full",
+        # Python still applies what each chunk Rust writes asks for and runs
+        # what follows it, and recalculates current state on request.
         rust_event_persistence = experimental.get("rust_event_persistence", "off")
         # YAML reads an unquoted `off` as a boolean.
         if rust_event_persistence is False:

synapse/storage/controllers/persist_events.py

@@ -82,7 +82,13 @@ from synapse.storage.databases.main.events import (
     persist_event_counter,
 )
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
-from synapse.synapse_rust.persist import Persister, PreparedChunk
+from synapse.storage.databases.main.rust_persist import finish_rust_persist_stream_ids
+from synapse.synapse_rust.persist import (
+    PersistBatch,
+    Persister,
+    PreparedChunk,
+    WrittenChunk,
+)
 from synapse.types import (
     PersistedEventPosition,
     RoomStreamToken,
@@ -90,7 +96,11 @@ from synapse.types import (
     get_domain_from_id,
 )
 from synapse.types.state import StateFilter
-from synapse.util.async_helpers import ObservableDeferred, yieldable_gather_results
+from synapse.util.async_helpers import (
+    ObservableDeferred,
+    delay_cancellation,
+    yieldable_gather_results,
+)
 from synapse.util.metrics import Measure
 from synapse.util.stringutils import shortstr
 
@@ -871,6 +881,14 @@ class EventsPersistenceStorageController:
         if not events_and_contexts:
             return replaced_events
 
+        if (
+            self._rust_persister is not None
+            and self.hs.config.experimental.rust_event_persistence_at_least("full")
+        ):
+            return await self._persist_event_batch_in_rust(
+                room_id, events_and_contexts, backfilled
+            )
+
         # Check if any of the events have a transaction ID that has already been
         # persisted, and if so we don't persist it again.
         #
@@ -892,6 +910,16 @@ class EventsPersistenceStorageController:
             if not events_and_contexts:
                 return replaced_events
 
+        await self._persist_chunks(room_id, events_and_contexts, backfilled)
+        return replaced_events
+
+    async def _persist_chunks(
+        self,
+        room_id: str,
+        events_and_contexts: list[EventPersistencePair],
+        backfilled: bool,
+    ) -> None:
+        """Persist the events, deduplicated, in chunks of 100."""
         chunks = [
             events_and_contexts[x : x + 100]
             for x in range(0, len(events_and_contexts), 100)
@@ -924,7 +952,98 @@ class EventsPersistenceStorageController:
                     rust_transaction=self._rust_transaction(prepared),
                 )
 
-        return replaced_events
+    async def _persist_event_batch_in_rust(
+        self,
+        room_id: str,
+        events_and_contexts: list[EventPersistencePair],
+        backfilled: bool,
+    ) -> dict[str, str]:
+        """`_persist_event_batch` at `rust_event_persistence: full`: Rust
+        deduplicates the batch, then prepares and writes it chunk by chunk,
+        and Python applies what each chunk Rust wrote asks for before the
+        next. A chunk Rust cannot prepare is handed back with the rest of
+        the batch and persisted chunk by chunk, as at lower levels."""
+        assert self._rust_persister is not None
+        persister = self._rust_persister
+        handle = persister.persist_batch(
+            room_id, events_and_contexts, backfilled=backfilled
+        )
+
+        async def persist() -> dict[str, str]:
+            while True:
+                try:
+                    # The Deferred follows the logcontext rules already.
+                    step = await persister.batch_next(handle)
+                except BaseException:
+                    # Rust finished the stream orderings of a chunk it
+                    # failed to write. Tell replication, as `get_next_mult`
+                    # does when its block raises.
+                    self.hs.get_replication_notifier().notify_replication()
+                    raise
+                if not isinstance(step, WrittenChunk):
+                    break
+                await self._apply_rust_batch_chunk(
+                    room_id, events_and_contexts, backfilled, step, handle
+                )
+
+            if step.remaining:
+                logger.debug(
+                    "Persisting %d events of %s chunk by chunk: %s",
+                    len(step.remaining),
+                    room_id,
+                    step.fallback_reason,
+                )
+                await self._persist_chunks(
+                    room_id,
+                    [events_and_contexts[i] for i in step.remaining],
+                    backfilled,
+                )
+            return step.replaced_events
+
+        # Rust carries on with a step if this is cancelled, and its effects
+        # must be applied. Waiting also keeps the room's queue from starting
+        # its next batch while this one, or the rest Rust handed back, is
+        # still being written.
+        return await delay_cancellation(persist())
+
+    async def _apply_rust_batch_chunk(
+        self,
+        room_id: str,
+        events_and_contexts: list[EventPersistencePair],
+        backfilled: bool,
+        step: WrittenChunk,
+        handle: PersistBatch,
+    ) -> None:
+        """Apply what a chunk of the batch Rust wrote asks for, and finish
+        its stream orderings. If that fails, release the batch."""
+        assert self._rust_persister is not None
+        effects = step.effects
+        try:
+            try:
+                prepared = step.prepared
+                chunk = [events_and_contexts[i] for i in step.indices]
+                # Counted as the chunk path counts, once prepared.
+                self._count_rust_prepare("full", "rust")
+                if prepared.state_res_stats is not None:
+                    self._state_resolution_handler.record_rust_state_res_metrics(
+                        room_id, prepared.state_res_stats
+                    )
+            except BaseException:
+                # Only the apply finishes the orderings otherwise.
+                finish_rust_persist_stream_ids(
+                    self.persist_events_store.store, effects, succeeded=False
+                )
+                raise
+            await self.persist_events_store._apply_rust_chunk(
+                room_id,
+                chunk,
+                effects,
+                backfilled=backfilled,
+                new_forward_extremities=prepared.new_forward_extremities,
+            )
+        except BaseException:
+            await self._rust_persister.batch_abort(handle)
+            raise
 
     def _rust_transaction(self, prepared: _PreparedChunk) -> RustTransaction | None:
         """Whether to write the chunk in Rust: at `rust_event_persistence:
@@ -1071,10 +1190,9 @@ class EventsPersistenceStorageController:
         with Measure(
             self._clock, name="rust_prepare_chunk", server_name=self.server_name
         ):
-            rust = await make_deferred_yieldable(
-                self._rust_persister.prepare_chunk(
-                    room_id, chunk, backfilled, shadow=shadow
-                )
+            # The Deferred follows the logcontext rules already.
+            rust = await self._rust_persister.prepare_chunk(
+                room_id, chunk, backfilled, shadow=shadow
             )
         if rust.state_res_fallback is not None:
             logger.info(

synapse/storage/databases/main/events.py

@@ -57,7 +57,6 @@ from synapse.events import (
 )
 from synapse.events.py_protocol import MSC4242Event, supports_msc4242_state_dag
 from synapse.events.snapshot import EventPersistencePair
-from synapse.logging.context import make_deferred_yieldable
 from synapse.logging.opentracing import trace
 from synapse.metrics import SERVER_NAME_LABEL
 from synapse.storage._base import db_to_json, make_in_list_sql_clause
@@ -98,7 +97,7 @@ from synapse.util.stringutils import non_null_str_or_none
 if TYPE_CHECKING:
     from synapse.server import HomeServer
     from synapse.storage.databases.main import DataStore
-    from synapse.synapse_rust.persist import Persister, PreparedChunk
+    from synapse.synapse_rust.persist import PersistEffects, Persister, PreparedChunk
 
 
 logger = logging.getLogger(__name__)
@@ -1195,18 +1194,17 @@ class PersistEventsStore:
         # caller waits for both, as Python's transaction cannot be cancelled.
 
         async def persist_and_apply() -> None:
+            # The persister's Deferreds follow the logcontext rules already.
             persister = rust_transaction.persister
             allocate = rust_transaction.allocate_stream_orderings
             if allocate:
                 try:
-                    effects = await make_deferred_yieldable(
-                        persister.persist_chunk(
-                            room_id,
-                            events_and_contexts,
-                            rust_transaction.prepared,
-                            backfilled=backfilled,
-                            inhibit_local_membership_updates=inhibit_local_membership_updates,
-                        )
+                    effects = await persister.persist_chunk(
+                        room_id,
+                        events_and_contexts,
+                        rust_transaction.prepared,
+                        backfilled=backfilled,
+                        inhibit_local_membership_updates=inhibit_local_membership_updates,
                     )
                 except BaseException:
                     # Rust finished the stream orderings it took. Tell
@@ -1215,44 +1213,71 @@ class PersistEventsStore:
                     self.hs.get_replication_notifier().notify_replication()
                     raise
             else:
-                effects = await make_deferred_yieldable(
-                    persister.persist_events_txn(
-                        room_id,
-                        events_and_contexts,
-                        rust_transaction.prepared,
-                        inhibit_local_membership_updates=inhibit_local_membership_updates,
-                    )
-                )
-            # The copy of a duplicated event that Rust wrote, as Python's
-            # `_add_to_cache` reads the deduplicated chunk.
-            written = self._filter_events_and_contexts_for_duplicates(
-                events_and_contexts
-            )
-            events = {event.event_id: event for event, _ in written}
-            if not allocate:
-                apply_rust_persist_effects(self.store, effects, events)
-                return
-            # What follows the write runs before the stream orderings are
-            # finished, as it runs inside `get_next_mult` in Python: the
-            # events must not be announced before, say, the MSC4293
-            # redactions they start are recorded.
-            try:
-                apply_rust_persist_effects(
-                    self.store, effects, events, finish_stream_ids=False
-                )
-                await self._after_persisting_events(
+                effects = await persister.persist_events_txn(
                     room_id,
                     events_and_contexts,
-                    use_negative_stream_ordering=backfilled,
+                    rust_transaction.prepared,
+                    inhibit_local_membership_updates=inhibit_local_membership_updates,
+                )
+            if allocate:
+                await self._apply_rust_chunk(
+                    room_id,
+                    events_and_contexts,
+                    effects,
+                    backfilled=backfilled,
                     new_forward_extremities=new_forward_extremities,
                 )
-            except BaseException:
-                finish_rust_persist_stream_ids(self.store, effects, succeeded=False)
-                raise
-            finish_rust_persist_stream_ids(self.store, effects, succeeded=True)
+            else:
+                apply_rust_persist_effects(
+                    self.store, effects, self._events_written(events_and_contexts)
+                )
 
         await delay_cancellation(persist_and_apply())
 
+    def _events_written(
+        self, events_and_contexts: list[EventPersistencePair]
+    ) -> dict[str, EventBase]:
+        """The copy of each event that a write of the chunk stores, as
+        Python's `_add_to_cache` reads the deduplicated chunk."""
+        written = self._filter_events_and_contexts_for_duplicates(events_and_contexts)
+        return {event.event_id: event for event, _ in written}
+
+    async def _apply_rust_chunk(
+        self,
+        room_id: str,
+        events_and_contexts: list[EventPersistencePair],
+        effects: "PersistEffects",
+        *,
+        backfilled: bool,
+        new_forward_extremities: set[str] | None,
+    ) -> None:
+        """Apply the effects of a chunk Rust wrote, having taken its stream
+        orderings, then run `_after_persisting_events`, then finish the
+        orderings.
+
+        What follows the write runs before the orderings are finished, as
+        it runs inside `get_next_mult` in Python: the events must not be
+        announced before, say, the MSC4293 redactions they start are
+        recorded. The orderings are finished whatever happens.
+        """
+        try:
+            apply_rust_persist_effects(
+                self.store,
+                effects,
+                self._events_written(events_and_contexts),
+                finish_stream_ids=False,
+            )
+            await self._after_persisting_events(
+                room_id,
+                events_and_contexts,
+                use_negative_stream_ordering=backfilled,
+                new_forward_extremities=new_forward_extremities,
+            )
+        except BaseException:
+            finish_rust_persist_stream_ids(self.store, effects, succeeded=False)
+            raise
+        finish_rust_persist_stream_ids(self.store, effects, succeeded=True)
+
     def _persist_events_txn(
         self,
         txn: LoggingTransaction,

synapse/synapse_rust/persist.pyi

@@ -106,6 +106,71 @@ class Persister:
     ) -> PersistEffects:
         """As `persist_chunk`, but blocks. For tests."""
 
+    def persist_batch(
+        self,
+        room_id: str,
+        events_and_contexts: list[tuple[EventBase, EventContext]],
+        *,
+        backfilled: bool,
+    ) -> PersistBatch:
+        """`_persist_event_batch` for a queued batch of `room_id`'s events,
+        one chunk at a time: await `batch_next` until it resolves with a
+        `PersistedBatch`, applying the effects of each `WrittenChunk` before
+        the next step, and call `batch_abort` if they cannot be applied.
+        Converts the events, without I/O."""
+    def batch_next(
+        self, batch: PersistBatch
+    ) -> Deferred[WrittenChunk | PersistedBatch]:
+        """The batch's next step. The first drops the events already
+        persisted under the same transaction ID. Each releases the state
+        groups of the chunk before, then prepares the next chunk, marks the
+        batch's state groups as being persisted, takes the chunk's stream
+        orderings and writes it.
+
+        The chunk's stream IDs are handed over unfinished: must not be
+        cancelled. If the step fails, Rust finishes them and releases the
+        state groups."""
+    def batch_next_blocking(self, batch: PersistBatch) -> WrittenChunk | PersistedBatch:
+        """As `batch_next`, but blocks. For tests."""
+    def batch_abort(self, batch: PersistBatch) -> Deferred[None]:
+        """Give up the batch, releasing its state groups as a failed
+        persist does."""
+    def batch_abort_blocking(self, batch: PersistBatch) -> None:
+        """As `batch_abort`, but blocks. For tests."""
+
+class PersistBatch:
+    """A batch being persisted. See `Persister.persist_batch`."""
+
+class WrittenChunk:
+    """A chunk `Persister.batch_next` wrote."""
+
+    @property
+    def effects(self) -> PersistEffects:
+        """To apply with `apply_rust_persist_effects`, with the chunk's
+        stream IDs unfinished."""
+    @property
+    def prepared(self) -> PreparedChunk:
+        """What Rust prepared for the chunk."""
+    @property
+    def indices(self) -> list[int]:
+        """The chunk's positions in the batch."""
+
+class PersistedBatch:
+    """What `Persister.persist_batch` did, once the batch is done."""
+
+    @property
+    def replaced_events(self) -> dict[str, str]:
+        """The events not persisted as they repeat an event already
+        persisted with the same transaction ID, mapped to that event."""
+    @property
+    def remaining(self) -> list[int]:
+        """The positions in the batch of the events left for Python to
+        persist, in order: the chunk Rust could not prepare and those after
+        it. Empty if Rust persisted them all."""
+    @property
+    def fallback_reason(self) -> str | None:
+        """Why Rust left the `remaining` events."""
+
 class PreparedChunk:
     """What `Persister.prepare_chunk` computed. If `state_res_fallback` is
     set, the rest is empty and Python should compute the chunk itself."""

tests/storage/test_rust_persist_prepare.py

@@ -13,7 +13,7 @@
 #
 
 """How `_persist_event_batch` prepares chunks in Rust
-(`rust/src/persist/prepare.rs`) at the `shadow` and `prepare` levels of
+(`rust/src/persist/prepare.rs`) at the `shadow`, `prepare` and `full` levels of
 `rust_event_persistence`. These need the native Rust Postgres driver:
 
     SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.storage.test_rust_persist_prepare
@@ -21,7 +21,7 @@
 
 from typing import Any
 from unittest import SkipTest
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import attr
 
@@ -118,6 +118,30 @@ class _PrepareTestBase(unittest.HomeserverTestCase):
             if after[result] != self._counts_before[result]
         }
 
+    def _mock_persister(self, **kwargs: Any) -> Mock:
+        """A mock persister. At `full` it hands each batch back to be
+        persisted chunk by chunk, so the mocked `prepare_chunk` is used."""
+
+        def persist_batch(
+            room_id: str, events_and_contexts: list, **kwargs: Any
+        ) -> Mock:
+            return Mock(events=len(events_and_contexts))
+
+        def batch_next(handle: Mock) -> defer.Deferred:
+            return defer.succeed(
+                Mock(
+                    replaced_events={},
+                    remaining=list(range(handle.events)),
+                    fallback_reason="mocked",
+                )
+            )
+
+        return Mock(
+            persist_batch=Mock(side_effect=persist_batch),
+            batch_next=Mock(side_effect=batch_next),
+            **kwargs,
+        )
+
     def _send(self, body: str, expect_code: int = 200) -> str | None:
         return self.helper.send(
             self.room_id, body, tok=self.alice_tok, expect_code=expect_code
@@ -308,7 +332,7 @@ class ShadowTestCase(_PrepareTestBase):
 
     def test_rust_rejection_python_accepts(self) -> None:
         self.persistence._rust_persist_strict = False
-        self.persistence._rust_persister = Mock(
+        self.persistence._rust_persister = self._mock_persister(
             prepare_chunk=Mock(return_value=defer.fail(SynapseError(500, "no")))
         )
         with self.assertLogs("synapse.storage.controllers.persist_events", "WARNING"):
@@ -317,7 +341,7 @@ class ShadowTestCase(_PrepareTestBase):
 
     def test_rust_error_is_logged(self) -> None:
         self.persistence._rust_persist_strict = False
-        self.persistence._rust_persister = Mock(
+        self.persistence._rust_persister = self._mock_persister(
             prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom")))
         )
         with self.assertLogs("synapse.storage.controllers.persist_events", "ERROR"):
@@ -326,7 +350,7 @@ class ShadowTestCase(_PrepareTestBase):
 
 
 class PrepareTestCase(_PrepareTestBase):
-    level = "prepare"
+    level: RustEventPersistenceLevel = "prepare"
 
     def test_room_lifecycle(self) -> None:
         chunks = self._room_lifecycle()
@@ -424,7 +448,7 @@ class PrepareTestCase(_PrepareTestBase):
     def test_rejection_is_raised(self) -> None:
         """A chunk Rust rejects is not computed again in Python."""
         self.persistence._rust_persist_strict = False
-        self.persistence._rust_persister = Mock(
+        self.persistence._rust_persister = self._mock_persister(
             prepare_chunk=Mock(return_value=defer.fail(SynapseError(500, "no")))
         )
         in_python = Mock(side_effect=AssertionError("computed in Python"))
@@ -434,7 +458,7 @@ class PrepareTestCase(_PrepareTestBase):
         in_python.assert_not_called()
 
     def test_fallback(self) -> None:
-        self.persistence._rust_persister = Mock(
+        self.persistence._rust_persister = self._mock_persister(
             prepare_chunk=Mock(
                 return_value=defer.succeed(Mock(state_res_fallback="no index"))
             )
@@ -448,7 +472,7 @@ class PrepareTestCase(_PrepareTestBase):
 
     def test_error_uses_python(self) -> None:
         self.persistence._rust_persist_strict = False
-        self.persistence._rust_persister = Mock(
+        self.persistence._rust_persister = self._mock_persister(
             prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom")))
         )
         with self.assertLogs("synapse.storage.controllers.persist_events", "ERROR"):
@@ -460,8 +484,55 @@ class PrepareTestCase(_PrepareTestBase):
         )
 
     def test_error_fails_strict_mode(self) -> None:
-        self.persistence._rust_persister = Mock(
+        self.persistence._rust_persister = self._mock_persister(
             prepare_chunk=Mock(return_value=defer.fail(RuntimeError("boom")))
         )
         self._send("hello", expect_code=500)
         self.assertEqual(self._counted(), {"error": 1})
+
+
+class FullTestCase(PrepareTestCase):
+    """At `full`, Rust persists each batch, chunk by chunk, and hands back
+    to the chunk path any chunk it cannot prepare. The prepare counter
+    counts the same."""
+
+    level: RustEventPersistenceLevel = "full"
+
+    def test_hand_back_in_the_middle(self) -> None:
+        """A batch whose second chunk Rust cannot prepare: Rust writes the
+        first and hands back the second alone, which Python persists."""
+        room_id = self.helper.create_room_as(
+            self.alice, room_version=RoomVersions.V1.identifier, tok=self.alice_tok
+        )
+        # Siblings, so the messages need no state resolution and the topics,
+        # in the second chunk, need it on state resolution v1.
+        messages = [
+            self._new_event(
+                room_id, EventTypes.Message, {"msgtype": "m.text", "body": str(i)}
+            )
+            for i in range(100)
+        ]
+        topics = [
+            self._new_event(room_id, EventTypes.Topic, {"topic": t}, state_key="")
+            for t in ("A", "B")
+        ]
+        before = self._counts()
+        persist_chunks = self.persistence._persist_chunks
+        handed_back: list[list[str]] = []
+
+        async def recording(
+            room_id: str, events_and_contexts: list, backfilled: bool
+        ) -> None:
+            handed_back.append([event.event_id for event, _ in events_and_contexts])
+            await persist_chunks(room_id, events_and_contexts, backfilled)
+
+        with patch.object(self.persistence, "_persist_chunks", side_effect=recording):
+            self.get_success(self.persistence.persist_events([*messages, *topics]))
+        self.assertEqual(handed_back, [[event.event_id for event, _ in topics]])
+        after = self._counts()
+        counted = {r: after[r] - before[r] for r in _RESULTS if after[r] != before[r]}
+        self.assertEqual(counted, {"rust": 1, "fallback": 1})
+        for event, _ in [*messages, *topics]:
+            self.assertIsNotNone(
+                self.get_success(self.store.get_event(event.event_id, allow_none=True))
+            )

tests/storage/test_rust_persist_prereqs.py

@@ -40,7 +40,11 @@ from synapse.config.experimental import ExperimentalConfig
 from synapse.config.homeserver import HomeServerConfig
 from synapse.events import EventBase, make_event_from_dict
 from synapse.events.snapshot import EventContext
-from synapse.logging.context import make_deferred_yieldable
+from synapse.logging.context import (
+    LoggingContext,
+    current_context,
+    make_deferred_yieldable,
+)
 from synapse.metrics import SERVER_NAME_LABEL, _PreCollectHooks
 from synapse.rest import admin
 from synapse.rest.client import login, room
@@ -52,6 +56,7 @@ from synapse.storage.engines import PostgresEngine
 from synapse.storage.rust_pool import RustConnectionPool
 from synapse.synapse_rust.persist import (
     PersistEffects,
+    WrittenChunk,
     persist_context_for_test,
     persist_gil_stats,
     persistence_json_for_test,
@@ -871,6 +876,8 @@ class _CountingPersister:
         self.transactions = 0
         # Chunks written with Rust taking their stream orderings.
         self.chunks = 0
+        # Batches persisted in Rust.
+        self.batches = 0
         self.last_effects: PersistEffects | None = None
 
     def prepare_chunk(self, *args: Any, **kwargs: Any) -> Any:
@@ -895,6 +902,19 @@ class _CountingPersister:
 
         return self._persister.persist_chunk(*args, **kwargs).addCallback(keep)
 
+    def persist_batch(self, *args: Any, **kwargs: Any) -> Any:
+        self.batches += 1
+        return self._persister.persist_batch(*args, **kwargs)
+
+    def batch_next(self, *args: Any, **kwargs: Any) -> Any:
+        def keep(step: Any) -> Any:
+            if isinstance(step, WrittenChunk):
+                self.chunks += 1
+                self.last_effects = step.effects
+            return step
+
+        return self._persister.batch_next(*args, **kwargs).addCallback(keep)
+
     def __getattr__(self, name: str) -> Any:
         return getattr(self._persister, name)
 
@@ -1031,7 +1051,8 @@ class RustTransactionTestCase(unittest.HomeserverTestCase):
             def persist_events_txn(self, *args: Any, **kwargs: Any) -> Any:
                 d: defer.Deferred = defer.Deferred()
                 gate.append((args, kwargs, d))
-                return d
+                # As the persister's Deferreds do.
+                return make_deferred_yieldable(d)
 
             def __getattr__(self, name: str) -> Any:
                 return getattr(blocking, name)
@@ -1125,16 +1146,20 @@ class RustFullTestCase(unittest.HomeserverTestCase):
         self.token = self.login("alice", "pass")
         self.room_id = self.helper.create_room_as(self.user, tok=self.token)
 
-    def _build(self, body: str) -> tuple[EventBase, EventContext]:
+    def _build(
+        self, body: str, *, txn_id: str | None = None
+    ) -> tuple[EventBase, EventContext]:
+        """A message, sent with `txn_id` from `DEVICE` if given."""
         event, unpersisted_context = self.get_success(
             self.hs.get_event_creation_handler().create_event(
-                create_requester(self.user),
+                create_requester(self.user, device_id="DEVICE" if txn_id else None),
                 {
                     "type": EventTypes.Message,
                     "content": {"msgtype": "m.text", "body": body},
                     "room_id": self.room_id,
                     "sender": self.user,
                 },
+                txn_id=txn_id,
             )
         )
         context = self.get_success(unpersisted_context.persist(event))
@@ -1263,6 +1288,207 @@ class RustFullTestCase(unittest.HomeserverTestCase):
         assert failed is not None
         self.assertGreaterEqual(id_gen.get_current_token(), failed)
 
+    def test_batch_is_one_rust_call(self) -> None:
+        """A batch of 150 events is one Rust call, written in two chunks,
+        with none of it in Python's store entry point."""
+        # Siblings on the room's extremity: an event cannot be built on
+        # one not yet persisted.
+        pairs = [self._build(f"message {i}") for i in range(150)]
+        batches, chunks = self.persister.batches, self.persister.chunks
+        with patch.object(
+            PersistEventsStore,
+            "_persist_events_and_state_updates",
+            side_effect=AssertionError("persisted in Python"),
+        ):
+            self.get_success(self.persistence.persist_events(pairs))
+        self.assertEqual(self.persister.batches, batches + 1)
+        self.assertEqual(self.persister.chunks, chunks + 2)
+        effects = self.persister.last_effects
+        assert effects is not None
+        self.assertEqual(len(effects.stream_ids), 50)
+
+        orderings = [self._event_row(event.event_id)[0] for event, _ in pairs]
+        self.assertEqual(orderings, sorted(orderings))
+        self.assertGreaterEqual(
+            self.store._stream_id_gen.get_current_token(), orderings[-1]
+        )
+
+    def test_transaction_ids(self) -> None:
+        """An event repeating the transaction ID of an earlier one, in the
+        same batch or an earlier one, is not persisted and maps to it."""
+        self.login("alice", "pass", device_id="DEVICE")
+        first = self._build("one", txn_id="txn")
+        again = self._build("one again", txn_id="txn")
+        persisted, _ = self.get_success(self.persistence.persist_events([first, again]))
+        self.assertEqual(
+            [event.event_id for event in persisted], [first[0].event_id] * 2
+        )
+
+        later = self._build("one later", txn_id="txn")
+        persisted, _ = self.get_success(self.persistence.persist_events([later]))
+        self.assertEqual([event.event_id for event in persisted], [first[0].event_id])
+        for event, _ in (again, later):
+            self.assertIsNone(
+                self.get_success(self.store.get_event(event.event_id, allow_none=True))
+            )
+
+    def test_state_groups_held_while_writing(self) -> None:
+        """The batch's state groups are marked as being persisted while each
+        chunk is written, and no longer once it is."""
+        held: list[list[tuple[int, str]]] = []
+        persist_events = self.hs.get_datastores().persist_events
+        assert persist_events is not None
+        apply_chunk = persist_events._apply_rust_chunk
+
+        async def checking(*args: Any, **kwargs: Any) -> None:
+            held.append(
+                await self.store.db_pool.simple_select_list(
+                    "state_groups_persisting",
+                    None,
+                    ("state_group", "instance_name"),
+                )
+            )
+            await apply_chunk(*args, **kwargs)
+
+        event, context = self._build("hello")
+        with patch.object(persist_events, "_apply_rust_chunk", side_effect=checking):
+            self.get_success(self.persistence.persist_events([(event, context)]))
+        assert context.state_group is not None
+        assert context.state_group_before_event is not None
+        groups = {context.state_group, context.state_group_before_event}
+        self.assertEqual(
+            [sorted(rows) for rows in held],
+            [sorted((group, "master") for group in groups)],
+        )
+        rows = self.get_success(
+            self.store.db_pool.simple_select_list(
+                "state_groups_persisting", None, ("state_group",)
+            )
+        )
+        self.assertEqual(rows, [])
+
+    def test_deleted_state_group_fails(self) -> None:
+        """A chunk whose state group was deleted meanwhile is not written."""
+        event, context = self._build("hello")
+        self.get_success(
+            self.store.db_pool.simple_delete_one(
+                "state_groups", {"id": context.state_group}
+            )
+        )
+        failure = self.get_failure(
+            self.persistence.persist_events([(event, context)]), RuntimeError
+        )
+        self.assertIn("state groups have been deleted", str(failure.value))
+        self.assertIsNone(
+            self.get_success(self.store.get_event(event.event_id, allow_none=True))
+        )
+
+    def _pending_deletion(self, state_group: int) -> tuple[int, int] | None:
+        """`(sequence_number, insertion_ts)` of the group's pending deletion."""
+        row = self.get_success(
+            self.store.db_pool.simple_select_one(
+                "state_groups_pending_deletion",
+                {"state_group": state_group},
+                ("sequence_number", "insertion_ts"),
+                allow_none=True,
+            )
+        )
+        return (row[0], row[1]) if row is not None else None
+
+    def test_pending_deletion_after_persisting(self) -> None:
+        """Once a chunk is written, its state groups are no longer pending
+        deletion. If the write fails, their pending deletion is pushed
+        back, as they may now be referenced."""
+        event, context = self._build("hello")
+        assert context.state_group is not None
+        self.get_success(
+            self.store.db_pool.simple_insert(
+                "state_groups_pending_deletion",
+                {"state_group": context.state_group, "insertion_ts": 0},
+            )
+        )
+        before = self._pending_deletion(context.state_group)
+        assert before is not None
+
+        failing, failing_context = self._build("partial")
+        failing_context.partial_state = True
+        self.reactor.advance(1)
+        self.get_failure(
+            self.persistence.persist_events([(failing, failing_context)]),
+            PartialStateConflictError,
+        )
+        after = self._pending_deletion(context.state_group)
+        assert after is not None
+        self.assertGreater(after[0], before[0])
+        self.assertEqual(after[1], self.clock.time_msec())
+
+        self.get_success(self.persistence.persist_events([(event, context)]))
+        self.assertIsNone(self._pending_deletion(context.state_group))
+
+    def test_replaced_event_and_chunks(self) -> None:
+        """A batch whose first event repeats a persisted transaction ID, and
+        whose other 150 go in two chunks: each chunk hands Python the events
+        it wrote."""
+        self.login("alice", "pass", device_id="DEVICE")
+        original = self._build("original", txn_id="txn")
+        self.get_success(self.persistence.persist_events([original]))
+
+        repeat = self._build("repeat", txn_id="txn")
+        new = [self._build(f"message {i}") for i in range(150)]
+        chunks = self.persister.chunks
+        persisted, _ = self.get_success(self.persistence.persist_events([repeat, *new]))
+        self.assertEqual(self.persister.chunks, chunks + 2)
+        effects = self.persister.last_effects
+        assert effects is not None
+        self.assertEqual(len(effects.stream_ids), 50)
+        self.assertEqual(
+            [event.event_id for event in persisted],
+            [original[0].event_id] + [event.event_id for event, _ in new],
+        )
+        self.assertIsNone(
+            self.get_success(self.store.get_event(repeat[0].event_id, allow_none=True))
+        )
+        orderings = [self._event_row(event.event_id)[0] for event, _ in new]
+        self.assertEqual(orderings, sorted(orderings))
+        for i, (event, _) in enumerate(new):
+            stored = self.get_success(self.store.get_event(event.event_id))
+            self.assertEqual(stored.content["body"], f"message {i}")
+
+    def test_same_event_in_two_chunks(self) -> None:
+        """An event in two chunks of a batch is written once, and every
+        ordering the batch took is finished."""
+        first = self._build("first")
+        rest = [self._build(f"message {i}") for i in range(100)]
+        self.get_success(self.persistence.persist_events([first, *rest, first]))
+        rows = self.get_success(
+            self.store.db_pool.simple_select_onecol(
+                "events", {"event_id": first[0].event_id}, "stream_ordering"
+            )
+        )
+        self.assertEqual(len(rows), 1)
+        effects = self.persister.last_effects
+        assert effects is not None
+        self.assertEqual(len(effects.stream_ids), 2)
+        self.assertGreaterEqual(
+            self.store._stream_id_gen.get_current_token(), max(effects.stream_ids)
+        )
+
+    def test_logcontext_kept(self) -> None:
+        """Awaiting the batch resumes in the caller's logcontext, so what
+        follows it is accounted to the caller."""
+        event, context = self._build("hello")
+        seen = []
+
+        async def persist() -> None:
+            with LoggingContext(name="persist", server_name="test") as ctx:
+                await self.persistence._persist_event_batch_in_rust(
+                    self.room_id, [(event, context)], False
+                )
+                seen.append(current_context() is ctx)
+
+        self.get_success(defer.ensureDeferred(persist()))
+        self.assertEqual(seen, [True])
+
     def test_gil_count(self) -> None:
         """Taking the stream orderings and setting them on the events takes
         the GIL at no persist site."""
@@ -1360,64 +1586,70 @@ class RustFullTestCase(unittest.HomeserverTestCase):
         self.assertGreaterEqual(
             self.store._stream_id_gen.get_current_token(), stream_ordering
         )
+        # `batch_abort` released the state groups.
+        rows = self.get_success(
+            self.store.db_pool.simple_select_list(
+                "state_groups_persisting", None, ("state_group",)
+            )
+        )
+        self.assertEqual(rows, [])
 
     def test_cancelled_persist_applies_its_effects(self) -> None:
-        """Cancelling the persist while Rust takes the orderings and writes
-        the chunk waits for both, and the orderings are finished."""
+        """Cancelling the batch while Rust writes a chunk waits for Rust,
+        so the room's queue cannot start another batch alongside, and the
+        chunk's effects are applied and its orderings finished."""
         gate: list[tuple[Any, Any, defer.Deferred]] = []
         blocking = self.persister._persister
 
         class Gated:
-            def persist_chunk(self, *args: Any, **kwargs: Any) -> Any:
+            def batch_next(self, *args: Any, **kwargs: Any) -> Any:
+                # Only the first step, which writes the chunk.
+                if gate:
+                    return blocking.batch_next(*args, **kwargs)
                 d: defer.Deferred = defer.Deferred()
                 gate.append((args, kwargs, d))
-                return d
+                # As the persister's Deferreds do.
+                return make_deferred_yieldable(d)
 
             def __getattr__(self, name: str) -> Any:
                 return getattr(blocking, name)
 
         self.persister._persister = Gated()
 
-        persist_events = self.hs.get_datastores().persist_events
-        assert persist_events is not None
-        original = persist_events._persist_events_in_rust
+        original = self.persistence._persist_event_batch_in_rust
         cancelled = []
 
-        async def cancel_then_wait(*args: Any, **kwargs: Any) -> None:
-            d = defer.ensureDeferred(original(*args, **kwargs))
+        async def cancel_then_wait(*args: Any) -> dict[str, str]:
+            d = defer.ensureDeferred(original(*args))
             d.cancel()
             try:
-                await make_deferred_yieldable(d)
+                return await make_deferred_yieldable(d)
             except defer.CancelledError:
                 cancelled.append(True)
+                return {}
 
+        event, context = self._build("hello")
         with patch.object(
-            persist_events, "_persist_events_in_rust", side_effect=cancel_then_wait
+            self.persistence,
+            "_persist_event_batch_in_rust",
+            side_effect=cancel_then_wait,
         ):
-            channel = make_request(
-                self.reactor,
-                self.site,
-                "PUT",
-                f"/rooms/{self.room_id}/send/m.room.message/txn1",
-                {"msgtype": "m.text", "body": "hello"},
-                access_token=self.token,
-                await_result=False,
+            persisted = defer.ensureDeferred(
+                self.persistence.persist_events([(event, context)])
             )
             self.pump()
             self.assertEqual(len(gate), 1)
-            # Cancelled, but waiting for the write.
+            # Cancelled, but waiting for Rust.
             self.assertEqual(cancelled, [])
             args, kwargs, d = gate[0]
-            blocking.persist_chunk(*args, **kwargs).chainDeferred(d)
-            channel.await_result()
-        self.assertEqual(channel.code, 200, channel.json_body)
-        event_id = channel.json_body["event_id"]
+            blocking.batch_next(*args, **kwargs).chainDeferred(d)
+            self.get_success(persisted)
         self.assertEqual(cancelled, [True])
 
         effects = self.persister.last_effects
         assert effects is not None
         (stream_ordering,) = effects.stream_ids
-        self.assertEqual(self._event_row(event_id)[0], stream_ordering)
+        self.assertEqual(self._event_row(event.event_id)[0], stream_ordering)
         self.assertGreaterEqual(
             self.store._stream_id_gen.get_current_token(), stream_ordering
         )