Read no state groups when persisting a batch of new events

After `batch_persist_unpersisted_contexts`, an event's `EventContext` only
had state group IDs, so every read of the state before the event went to
the state store. For a new room that was eight state-group reads per
`/createRoom`: six from `deduplicate_state_event`, one from the
create-event check in `persist_and_notify_client_events`, and one from the
persister working out the new current state.

Keep the state map the batch was built from on the `EventContext`, as a
private attribute that `get_prev_state_ids` prefers and that is not
serialized, so a context sent to another worker reads the state group as
before. A context that arrives with a state group of its own, which the
batch then replaces, drops its map. `store_state_deltas_for_batched` also prefills the state group
caches with the state after the last event of the batch, through the
helper `store_state_group` now shares.

On the trial harness (SQLite) a plain room drops from 35 to 27 database
interactions, and `/createRoom` from 17 to 10 database transactions.

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

synapse/events/snapshot.py

@@ -130,6 +130,12 @@ class EventContext(UnpersistedEventContextBase):
 
         partial_state: if True, we may be storing this event with a temporary,
             incomplete state.
+
+        _state_map_before_event: The state before the event, i.e. the state at
+            ``state_group_before_event``, if it was known when the context was
+            built. ``get_prev_state_ids`` returns it instead of reading the state
+            group. It is not serialized, so a context passed to another worker
+            reads the state group.
     """
 
     _storage: "StorageControllers"
@@ -141,6 +147,7 @@ class EventContext(UnpersistedEventContextBase):
     app_service: ApplicationService | None = None
 
     partial_state: bool = False
+    _state_map_before_event: StateMap[str] | None = None
 
     @staticmethod
     def with_state(
@@ -301,6 +308,10 @@ class EventContext(UnpersistedEventContextBase):
 
         if self.state_group_before_event is None:
             return {}
+        if self._state_map_before_event is not None:
+            if state_filter is None:
+                return self._state_map_before_event
+            return state_filter.filter_state(self._state_map_before_event)
         return await self._storage.state.get_state_ids_for_group(
             self.state_group_before_event, state_filter
         )
@@ -398,6 +409,7 @@ class UnpersistedEventContext(UnpersistedEventContextBase):
                 state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
                 partial_state=unpersisted_context.partial_state,
                 state_group_deltas=state_group_deltas,
+                state_map_before_event=unpersisted_context.state_map_before_event,
             )
             events_and_persisted_context.append((event, context))
         return events_and_persisted_context

synapse/storage/databases/state/store.py

@@ -443,6 +443,29 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
                 fetched_keys=non_member_types,
             )
 
+    def _prefill_state_group_caches_txn(
+        self, txn: LoggingTransaction, state_group: int, state: StateMap[str]
+    ) -> None:
+        """Fill the state group caches with the full state of a new group once
+        `txn` commits.
+
+        It's fine to use the sequence like this as the state group map is
+        immutable. (If the map wasn't immutable then this prefill could race with
+        another update.)
+        """
+        txn.call_after(
+            self._state_group_members_cache.update,
+            self._state_group_members_cache.sequence,
+            key=state_group,
+            value={s: ev for (s, ev) in state.items() if s[0] == EventTypes.Member},
+        )
+        txn.call_after(
+            self._state_group_cache.update,
+            self._state_group_cache.sequence,
+            key=state_group,
+            value={s: ev for (s, ev) in state.items() if s[0] != EventTypes.Member},
+        )
+
     @trace
     @tag_args
     async def store_state_deltas_for_batched(
@@ -529,6 +552,11 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
 
             state_group_iter = iter(state_groups)
             for event, context in events_and_context:
+                if context.state_group_before_event not in (None, sg_before):
+                    # The context came with a state group of its own, which the
+                    # chain replaces, so its state map may not match the new one.
+                    context.state_map_before_event = None
+
                 if not event.is_state():
                     context.state_group_after_event = sg_before
                     context.state_group_before_event = sg_before
@@ -587,6 +615,21 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
                 ],
             )
 
+            # Prefill the state group caches with the state after the last
+            # event, which persisting the batch reads to work out the room's
+            # new current state.
+            last_event, last_context = events_and_context[-1]
+            if last_context.state_map_before_event is not None:
+                last_state = dict(last_context.state_map_before_event)
+                if last_event.is_state():
+                    last_state[(last_event.type, last_event.state_key)] = (
+                        last_event.event_id
+                    )
+                assert last_context.state_group_after_event is not None
+                self._prefill_state_group_caches_txn(
+                    txn, last_context.state_group_after_event, last_state
+                )
+
             return events_and_context
 
         return await self.db_pool.runInteraction(
@@ -710,34 +753,7 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
                 ],
             )
 
-            # Prefill the state group caches with this group.
-            # It's fine to use the sequence like this as the state group map
-            # is immutable. (If the map wasn't immutable then this prefill could
-            # race with another update)
-
-            current_member_state_ids = {
-                s: ev
-                for (s, ev) in current_state_ids.items()
-                if s[0] == EventTypes.Member
-            }
-            txn.call_after(
-                self._state_group_members_cache.update,
-                self._state_group_members_cache.sequence,
-                key=state_group,
-                value=current_member_state_ids,
-            )
-
-            current_non_member_state_ids = {
-                s: ev
-                for (s, ev) in current_state_ids.items()
-                if s[0] != EventTypes.Member
-            }
-            txn.call_after(
-                self._state_group_cache.update,
-                self._state_group_cache.sequence,
-                key=state_group,
-                value=current_non_member_state_ids,
-            )
+            self._prefill_state_group_caches_txn(txn, state_group, current_state_ids)
 
             return state_group
 

tests/handlers/test_room.py

@@ -465,3 +465,19 @@ class CreateRoomPersistBatchesTestCase(unittest.HomeserverTestCase):
                 )
             )
             self.assertEqual(len(groups), len(state) + 1)
+
+    def test_no_state_group_reads(self) -> None:
+        """Creating a room reads no state groups from the database: the batch
+        knows the state before each event and after the last one."""
+        state_store = self.hs.get_datastores().state
+        reads: list[object] = []
+        orig = state_store._get_state_groups_from_groups
+
+        async def record(*args: Any, **kwargs: Any) -> Any:
+            reads.append(args)
+            return await orig(*args, **kwargs)
+
+        for room_version in ("10", "12"):
+            with patch.object(state_store, "_get_state_groups_from_groups", record):
+                self._create_room(room_version)
+            self.assertEqual(reads, [])

tests/rest/client/test_rooms.py

@@ -799,7 +799,7 @@ class RoomsCreateTestCase(RoomBase):
         self.assertEqual(HTTPStatus.OK, channel.code, channel.result)
         self.assertTrue("room_id" in channel.json_body)
         assert channel.resource_usage is not None
-        self.assertEqual(17, channel.resource_usage.db_txn_count)
+        self.assertEqual(10, channel.resource_usage.db_txn_count)
 
     def test_post_room_initial_state(self) -> None:
         # POST with initial_state config key, expect new room id
@@ -812,7 +812,7 @@ class RoomsCreateTestCase(RoomBase):
         self.assertEqual(HTTPStatus.OK, channel.code, channel.result)
         self.assertTrue("room_id" in channel.json_body)
         assert channel.resource_usage is not None
-        self.assertEqual(18, channel.resource_usage.db_txn_count)
+        self.assertEqual(10, channel.resource_usage.db_txn_count)
 
     def test_post_room_topic(self) -> None:
         # POST with topic key, expect new room id