diff --git a/rust/src/persist/prepare.rs b/rust/src/persist/prepare.rs index fb22f5d6bc..dcf70c33bc 100644 --- a/rust/src/persist/prepare.rs +++ b/rust/src/persist/prepare.rs @@ -151,75 +151,88 @@ pub async fn prepare_chunk( ) -> anyhow::Result { 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::() { + Ok(fallback) => return Ok(Err(Prepared::StateResFallback(fallback.0))), + Err(e) => return Err(e), + }, + } }; - match result { - Ok(state) => state, - Err(e) => match e.downcast::() { - 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, diff --git a/rust/src/persist/python.rs b/rust/src/persist/python.rs index 5b7b995b47..b37594da47 100644 --- a/rust/src/persist/python.rs +++ b/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 { 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 { "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 { 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.