diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py index fbe8f3377a..ffd4d272b9 100644 --- a/synapse/events/snapshot.py +++ b/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 diff --git a/synapse/storage/databases/state/store.py b/synapse/storage/databases/state/store.py index 847892c009..79f4ab0097 100644 --- a/synapse/storage/databases/state/store.py +++ b/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 diff --git a/tests/handlers/test_room.py b/tests/handlers/test_room.py index 7cd97f9aab..769eac7334 100644 --- a/tests/handlers/test_room.py +++ b/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, []) diff --git a/tests/rest/client/test_rooms.py b/tests/rest/client/test_rooms.py index 6499021f33..0741cc8058 100644 --- a/tests/rest/client/test_rooms.py +++ b/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