misc cleanup

4 files changed, 20 insertions, 14 deletions

diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index fcfdaa074f..1f2908f9b0 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -580,7 +580,7 @@ class EventCreationHandler:
         the event using the parameter state_map, thus this parameter must be provided
         if for_batch is True. The subsequently created event and context are suitable
         for being batched up and bulk persisted to the database with other similarly
-        created events. todo: must batch state groups
+        created events. todo: remind must batch state groups if for_batch?
 
         Creates an FrozenEvent object, filling out auth_events, prev_events,
         etc.
@@ -1066,9 +1066,9 @@ class EventCreationHandler:
         for_batch: bool = False,
     ) -> Tuple[EventBase, EventContext]:
         """Create a new event for a local client. If bool for_batch is true, will
-        create an event using the prev_event_ids, and will create an event context todo: note that event context has not state group?
-        for
-        the event using the parameters state_map, thus this parameter if for_batch is True.
+        create an event using the prev_event_ids, and will create an event context
+        (todo: note that event context has not state group if batched?)
+        for the event using the parameters state_map, thus this parameter if for_batch is True.
         The subsequently created event and context are suitable for being batched up
         and bulk persisted to the database with other similarly created events.
 
@@ -1167,7 +1167,9 @@ class EventCreationHandler:
             event = await builder.build(
                 prev_event_ids=prev_event_ids, auth_event_ids=auth_ids, depth=depth
             )
-            context = EventContext(self._storage_controllers)
+            context = await self.state.compute_event_context_for_batched(
+                event, state_map
+            )
         else:
             event = await builder.build(
                 prev_event_ids=prev_event_ids,
diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index 168012a9c6..9cd9d564e5 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -1134,7 +1134,11 @@ class RoomCreationHandler:
         # through a different code path
         depth += 1
         state_map[(EventTypes.Member, creator.user.to_string())] = member_event_id
-        event_to_state_group = await self._storage_controllers.state.get_state_group_for_events([member_event_id])
+        event_to_state_group = (
+            await self._storage_controllers.state.get_state_group_for_events(
+                [member_event_id]
+            )
+        )
         current_state_group = event_to_state_group[member_event_id]
 
         events_to_send = []
@@ -1243,15 +1247,10 @@ class RoomCreationHandler:
 
         assert self.hs.datastores is not None
         assert current_state_group is not None
-        state_groups = await self.hs.datastores.state.store_state_deltas_for_batched(
+        await self.hs.datastores.state.store_state_deltas_for_batched(
             events_to_send, room_id, prev_group=current_state_group
         )
 
-        index = 0
-        for _, context in events_to_send:
-            context._state_group = state_groups[index]
-            index += 1
-
         last_event = await self.event_creation_handler.handle_new_client_event(
             creator,
             events_to_send,
diff --git a/synapse/storage/databases/state/store.py b/synapse/storage/databases/state/store.py
index 97474bc60d..b554d80d7b 100644
--- a/synapse/storage/databases/state/store.py
+++ b/synapse/storage/databases/state/store.py
@@ -520,8 +520,10 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
         current_state_ids: Optional[StateMap[str]],
     ) -> int:
         """Store a new set of state, returning a newly assigned state group.
+
         At least one of `current_state_ids` and `prev_group` must be provided. Whenever
         `prev_group` is not None, `delta_ids` must also not be None.
+
         Args:
             event_id: The event ID for which the state was calculated
             room_id
@@ -531,6 +533,7 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
                 `current_state_ids`.
             current_state_ids: The state to store. Map of (type, state_key)
                 to event_id.
+
         Returns:
             The state group ID
         """
@@ -545,7 +548,9 @@ class StateGroupDataStore(StateBackgroundUpdateStore, SQLBaseStore):
             txn: LoggingTransaction, prev_group: int, delta_ids: StateMap[str]
         ) -> Optional[int]:
             """Try and persist the new group as a delta.
+
             Requires that we have the state as a delta from a previous state group.
+
             Returns:
                 The state group if successfully created, or None if the state
                 needs to be persisted as a full state.
diff --git a/tests/rest/client/test_rooms.py b/tests/rest/client/test_rooms.py
index 6042cc714b..d45f9f5e43 100644
--- a/tests/rest/client/test_rooms.py
+++ b/tests/rest/client/test_rooms.py
@@ -715,7 +715,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(32, channel.resource_usage.db_txn_count)
+        self.assertEqual(30, channel.resource_usage.db_txn_count)
 
     def test_post_room_initial_state(self) -> None:
         # POST with initial_state config key, expect new room id
@@ -728,7 +728,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(34, channel.resource_usage.db_txn_count)
+        self.assertEqual(30, channel.resource_usage.db_txn_count)
 
     def test_post_room_visibility_key(self) -> None:
         # POST with visibility config key, expect new room id