Run the persister's independent steps at the same time

Preparing a chunk worked out the state, then the chain cover index, and
the batch then marked its state groups as persisting: each one waited
for the one before, though none uses another's result.

The chain cover index, a transaction of its own of about six round
trips, is now worked out at the same time as the state and the sliding
sync changes, on its own connection. Marking the state groups runs at
the same time as preparing. A state error or a fallback to Python is
reported as before, ahead of anything the other step found. If the
chunk is handed back to Python, a mark that went ahead is released as a
failure: that only pushes back any pending deletion of the groups, as
Python's own failure path would, and a failure to release is logged
rather than failing the hand-back. The chain cover index is worked out
even when the state step then falls back, which costs chain IDs and can
leave a row in `event_auth_chain_to_calculate` that Python's persist
of the chunk writes and clears again.

The E2EE monitor's room on jki.re now waits on about 44 round trips in
turn, from about 60.

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

rust/src/persist/prepare.rs

@@ -151,75 +151,88 @@ pub async fn prepare_chunk(
 ) -> anyhow::Result<Prepared> {
     let room_id = input.room_id.as_str();
     let now_ms = || i64::try_from(options.clock.now_millis());
-    let state = if input.backfilled {
-        StateAndExtremities::default()
-    } else {
-        let result = if input.state_dag {
-            state_delta::process_state_dag_forward_extremities_and_state_delta(
-                reader,
-                resolver,
-                room_id,
-                options.server_name,
-                &input.events,
-                &input.event_map,
-                now_ms()?,
-                counts,
-            )
-            .await
+
+    // The state and the chain cover index need nothing from each other, so
+    // they are worked out at the same time, on their own connections.
+    let state_part = async {
+        let state = if input.backfilled {
+            StateAndExtremities::default()
         } else {
-            state_delta::calculate_new_forward_extremities_and_state_delta(
-                reader,
-                resolver,
-                room_id,
-                options.server_name,
-                &input.events,
-                &input.event_map,
-                now_ms()?,
-                counts,
-            )
-            .await
+            let result = if input.state_dag {
+                state_delta::process_state_dag_forward_extremities_and_state_delta(
+                    reader,
+                    resolver,
+                    room_id,
+                    options.server_name,
+                    &input.events,
+                    &input.event_map,
+                    now_ms()?,
+                    counts,
+                )
+                .await
+            } else {
+                state_delta::calculate_new_forward_extremities_and_state_delta(
+                    reader,
+                    resolver,
+                    room_id,
+                    options.server_name,
+                    &input.events,
+                    &input.event_map,
+                    now_ms()?,
+                    counts,
+                )
+                .await
+            };
+            match result {
+                Ok(state) => state,
+                Err(e) => match e.downcast::<StateResFallback>() {
+                    Ok(fallback) => return Ok(Err(Prepared::StateResFallback(fallback.0))),
+                    Err(e) => return Err(e),
+                },
+            }
         };
-        match result {
-            Ok(state) => state,
-            Err(e) => match e.downcast::<StateResFallback>() {
-                Ok(fallback) => return Ok(Prepared::StateResFallback(fallback.0)),
-                Err(e) => return Err(e),
-            },
+
+        let mut sliding_sync = None;
+        let mut sticky_events_to_un_soft_fail = HashSet::new();
+        if let Some(delta) = &state.state_delta {
+            sliding_sync = Some(
+                sliding_sync::calculate_sliding_sync_table_changes(
+                    reader,
+                    room_id,
+                    options.server_name,
+                    &input.sliding_sync_events,
+                    delta,
+                )
+                .await?,
+            );
+            if options.sticky {
+                sticky_events_to_un_soft_fail = sticky::compute_sticky_events_to_un_soft_fail(
+                    reader,
+                    room_id,
+                    &input.sticky_events,
+                    delta,
+                    now_ms()?,
+                )
+                .await?;
+            }
         }
+        anyhow::Ok(Ok((state, sliding_sync, sticky_events_to_un_soft_fail)))
     };
 
-    let mut sliding_sync = None;
-    let mut sticky_events_to_un_soft_fail = HashSet::new();
-    if let Some(delta) = &state.state_delta {
-        sliding_sync = Some(
-            sliding_sync::calculate_sliding_sync_table_changes(
-                reader,
-                room_id,
-                options.server_name,
-                &input.sliding_sync_events,
-                delta,
-            )
-            .await?,
-        );
-        if options.sticky {
-            sticky_events_to_un_soft_fail = sticky::compute_sticky_events_to_un_soft_fail(
-                reader,
-                room_id,
-                &input.sticky_events,
-                delta,
-                now_ms()?,
-            )
-            .await?;
-        }
-    }
-
-    let new_event_links = chain_cover::calculate_chain_cover_index_for_events(
+    let chain_cover = chain_cover::calculate_chain_cover_index_for_events(
         pool,
         room_id,
         &input.chain_cover_events,
         options.dry_run,
-    )
-    .await?;
+    );
+
+    let (state_part, new_event_links) = futures::future::join(state_part, chain_cover).await;
+    // The state's outcome first, as when it was worked out first.
+    let (state, sliding_sync, sticky_events_to_un_soft_fail) = match state_part? {
+        Ok(state_part) => state_part,
+        Err(fallback) => return Ok(fallback),
+    };
+    let new_event_links = new_event_links?;
 
     Ok(Prepared::Chunk(Box::new(PreparedChunk {
         state,

rust/src/persist/python.rs

@@ -741,6 +741,22 @@ impl BatchState {
 
     /// `_finish_persisting_txn` for the chunk last written, if its state
     /// groups are still marked.
+    /// Release the mark of a chunk handed back to Python after the marking,
+    /// which runs alongside preparing, succeeded. As a failure: that only
+    /// pushes back any pending deletion of the groups, as Python's own
+    /// failure path does, where a success would cancel it before Python has
+    /// persisted anything. A failure to release is logged: Python persists
+    /// the chunk either way.
+    async fn release_early_mark(&mut self) {
+        self.guard_held = true;
+        if let Err(e) = self.release_guard(true).await {
+            log::warn!(
+                "Failed to release the state groups of a chunk of {} handed back to Python: {e}",
+                self.room_id
+            );
+        }
+    }
+
     async fn release_guard(&mut self, error: bool) -> PyResult<()> {
         if !std::mem::take(&mut self.guard_held) {
             return Ok(());
@@ -890,13 +906,36 @@ async fn batch_step(state: &mut BatchState) -> PyResult<BatchStep> {
 
     let inner = &state.inner;
     let pool = &state.pool;
-    let prepared = async {
+    let prepare = async {
         let input = PrepareInput::new(state.room_id.clone(), &chunk, state.backfilled)?;
         prepare_chunk(inner, pool, &input, false).await
-    }
-    .await;
+    };
+    // Marking the state groups needs nothing from preparing, so it runs at
+    // the same time, on its own connection.
+    let guarded = !deduplicated.state_groups.is_empty();
+    let mark = async {
+        if !guarded {
+            return Ok(());
+        }
+        batch::mark_state_groups_as_persisting(
+            pool,
+            inner.runtime.clock(),
+            &inner.instance_name,
+            &deduplicated.state_groups,
+        )
+        .await
+    };
+    let (prepared, marked) = futures::future::join(prepare, mark).await;
+    let marked_ok = guarded && marked.is_ok();
     let (prepared, counts) = match prepared {
-        Ok((Prepared::StateResFallback(reason), _)) => return Ok(hand_back(reason)),
+        Ok((Prepared::StateResFallback(reason), _)) => {
+            // Python persists the chunk, and marks the groups itself.
+            let done = hand_back(reason);
+            if marked_ok {
+                state.release_early_mark().await;
+            }
+            return Ok(done);
+        }
         // The chunk path prepares it again, and raises or falls back to
         // Python as the error and the strictness say.
         Err(e) => {
@@ -904,7 +943,11 @@ async fn batch_step(state: &mut BatchState) -> PyResult<BatchStep> {
                 "Failed to prepare events for {} in Rust, handing them back: {e:#}",
                 state.room_id
             );
-            return Ok(hand_back(format!("{e:#}")));
+            let done = hand_back(format!("{e:#}"));
+            if marked_ok {
+                state.release_early_mark().await;
+            }
+            return Ok(done);
         }
         Ok(prepared) => prepared,
     };
@@ -914,16 +957,8 @@ async fn batch_step(state: &mut BatchState) -> PyResult<BatchStep> {
         unreachable!("handed back above");
     };
 
-    let guarded = !deduplicated.state_groups.is_empty();
+    marked.map_err(write_error_to_py)?;
     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;
     }
     // The follow-ups read the events after the write sets their orderings.