Draft: Optimize MSC2716 /batch_send endpoint v1

scripts-dev/complement.sh

@@ -71,4 +71,4 @@ fi
 
 # Run the tests!
 echo "Images built; running complement"
-go test -v -tags synapse_blacklist,msc2403 -count=1 $EXTRA_COMPLEMENT_ARGS ./tests/...
+go test -v -tags synapse_blacklist,msc2403,msc2716 -count=1 $EXTRA_COMPLEMENT_ARGS ./tests/

synapse/handlers/message.py

@@ -884,6 +884,30 @@ async def create_and_send_nonmember_event(
         assert ev.internal_metadata.stream_ordering
         return ev, ev.internal_metadata.stream_ordering
 
+    async def strip_auth_event_ids_for_given_event_builder(
+        self,
+        builder: EventBuilder,
+        prev_event_ids: List[str],
+        auth_event_ids: List[str],
+        depth: Optional[int] = None,
+    ) -> List[str]:
+        temp_event = await builder.build(
+            prev_event_ids=prev_event_ids,
+            auth_event_ids=auth_event_ids,
+            depth=depth,
+        )
+        auth_events = await self.store.get_events_as_list(auth_event_ids)
+        # Create a StateMap[str]
+        auth_event_state_map = {(e.type, e.state_key): e.event_id for e in auth_events}
+        # Actually strip down and use the necessary auth events
+        stripped_auth_event_ids = self._event_auth_handler.compute_auth_events(
+            event=temp_event,
+            current_state_ids=auth_event_state_map,
+            for_verification=False,
+        )
+
+        return stripped_auth_event_ids
+
     @measure_func("create_new_client_event")
     async def create_new_client_event(
         self,
@@ -925,29 +949,19 @@ async def create_new_client_event(
         # For example, we don't need extra m.room.member that don't match event.sender
         full_state_ids_at_event = None
         if auth_event_ids is not None:
-            # If auth events are provided, prev events must be also.
+            # If auth events are provided, prev events must also be provided.
             # prev_event_ids could be an empty array though.
             assert prev_event_ids is not None
 
             # Copy the full auth state before it stripped down
             full_state_ids_at_event = auth_event_ids.copy()
 
-            temp_event = await builder.build(
+            auth_event_ids = await self.strip_auth_event_ids_for_given_event_builder(
+                builder=builder,
                 prev_event_ids=prev_event_ids,
                 auth_event_ids=auth_event_ids,
                 depth=depth,
             )
-            auth_events = await self.store.get_events_as_list(auth_event_ids)
-            # Create a StateMap[str]
-            auth_event_state_map = {
-                (e.type, e.state_key): e.event_id for e in auth_events
-            }
-            # Actually strip down and use the necessary auth events
-            auth_event_ids = self._event_auth_handler.compute_auth_events(
-                event=temp_event,
-                current_state_ids=auth_event_state_map,
-                for_verification=False,
-            )
 
         if prev_event_ids is not None:
             assert (
@@ -991,6 +1005,8 @@ async def create_new_client_event(
             and full_state_ids_at_event
             and builder.internal_metadata.is_historical()
         ):
+            # Add explicit state to the insertion event so the rest of the batch
+            # can inherit the same state and `state_group`
             old_state = await self.store.get_events_as_list(full_state_ids_at_event)
             context = await self.state.compute_event_context(event, old_state=old_state)
         else:

synapse/handlers/room_batch.py

@@ -1,8 +1,11 @@
 import logging
-from typing import TYPE_CHECKING, List, Tuple
+from typing import TYPE_CHECKING, Dict, List, NamedTuple, Tuple
 
 from synapse.api.constants import EventContentFields, EventTypes
 from synapse.appservice import ApplicationService
+from synapse.events import EventBase
+from synapse.events.snapshot import EventContext
+from synapse.events.validator import EventValidator
 from synapse.http.servlet import assert_params_in_dict
 from synapse.types import JsonDict, Requester, UserID, create_requester
 from synapse.util.stringutils import random_string
@@ -16,10 +19,13 @@
 class RoomBatchHandler:
     def __init__(self, hs: "HomeServer"):
         self.hs = hs
+        self.config = hs.config
         self.store = hs.get_datastore()
         self.state_store = hs.get_storage().state
         self.event_creation_handler = hs.get_event_creation_handler()
         self.room_member_handler = hs.get_room_member_handler()
+        self.validator = EventValidator()
+        self.event_builder_factory = hs.get_event_builder_factory()
         self.auth = hs.get_auth()
 
     async def inherit_depth_from_prev_ids(self, prev_event_ids: List[str]) -> int:
@@ -290,6 +296,18 @@ async def persist_historical_events(
         """
         assert app_service_requester.app_service
 
+        room_version_obj = await self.store.get_room_version(room_id)
+
+        # Map from event type to EventContext that we can re-use and create
+        # another event in the batch with the same type
+        event_type_to_context_cache: Dict[str, EventContext] = {}
+
+        # Map from (event.sender, event.type) to auth_event_ids that we can re-use and create
+        # another event in the batch with the same sender and type
+        event_sender_and_type_to_auth_event_ids_cache: Dict[
+            Tuple(str, str), List[str]
+        ] = {}
+
         # Make the historical event chain float off on its own by specifying no
         # prev_events for the first event in the chain which causes the HS to
         # ask for the state at the start of the batch later.
@@ -309,38 +327,93 @@ async def persist_historical_events(
                 "origin_server_ts": ev["origin_server_ts"],
                 "content": ev["content"],
                 "room_id": room_id,
-                "sender": ev["sender"],  # requester.user.to_string(),
+                "sender": ev["sender"],
                 "prev_events": prev_event_ids.copy(),
             }
 
             # Mark all events as historical
             event_dict["content"][EventContentFields.MSC2716_HISTORICAL] = True
 
-            event, context = await self.event_creation_handler.create_event(
-                await self.create_requester_for_user_id_from_app_service(
-                    ev["sender"], app_service_requester.app_service
-                ),
-                event_dict,
-                # Only the first event in the chain should be floating.
-                # The rest should hang off each other in a chain.
-                allow_no_prev_events=index == 0,
-                prev_event_ids=event_dict.get("prev_events"),
-                auth_event_ids=auth_event_ids,
-                historical=True,
-                depth=inherited_depth,
-            )
+            # We can skip a bunch of context and state calculations if we
+            # already have an event with the same type to base off of.
+            cached_context = event_type_to_context_cache.get(ev["type"])
 
-            assert context._state_group
+            if cached_context is None:
+                event, context = await self.event_creation_handler.create_event(
+                    await self.create_requester_for_user_id_from_app_service(
+                        ev["sender"], app_service_requester.app_service
+                    ),
+                    event_dict,
+                    # Only the first event in the chain should be floating.
+                    # The rest should hang off each other in a chain.
+                    allow_no_prev_events=index == 0,
+                    prev_event_ids=event_dict.get("prev_events"),
+                    auth_event_ids=auth_event_ids,
+                    historical=True,
+                    depth=inherited_depth,
+                )
 
-            # Normally this is done when persisting the event but we have to
-            # pre-emptively do it here because we create all the events first,
-            # then persist them in another pass below. And we want to share
-            # state_groups across the whole batch so this lookup needs to work
-            # for the next event in the batch in this loop.
-            await self.store.store_state_group_id_for_event_id(
-                event_id=event.event_id,
-                state_group_id=context._state_group,
-            )
+                # Cache the context so we can re-use it for events in
+                # the batch that have the same type.
+                event_type_to_context_cache[event.type] = context
+                # Cache the auth_event_ids so we can re-use it for events in
+                # the batch that have the same sender and type.
+                event_sender_and_type_to_auth_event_ids_cache[
+                    (event.sender, event.type)
+                ] = event.auth_event_ids()
+
+                # Normally this is done when persisting the event but we have to
+                # pre-emptively do it here because we create all the events first,
+                # then persist them in another pass below. And we want to share
+                # state_groups across the whole batch so this lookup needs to work
+                # for the next event in the batch in this loop.
+                await self.store.store_state_group_id_for_event_id(
+                    event_id=event.event_id,
+                    state_group_id=context._state_group,
+                )
+            else:
+                # Create an event with a lot less overhead than create_event
+                builder = self.event_builder_factory.for_room_version(
+                    room_version_obj, event_dict
+                )
+                builder.internal_metadata.historical = True
+
+                # TODO: Can we get away without this? Can't we just rely on validate_new below?
+                # self.validator.validate_builder(builder)
+
+                # TODO: Do we need to sanity check the number of prev_events?
+
+                resultant_auth_event_ids = (
+                    event_sender_and_type_to_auth_event_ids_cache.get(
+                        (ev["sender"], ev["type"])
+                    )
+                )
+                if resultant_auth_event_ids is None:
+                    resultant_auth_event_ids = await self.event_creation_handler.strip_auth_event_ids_for_given_event_builder(
+                        builder=builder,
+                        prev_event_ids=prev_event_ids,
+                        auth_event_ids=auth_event_ids,
+                        depth=inherited_depth,
+                    )
+                    # Cache the auth_event_ids so we can re-use it for events in
+                    # the batch that have the same sender and type.
+                    event_sender_and_type_to_auth_event_ids_cache[
+                        (ev["sender"], ev["type"])
+                    ] = resultant_auth_event_ids
+
+                event = await builder.build(
+                    prev_event_ids=event_dict.get("prev_events"),
+                    auth_event_ids=resultant_auth_event_ids.copy(),
+                    depth=inherited_depth,
+                )
+                # We can re-use the context per-event type because it will
+                # calculate out to be the same for all events in the batch. We
+                # also get the benefit of sharing the same state_group.
+                context = cached_context
+
+                # TODO: Do we need to check `third_party_event_rules.check_event_allowed(...)`?
+
+                self.validator.validate_new(event, self.config)
 
             logger.debug(
                 "RoomBatchSendEventRestServlet inserting event=%s, prev_event_ids=%s, auth_event_ids=%s",

synapse/storage/databases/main/room_batch.py

@@ -14,6 +14,7 @@
 
 from typing import Optional
 
+from typing import List
 from synapse.storage._base import SQLBaseStore