Recalculate a room's current state in Rust at rust_event_persistence: full

`_update_current_state`, the queue's `_UpdateCurrentStateTask`, becomes
one call to `Persister.update_current_state`:

- `calculate_current_state` reads the room's forward extremities (its
  state DAG extremities in an MSC4242 room) and takes the state of
  their one state group, or resolves several, after checking the groups
  still exist and putting off their deletion, as `resolve_state_groups`
  does.
- The delta from the current state and the sliding sync table changes
  are computed as for a chunk, with no events.
- If anything changed, one stream ID is taken from the events stream,
  as `get_next` does, and `_update_current_state_txn` runs in Rust.
  Python applies the effects, which finishes the ID.

A room on state resolution v1 whose extremities need resolving is handed
back, and Python recalculates its state as before.

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

changelog.d/rust-persist-batch.misc

@@ -0,0 +1 @@
+Persist each queued batch of events in Rust, and recalculate a room's current state there, at `experimental_features.rust_event_persistence: full`.

rust/src/persist/mod.rs

@@ -64,6 +64,7 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()>
     child_module.add_class::<python::PersistBatch>()?;
     child_module.add_class::<python::PersistedBatch>()?;
     child_module.add_class::<python::WrittenChunk>()?;
+    child_module.add_class::<python::CurrentStateUpdate>()?;
     child_module.add_class::<effects::PersistEffects>()?;
     child_module.add_class::<reader::PyPersistReaderForTest>()?;
     child_module.add_function(wrap_pyfunction!(

rust/src/persist/python.rs

@@ -34,7 +34,8 @@ use super::{
     metrics::{PersistCounts, PersistMetrics},
     prepare::{self, PrepareInput, PrepareOptions, Prepared},
     reader::{self, PersistReader, PyDelta, PySlidingSync},
-    state_delta::NativeStateResolver,
+    sliding_sync,
+    state_delta::{self, NativeStateResolver, StateResFallback},
     txn::{self, json, streams::TxnStream, PartialStateConflict, TxnEvent, TxnInput},
 };
 use crate::{
@@ -252,6 +253,16 @@ impl Persister {
         .boxed()
     }
 
+    /// The future of [`Self::update_current_state`].
+    fn start_update_current_state(
+        &self,
+        room_id: String,
+    ) -> PyResult<BoxFuture<'static, PyResult<CurrentStateUpdate>>> {
+        let pool = self.inner.pool.get()?;
+        let inner = Arc::clone(&self.inner);
+        Ok(async move { update_current_state(&inner, &pool, &room_id).await }.boxed())
+    }
+
     /// Run `future` to completion on the runtime, releasing the GIL. For the
     /// `_blocking` entry points.
     fn block_on<T: Send>(
@@ -475,6 +486,30 @@ impl Persister {
         self.block_on(py, future)
     }
 
+    /// `_UpdateCurrentStateTask`: recalculate the room's current state
+    /// from its forward extremities and write the change, taking one
+    /// stream ID from the events stream. Resolves with a
+    /// [`CurrentStateUpdate`], whose effects Python must apply. Must not be
+    /// cancelled, as [`Self::persist_events_txn`].
+    fn update_current_state<'py>(
+        &self,
+        py: Python<'py>,
+        room_id: String,
+    ) -> PyResult<Bound<'py, PyAny>> {
+        let future = self.start_update_current_state(room_id)?;
+        create_deferred_with(py, &self.inner.runtime, future, OnCancel::RunOn)
+    }
+
+    /// As [`Self::update_current_state`], but blocks.
+    fn update_current_state_blocking(
+        &self,
+        py: Python<'_>,
+        room_id: String,
+    ) -> PyResult<CurrentStateUpdate> {
+        let future = self.start_update_current_state(room_id)?;
+        self.block_on(py, future)
+    }
+
     /// As [`Self::prepare_chunk`], but blocks.
     #[pyo3(signature = (room_id, events_and_contexts, backfilled, *, shadow))]
     fn prepare_chunk_blocking(
@@ -977,6 +1012,160 @@ impl WrittenChunk {
     }
 }
 
+/// `_update_current_state`: recalculate `room_id`'s current state and
+/// write the change, if there is one. Reads are retried as Python's
+/// `runInteraction`s are.
+async fn update_current_state(
+    inner: &Arc<PersisterInner>,
+    pool: &ConnectionPool,
+    room_id: &str,
+) -> PyResult<CurrentStateUpdate> {
+    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 mut attempt = 1;
+    let computed = loop {
+        let result = async {
+            let (state, stats) =
+                state_delta::calculate_current_state(&reader, &resolver, room_id).await?;
+            let delta = state_delta::calculate_state_delta(&reader, room_id, &state).await?;
+            let changes = sliding_sync::calculate_sliding_sync_table_changes(
+                &reader,
+                room_id,
+                &inner.server_name,
+                &[],
+                &delta,
+            )
+            .await?;
+            anyhow::Ok((delta, changes, stats))
+        }
+        .await;
+        match result {
+            Err(e) if attempt < MAX_ATTEMPTS && is_retryable_error(&e) => {
+                log::warn!("[TXN RETRY] {{update_current_state}} {e:#} {attempt}/{MAX_ATTEMPTS}");
+                attempt += 1;
+            }
+            result => break result,
+        }
+    };
+    let (delta, changes, stats) = match computed {
+        Ok(computed) => computed,
+        Err(e) => {
+            return match e.downcast::<StateResFallback>() {
+                Ok(fallback) => Ok(CurrentStateUpdate {
+                    effects: None,
+                    state_res_stats: None,
+                    fallback_reason: Some(fallback.0),
+                }),
+                Err(e) => Err(persist_error_to_py(e)),
+            };
+        }
+    };
+    // `DeltaState.is_noop`.
+    if delta.to_delete.is_empty() && delta.to_insert.is_empty() && !delta.no_longer_in_room {
+        return Ok(CurrentStateUpdate {
+            effects: None,
+            state_res_stats: stats,
+            fallback_reason: None,
+        });
+    }
+
+    // `get_next`, as late as can be.
+    let event_stream = TxnStream::events(Arc::clone(&inner.events_stream));
+    let ids = allocate_stream_orderings(pool, &event_stream, 1)
+        .await
+        .map_err(persist_error_to_py)?;
+    let stream_id = *ids
+        .first()
+        .ok_or_else(|| PyRuntimeError::new_err("no stream ID"))?;
+
+    let no_chunk = prepare::PreparedChunk {
+        state: Default::default(),
+        new_event_links: HashMap::new(),
+        sliding_sync: None,
+        sticky_events_to_un_soft_fail: HashSet::new(),
+    };
+    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: &[],
+        prepared: &no_chunk,
+        instance_name: &inner.instance_name,
+        server_name: &inner.server_name,
+        clock: inner.runtime.clock(),
+        inhibit_local_membership_updates: false,
+        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, "update_current_state", |txn, _| {
+        let (input, delta, changes) = (&input, &delta, &changes);
+        async move {
+            let mut effects = PersistEffects::default();
+            txn::update_current_state_txn(txn, input, delta, changes, stream_id, &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` finishes
+    // it when its block raises.
+    event_stream.hand_over(&ids);
+    effects.add_stream_ids(ids, false);
+    Ok(CurrentStateUpdate {
+        effects: Some(effects),
+        state_res_stats: stats,
+        fallback_reason: None,
+    })
+}
+
+/// What [`Persister::update_current_state`] did.
+#[pyclass(frozen)]
+pub struct CurrentStateUpdate {
+    effects: Option<PersistEffects>,
+    state_res_stats: Option<StateResStats>,
+    fallback_reason: Option<String>,
+}
+
+#[pymethods]
+impl CurrentStateUpdate {
+    /// What Python applies with `apply_rust_persist_effects`, with the
+    /// stream ID unfinished. `None` if the state did not change, or if
+    /// Rust fell back.
+    #[getter]
+    fn effects(&self) -> Option<PersistEffects> {
+        self.effects.clone()
+    }
+
+    /// The stats of the state resolution, if the state was resolved.
+    #[getter]
+    fn state_res_stats(&self) -> Option<StateResStats> {
+        self.state_res_stats.clone()
+    }
+
+    /// Why the Rust state resolution could not resolve the room, in which
+    /// case nothing was written and Python should recalculate the state.
+    #[getter]
+    fn fallback_reason(&self) -> Option<String> {
+        self.fallback_reason.clone()
+    }
+}
+
 /// Where the stream orderings of a chunk come from.
 #[derive(Clone, Copy)]
 enum StreamOrderings {

rust/src/persist/state_delta.rs

@@ -551,6 +551,66 @@ pub async fn get_new_state_after_events(
     })
 }
 
+/// `_calculate_current_state`, for `update_current_state`: the room's
+/// current state as its forward extremities imply, its state DAG
+/// extremities in an MSC4242 room. The state of their one state group, or
+/// the resolution of several, with the stats of the resolution. Fails with
+/// [`StateResFallback`] if the Rust state resolution cannot resolve the
+/// room.
+pub async fn calculate_current_state(
+    reads: &dyn StateDeltaReads,
+    resolver: &dyn StateResolver,
+    room_id: &str,
+) -> anyhow::Result<(StateMap, Option<StateResStats>)> {
+    let room_version_id = reads
+        .room_version_id(room_id)
+        .await?
+        .ok_or_else(|| anyhow!("Room {room_id} not found"))?;
+    let room_version: &'static RoomVersion = room_version_id
+        .parse()
+        .map_err(|e| anyhow!("{e}: {room_version_id:?}"))?;
+    let latest_event_ids = if room_version.msc4242_state_dags {
+        reads.state_dag_extremities(room_id).await?
+    } else {
+        reads.latest_event_ids_in_room(room_id).await?
+    };
+    let latest_event_ids: Vec<Arc<str>> = latest_event_ids.into_iter().collect();
+
+    let mut groups: Vec<i64> = reads
+        .state_group_for_events(&latest_event_ids)
+        .await?
+        .into_values()
+        .collect();
+    groups.sort_unstable();
+    groups.dedup();
+    let mut state_groups_map = reads.state_for_groups(&groups).await?;
+
+    if groups.len() == 1 {
+        let state = state_groups_map.remove(&groups[0]).unwrap_or_default();
+        return Ok(((*state).clone(), None));
+    }
+
+    if room_version.state_res == StateResolutionVersions::V1 {
+        return Err(v1_fallback(room_version));
+    }
+    info!("Resolving state for {room_id} with groups {groups:?}");
+    // As `resolve_state_groups` does.
+    let missing_groups = resolver
+        .check_state_groups_and_bump_deletion(&groups)
+        .await?;
+    if !missing_groups.is_empty() {
+        bail!("State groups have been deleted: {missing_groups:?}");
+    }
+    let state_sets: Vec<Arc<StateMap>> = groups
+        .iter()
+        .map(|group| state_groups_map.get(group).cloned().unwrap_or_default())
+        .collect();
+    let (state, stats) = resolver
+        .resolve(room_id, room_version, &state_sets, &EventMap::new())
+        .await?;
+    Ok((state, Some(stats)))
+}
+
 /// `_calculate_state_delta`: the change from the room's current state to
 /// `current_state`.
 pub async fn calculate_state_delta(
@@ -1601,4 +1661,67 @@ mod tests {
         assert_eq!(counts.state_delta, 0);
         assert_eq!(counts.state_dag_forward_extremities[&1], 1);
     }
+
+    fn current_state_reads(latest: &[(&str, i64)]) -> FakeReads {
+        FakeReads {
+            room_version: Some(Arc::from("12")),
+            latest: latest.iter().map(|(id, _)| Arc::from(*id)).collect(),
+            event_groups: latest
+                .iter()
+                .map(|(id, group)| (Arc::from(*id), *group))
+                .collect(),
+            groups: [
+                (5, Arc::new(state(&[("m.room.topic", "", "$a")]))),
+                (6, Arc::new(state(&[("m.room.topic", "", "$b")]))),
+            ]
+            .into(),
+            ..Default::default()
+        }
+    }
+
+    #[test]
+    fn current_state_of_one_group() {
+        // Two extremities in the same group need no resolution.
+        let reads = current_state_reads(&[("$x", 5), ("$y", 5)]);
+        let resolver = FakeResolver::default();
+        let (current, stats) = block_on(calculate_current_state(&reads, &resolver, ROOM)).unwrap();
+        assert_eq!(current, state(&[("m.room.topic", "", "$a")]));
+        assert!(stats.is_none());
+        assert!(resolver.resolved.lock().unwrap().is_empty());
+    }
+
+    #[test]
+    fn current_state_of_several_groups_is_resolved() {
+        let reads = current_state_reads(&[("$x", 5), ("$y", 6)]);
+        let resolver = FakeResolver {
+            result: state(&[("m.room.topic", "", "$b")]),
+            ..Default::default()
+        };
+        let (current, stats) = block_on(calculate_current_state(&reads, &resolver, ROOM)).unwrap();
+        assert_eq!(current, state(&[("m.room.topic", "", "$b")]));
+        assert!(stats.is_some());
+        assert_eq!(*resolver.checked.lock().unwrap(), vec![vec![5, 6]]);
+        assert_eq!(*resolver.resolved.lock().unwrap(), vec![("12", 2)]);
+    }
+
+    #[test]
+    fn current_state_falls_back_on_v1_and_fails_on_deleted_groups() {
+        let mut reads = current_state_reads(&[("$x", 5), ("$y", 6)]);
+        reads.room_version = Some(Arc::from("1"));
+        let err = block_on(calculate_current_state(
+            &reads,
+            &FakeResolver::default(),
+            ROOM,
+        ))
+        .unwrap_err();
+        assert!(err.downcast_ref::<StateResFallback>().is_some(), "{err:#}");
+
+        let reads = current_state_reads(&[("$x", 5), ("$y", 6)]);
+        let resolver = FakeResolver {
+            deleted_groups: vec![6],
+            ..Default::default()
+        };
+        let err = block_on(calculate_current_state(&reads, &resolver, ROOM)).unwrap_err();
+        assert!(err.to_string().contains("deleted"), "{err:#}");
+    }
 }

rust/src/persist/txn/mod.rs

@@ -51,6 +51,8 @@ use super::{
     context::PersistContext,
     effects::{EventInvalidation, PersistEffects},
     prepare::PreparedChunk,
+    sliding_sync::SlidingSyncTableChanges,
+    state_delta::DeltaState,
 };
 use crate::{
     clock::Clock,
@@ -205,6 +207,31 @@ pub async fn persist_events_txn(
     result
 }
 
+/// `_update_current_state_txn`, for `update_current_state`: write the
+/// change `delta` to the room's current state, with `stream_id` from the
+/// events stream, adding to `effects` what Python would queue. `input`
+/// has no events. As [`persist_events_txn`], `effects` must start empty.
+pub async fn update_current_state_txn(
+    txn: &WriteTxn,
+    input: &TxnInput<'_>,
+    delta: &DeltaState,
+    changes: &SlidingSyncTableChanges,
+    stream_id: i64,
+    effects: &mut PersistEffects,
+) -> Result<(), WriteError> {
+    for stream in input.streams() {
+        stream.finish_pending();
+    }
+    let result =
+        current_state::update_current_state(txn, input, delta, changes, stream_id, effects).await;
+    if result.is_err() {
+        for stream in input.streams() {
+            stream.finish_pending();
+        }
+    }
+    result
+}
+
 async fn persist_events_txn_steps<'a>(
     txn: &WriteTxn,
     input: &TxnInput<'a>,

synapse/config/experimental.py

@@ -415,9 +415,8 @@ 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 at "full",
-        # Python still applies what each chunk Rust writes asks for and runs
-        # what follows it, and recalculates current state on request.
+        # At "full", Python still runs the per-room queue and applies what
+        # each chunk Rust writes asks for to its caches.
         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

@@ -796,6 +796,13 @@ class EventsPersistenceStorageController:
 
         Recalculates the current state for a room, and persists it.
         """
+        if (
+            self._rust_persister is not None
+            and self.hs.config.experimental.rust_event_persistence_at_least("full")
+            and await self._update_current_state_in_rust(room_id)
+        ):
+            return
+
         state = await self._calculate_current_state(room_id)
         delta = await self._calculate_state_delta(room_id, state)
         sliding_sync_table_changes = (
@@ -808,6 +815,50 @@ class EventsPersistenceStorageController:
             room_id, delta, sliding_sync_table_changes
         )
 
+    async def _update_current_state_in_rust(self, room_id: str) -> bool:
+        """`_update_current_state` in Rust, at `rust_event_persistence:
+        full`. False if Rust cannot resolve the room's state, and wrote
+        nothing."""
+        assert self._rust_persister is not None
+        persister = self._rust_persister
+
+        async def update_and_apply() -> bool:
+            # The Deferred follows the logcontext rules already.
+            try:
+                update = await persister.update_current_state(room_id)
+            except BaseException:
+                # Rust finished the stream ID, if it took one. Tell
+                # replication, as `get_next` does when its block raises.
+                self.hs.get_replication_notifier().notify_replication()
+                raise
+            # Recorded first, as Python records them before it writes.
+            if update.state_res_stats is not None:
+                self._state_resolution_handler.record_rust_state_res_metrics(
+                    room_id, update.state_res_stats
+                )
+            effects = update.effects
+            if effects is not None:
+                store = self.persist_events_store.store
+                try:
+                    apply_rust_persist_effects(store, effects, finish_stream_ids=False)
+                except BaseException:
+                    # Whatever failed, the stream ID must not stall the
+                    # stream.
+                    finish_rust_persist_stream_ids(store, effects, succeeded=False)
+                    raise
+                finish_rust_persist_stream_ids(store, effects, succeeded=True)
+            if update.fallback_reason is not None:
+                logger.info(
+                    "Recalculating the current state of %s in Python: %s",
+                    room_id,
+                    update.fallback_reason,
+                )
+                return False
+            return True
+
+        # The effects of a write that commits must be applied.
+        return await delay_cancellation(update_and_apply())
+
     async def _calculate_current_state(self, room_id: str) -> StateMap[str]:
         """Calculate the current state of a room, based on the forward extremities
 

synapse/synapse_rust/persist.pyi

@@ -140,6 +140,30 @@ class Persister:
     def batch_abort_blocking(self, batch: PersistBatch) -> None:
         """As `batch_abort`, but blocks. For tests."""
 
+    def update_current_state(self, room_id: str) -> Deferred[CurrentStateUpdate]:
+        """`_UpdateCurrentStateTask`: recalculate the room's current state
+        from its forward extremities and write the change, taking one
+        stream ID from the events stream. Must not be cancelled."""
+    def update_current_state_blocking(self, room_id: str) -> CurrentStateUpdate:
+        """As `update_current_state`, but blocks. For tests."""
+
+class CurrentStateUpdate:
+    """What `Persister.update_current_state` did."""
+
+    @property
+    def effects(self) -> PersistEffects | None:
+        """To apply with `apply_rust_persist_effects`, with the stream ID
+        unfinished. `None` if the state did not change, or if Rust fell
+        back."""
+    @property
+    def state_res_stats(self) -> StateResStats | None:
+        """The stats of the state resolution, if the state was resolved."""
+    @property
+    def fallback_reason(self) -> str | None:
+        """Why the Rust state resolution could not resolve the room, in
+        which case nothing was written and Python should recalculate the
+        state."""
+
 class PersistBatch:
     """A batch being persisted. See `Persister.persist_batch`."""
 

tests/storage/test_rust_persist_prereqs.py

@@ -1545,6 +1545,223 @@ class RustFullTestCase(unittest.HomeserverTestCase):
         self.get_success(defer.ensureDeferred(persist()))
         self.assertEqual(seen, [True])
 
+    def _current_topic(self) -> list[str]:
+        return self.get_success(
+            self.store.db_pool.simple_select_onecol(
+                "current_state_events",
+                {"room_id": self.room_id, "type": EventTypes.Topic},
+                "event_id",
+            )
+        )
+
+    def test_update_current_state(self) -> None:
+        """Rust recalculates the current state from the extremities, writes
+        the change with a stream ID it takes, and Python applies it."""
+        topic = self.helper.send_state(
+            self.room_id, EventTypes.Topic, {"topic": "hi"}, tok=self.token
+        )["event_id"]
+        self.get_success(
+            self.store.db_pool.simple_delete(
+                "current_state_events",
+                {"room_id": self.room_id, "type": EventTypes.Topic},
+                desc="test",
+            )
+        )
+        before = self.store._stream_id_gen.get_current_token()
+        with patch.object(
+            self.persistence,
+            "_calculate_current_state",
+            side_effect=AssertionError("calculated in Python"),
+        ):
+            self.get_success(self.persistence.update_current_state(self.room_id))
+        self.assertEqual(self._current_topic(), [topic])
+        after = self.store._stream_id_gen.get_current_token()
+        self.assertEqual(after, before + 1)
+        # The delta row is at the stream ID taken.
+        rows = self.get_success(
+            self.store.db_pool.simple_select_list(
+                "current_state_delta_stream",
+                {"room_id": self.room_id, "stream_id": after},
+                ("type", "event_id", "prev_event_id"),
+            )
+        )
+        # No previous event, as its row was gone.
+        self.assertEqual(rows, [(EventTypes.Topic, topic, None)])
+
+    def test_update_current_state_unchanged(self) -> None:
+        """A room whose current state is right takes no stream ID."""
+        before = self.store._stream_id_gen.get_current_token()
+        with patch.object(
+            self.persistence,
+            "_calculate_current_state",
+            side_effect=AssertionError("calculated in Python"),
+        ):
+            self.get_success(self.persistence.update_current_state(self.room_id))
+        self.assertEqual(self.store._stream_id_gen.get_current_token(), before)
+
+    def test_update_current_state_falls_back(self) -> None:
+        """A room on state resolution v1 whose extremities need resolving
+        is recalculated in Python."""
+        self.room_id = self.helper.create_room_as(
+            self.user, room_version=RoomVersions.V1.identifier, tok=self.token
+        )
+        # Two topics on the same extremity, so the room forks.
+        pairs = []
+        for topic in ("A", "B"):
+            event, unpersisted = self.get_success(
+                self.hs.get_event_creation_handler().create_event(
+                    create_requester(self.user),
+                    {
+                        "type": EventTypes.Topic,
+                        "state_key": "",
+                        "content": {"topic": topic},
+                        "room_id": self.room_id,
+                        "sender": self.user,
+                    },
+                )
+            )
+            pairs.append((event, self.get_success(unpersisted.persist(event))))
+        for pair in pairs:
+            self.get_success(self.persistence.persist_events([pair]))
+        calculate = self.persistence._calculate_current_state
+        id_gen = self.store._stream_id_gen
+        before = id_gen.get_current_token()
+        tokens = []
+
+        async def in_python(room_id: str) -> Any:
+            # Rust took no stream ID before handing back.
+            tokens.append(id_gen.get_current_token())
+            return await calculate(room_id)
+
+        with patch.object(
+            self.persistence, "_calculate_current_state", side_effect=in_python
+        ):
+            self.get_success(self.persistence.update_current_state(self.room_id))
+        self.assertEqual(tokens, [before])
+
+    def _fork_topics(self) -> None:
+        """Two topics on the same extremity, persisted one after the other,
+        so the room has two extremities in two state groups."""
+        pairs = []
+        for topic in ("A", "B"):
+            event, unpersisted = self.get_success(
+                self.hs.get_event_creation_handler().create_event(
+                    create_requester(self.user),
+                    {
+                        "type": EventTypes.Topic,
+                        "state_key": "",
+                        "content": {"topic": topic},
+                        "room_id": self.room_id,
+                        "sender": self.user,
+                    },
+                )
+            )
+            pairs.append((event, self.get_success(unpersisted.persist(event))))
+        for pair in pairs:
+            self.get_success(self.persistence.persist_events([pair]))
+
+    def _delete_current_state(self, event_type: str, state_key: str) -> None:
+        self.get_success(
+            self.store.db_pool.simple_delete(
+                "current_state_events",
+                {"room_id": self.room_id, "type": event_type, "state_key": state_key},
+                desc="test",
+            )
+        )
+
+    def test_update_current_state_resolves(self) -> None:
+        """Extremities in several state groups are resolved in Rust, to the
+        state Python resolves to, and the resolution is recorded."""
+        self._fork_topics()
+        expected = self.get_success(
+            self.persistence._calculate_current_state(self.room_id)
+        )
+        self._delete_current_state(EventTypes.Topic, "")
+        handler = self.hs.get_state_resolution_handler()
+        handler._state_res_metrics.clear()
+        with patch.object(
+            self.persistence,
+            "_calculate_current_state",
+            side_effect=AssertionError("calculated in Python"),
+        ):
+            self.get_success(self.persistence.update_current_state(self.room_id))
+        self.assertEqual(self._current_topic(), [expected[(EventTypes.Topic, "")]])
+        self.assertIn(self.room_id, handler._state_res_metrics)
+
+    def test_update_current_state_membership(self) -> None:
+        """A membership the update restores moves the membership tables on
+        and streams the state caches' invalidation."""
+        bob = self.register_user("bob", "pass")
+        bob_token = self.login("bob", "pass")
+        self.helper.join(self.room_id, bob, tok=bob_token)
+        self._delete_current_state(EventTypes.Member, bob)
+        self.get_success(
+            self.store.db_pool.simple_delete(
+                "sliding_sync_membership_snapshots",
+                {"room_id": self.room_id, "user_id": bob},
+                desc="test",
+            )
+        )
+        cache_position = self.store._cache_id_gen.get_current_token()  # type: ignore[union-attr]
+        self.get_success(self.persistence.update_current_state(self.room_id))
+
+        (membership,) = self.get_success(
+            self.store.db_pool.simple_select_onecol(
+                "current_state_events",
+                {"room_id": self.room_id, "type": EventTypes.Member, "state_key": bob},
+                "membership",
+            )
+        )
+        self.assertEqual(membership, "join")
+        (snapshot,) = self.get_success(
+            self.store.db_pool.simple_select_onecol(
+                "sliding_sync_membership_snapshots",
+                {"room_id": self.room_id, "user_id": bob},
+                "membership",
+            )
+        )
+        self.assertEqual(snapshot, "join")
+        cache_rows = self.get_success(
+            self.store.db_pool.execute(
+                "test",
+                "SELECT cache_func FROM cache_invalidation_stream_by_instance"
+                " WHERE stream_id > ?",
+                cache_position,
+            )
+        )
+        self.assertIn(("cs_cache_fake",), cache_rows)
+
+    def test_update_current_state_apply_failure(self) -> None:
+        """If applying the update's effects fails, its stream ID is still
+        finished."""
+
+        class ApplyFailed(Exception):
+            pass
+
+        self._delete_current_state(EventTypes.Topic, "")
+        self.helper.send_state(
+            self.room_id, EventTypes.Topic, {"topic": "hi"}, tok=self.token
+        )
+        self._delete_current_state(EventTypes.Topic, "")
+        with patch(
+            "synapse.storage.controllers.persist_events.apply_rust_persist_effects",
+            side_effect=ApplyFailed(),
+        ):
+            self.get_failure(
+                self.persistence.update_current_state(self.room_id), ApplyFailed
+            )
+        (stream_id,) = self.get_success(
+            self.store.db_pool.execute(
+                "test",
+                "SELECT MAX(stream_id) FROM current_state_delta_stream"
+                " WHERE room_id = ?",
+                self.room_id,
+            )
+        )[0]
+        self.assertGreaterEqual(
+            self.store._stream_id_gen.get_current_token(), stream_id
+        )
+
     def test_gil_count(self) -> None:
         """Taking the stream orderings and setting them on the events takes
         the GIL at no persist site."""