Add a class UnpersistedEventContext to allow for the batching up of storing state groups

changelog.d/14675.misc

@@ -0,0 +1 @@
+Add a class UnpersistedEventContext to allow for the batching up of storing state groups.

synapse/events/snapshot.py

@@ -16,11 +16,13 @@
 import attr
 from frozendict import frozendict
 
+from synapse.api.constants import EventTypes
 from synapse.appservice import ApplicationService
 from synapse.events import EventBase
 from synapse.types import JsonDict, StateMap
 
 if TYPE_CHECKING:
+    from synapse.state import StateHandler, _StateCacheEntry
     from synapse.storage.controllers import StorageControllers
     from synapse.storage.databases.main import DataStore
     from synapse.types.state import StateFilter
@@ -254,6 +256,97 @@ async def get_prev_state_ids(
         )
 
 
+@attr.s(slots=True, auto_attribs=True)
+class UnpersistedEventContext:
+    """
+    This is a version of an EventContext before the new state group (if any) has been
+    computed and stored. It contains information about the state before the event (which
+    also may be the information after the event, if the event is not a state event). The
+    UnpersistedEventContext must be converted into an EventContext by calling the method
+    'persist' on it before it is suitable to be sent to the DB for processing.
+
+        state_handler:
+            an instance of the class StateHandler, this is required for all
+            but outlier instances
+        state_group_before_event:
+            the state group at/before the event. This is required if `for_batch` is True.
+        state_map:
+             the current state of the room
+        for_batch:
+            whether the context is part of a group of events being created for batch persisting
+            to the DB.
+        state_entry:
+            if a state entry is created while calculating this context, it can be stored and
+            passed in when persisting the context to avoid duplicated work
+        app_service: If this event is being sent by a (local) application service, that
+            app service.
+
+    """
+
+    _storage: "StorageControllers"
+    state_handler: Optional["StateHandler"] = None
+    state_group_before_event: Optional[int] = None
+    state_map: Optional[StateMap[str]] = None
+    for_batch: bool = False
+    state_entry: Optional["_StateCacheEntry"] = None
+    app_service: Optional[ApplicationService] = None
+
+    @staticmethod
+    def for_outlier(
+        storage: "StorageControllers",
+    ) -> "UnpersistedEventContext":
+        """Returns an UnpersistedEventContext instance suitable for an outlier event"""
+        return UnpersistedEventContext(storage=storage)
+
+    async def persist(
+        self,
+        event: EventBase,
+    ) -> "EventContext":
+        """
+        Converts an UnpersistedEventContext into an EventContext, which is now suitable for
+        sending to the DB along with the event.
+        """
+        if event.internal_metadata.is_outlier():
+            return EventContext.for_outlier(self._storage)
+
+        assert self.state_handler is not None
+        if self.for_batch:
+            assert self.state_group_before_event is not None
+            assert self.state_map is not None
+            return await self.state_handler.compute_event_context_for_batched(
+                event, self.state_map, self.state_group_before_event
+            )
+        elif (
+            event.type == EventTypes.MSC2716_INSERTION
+            and self.state_map
+            and event.internal_metadata.is_historical()
+        ):
+            # TODO(faster_joins): check how MSC2716 works and whether we can have
+            #   partial state here
+            #   https://github.com/matrix-org/synapse/issues/13003
+            return await self.state_handler.compute_event_context(
+                event, state_ids_before_event=self.state_map, partial_state=False
+            )
+        return await self.state_handler.compute_event_context(
+            event, entry=self.state_entry
+        )
+
+    async def get_prev_state_ids(
+        self, state_filter: Optional["StateFilter"] = None
+    ) -> StateMap[str]:
+        """
+        Returns the state at/before the event (which also may be the same as after the event
+        if the event is not a state event.)
+        """
+        if self.state_map is not None:
+            return self.state_map
+        else:
+            assert self.state_group_before_event is not None
+            return await self._storage.state.get_state_ids_for_group(
+                self.state_group_before_event, state_filter
+            )
+
+
 def _encode_state_dict(
     state_dict: Optional[StateMap[str]],
 ) -> Optional[List[Tuple[str, str, str]]]:

synapse/events/third_party_rules.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, List, Optional, Tuple, Union
 
 from twisted.internet.defer import CancelledError
 
 from synapse.api.errors import ModuleFailedException, SynapseError
 from synapse.events import EventBase
-from synapse.events.snapshot import EventContext
+from synapse.events.snapshot import EventContext, UnpersistedEventContext
 from synapse.storage.roommember import ProfileInfo
 from synapse.types import Requester, StateMap
 from synapse.util.async_helpers import delay_cancellation, maybe_awaitable
@@ -231,7 +231,7 @@ def register_third_party_rules_callbacks(
             self._on_threepid_bind_callbacks.append(on_threepid_bind)
 
     async def check_event_allowed(
-        self, event: EventBase, context: EventContext
+        self, event: EventBase, context: Union[EventContext, UnpersistedEventContext]
     ) -> Tuple[bool, Optional[dict]]:
         """Check if a provided event should be allowed in the given context.
 

synapse/handlers/federation.py

@@ -56,7 +56,7 @@
 from synapse.crypto.event_signing import compute_event_signature
 from synapse.event_auth import validate_event_for_room_version
 from synapse.events import EventBase
-from synapse.events.snapshot import EventContext
+from synapse.events.snapshot import EventContext, UnpersistedEventContext
 from synapse.events.validator import EventValidator
 from synapse.federation.federation_client import InvalidResponseError
 from synapse.http.servlet import assert_params_in_dict
@@ -935,15 +935,20 @@ async def on_make_join_request(
         )
 
         try:
-            event, context = await self.event_creation_handler.create_new_client_event(
+            (
+                event,
+                unpersisted_context,
+            ) = await self.event_creation_handler.create_new_client_event(
                 builder=builder
             )
         except SynapseError as e:
             logger.warning("Failed to create join to %s because %s", room_id, e)
             raise
 
         # Ensure the user can even join the room.
-        await self._federation_event_handler.check_join_restrictions(context, event)
+        await self._federation_event_handler.check_join_restrictions(
+            unpersisted_context, event
+        )
 
         # The remote hasn't signed it yet, obviously. We'll do the full checks
         # when we get the event back in `on_send_join_request`
@@ -1123,7 +1128,7 @@ async def on_make_leave_request(
             },
         )
 
-        event, context = await self.event_creation_handler.create_new_client_event(
+        event, _ = await self.event_creation_handler.create_new_client_event(
             builder=builder
         )
 
@@ -1173,12 +1178,13 @@ async def on_make_knock_request(
             },
         )
 
-        event, context = await self.event_creation_handler.create_new_client_event(
-            builder=builder
-        )
+        (
+            event,
+            unpersisted_context,
+        ) = await self.event_creation_handler.create_new_client_event(builder=builder)
 
         event_allowed, _ = await self.third_party_event_rules.check_event_allowed(
-            event, context
+            event, unpersisted_context
         )
         if not event_allowed:
             logger.warning("Creation of knock %s forbidden by third-party rules", event)
@@ -1343,14 +1349,22 @@ async def exchange_third_party_invite(
             )
 
             EventValidator().validate_builder(builder)
-            event, context = await self.event_creation_handler.create_new_client_event(
+            (
+                event,
+                unpersisted_context,
+            ) = await self.event_creation_handler.create_new_client_event(
                 builder=builder
             )
 
-            event, context = await self.add_display_name_to_third_party_invite(
-                room_version_obj, event_dict, event, context
+            (
+                event,
+                unpersisted_context,
+            ) = await self.add_display_name_to_third_party_invite(
+                room_version_obj, event_dict, event, unpersisted_context
             )
 
+            context = await unpersisted_context.persist(event)
+
             EventValidator().validate_new(event, self.config)
 
             # We need to tell the transaction queue to send this out, even
@@ -1400,13 +1414,17 @@ async def on_exchange_third_party_invite_request(
             room_version_obj, event_dict
         )
 
-        event, context = await self.event_creation_handler.create_new_client_event(
-            builder=builder
-        )
-        event, context = await self.add_display_name_to_third_party_invite(
-            room_version_obj, event_dict, event, context
+        (
+            event,
+            unpersisted_context,
+        ) = await self.event_creation_handler.create_new_client_event(builder=builder)
+
+        event, unpersisted_context = await self.add_display_name_to_third_party_invite(
+            room_version_obj, event_dict, event, unpersisted_context
         )
 
+        context = await unpersisted_context.persist(event)
+
         try:
             validate_event_for_room_version(event)
             await self._event_auth_handler.check_auth_rules_from_context(event)
@@ -1428,8 +1446,8 @@ async def add_display_name_to_third_party_invite(
         room_version_obj: RoomVersion,
         event_dict: JsonDict,
         event: EventBase,
-        context: EventContext,
-    ) -> Tuple[EventBase, EventContext]:
+        context: UnpersistedEventContext,
+    ) -> Tuple[EventBase, UnpersistedEventContext]:
         key = (
             EventTypes.ThirdPartyInvite,
             event.content["third_party_invite"]["signed"]["token"],
@@ -1463,11 +1481,14 @@ async def add_display_name_to_third_party_invite(
             room_version_obj, event_dict
         )
         EventValidator().validate_builder(builder)
-        event, context = await self.event_creation_handler.create_new_client_event(
-            builder=builder
-        )
+
+        (
+            event,
+            unpersisted_context,
+        ) = await self.event_creation_handler.create_new_client_event(builder=builder)
+
         EventValidator().validate_new(event, self.config)
-        return event, context
+        return event, unpersisted_context
 
     async def _check_signature(self, event: EventBase, context: EventContext) -> None:
         """

synapse/handlers/federation_event.py

@@ -27,6 +27,7 @@
     Sequence,
     Set,
     Tuple,
+    Union,
 )
 
 from prometheus_client import Counter, Histogram
@@ -58,7 +59,7 @@
     validate_event_for_room_version,
 )
 from synapse.events import EventBase
-from synapse.events.snapshot import EventContext
+from synapse.events.snapshot import EventContext, UnpersistedEventContext
 from synapse.federation.federation_client import InvalidResponseError, PulledPduInfo
 from synapse.logging.context import nested_logging_context
 from synapse.logging.opentracing import (
@@ -425,7 +426,7 @@ async def on_send_membership_event(
         return event, context
 
     async def check_join_restrictions(
-        self, context: EventContext, event: EventBase
+        self, context: Union[EventContext, UnpersistedEventContext], event: EventBase
     ) -> None:
         """Check that restrictions in restricted join rules are matched