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

synapse/events/snapshot.py

@@ -279,6 +279,7 @@ class UnpersistedEventContext:
             app service.
 
     """
+
     _storage: "StorageControllers"
     state_handler: Optional["StateHandler"] = None
     state_group_before_event: Optional[int] = None

synapse/events/third_party_rules.py

@@ -12,7 +12,7 @@
 # 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
 
@@ -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

@@ -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`
@@ -1173,7 +1178,6 @@ async def on_make_knock_request(
             },
         )
 
-
         (
             event,
             unpersisted_context,
@@ -1345,7 +1349,10 @@ async def exchange_third_party_invite(
             )
 
             EventValidator().validate_builder(builder)
-            event, unpersisted_context = await self.event_creation_handler.create_new_client_event(
+            (
+                event,
+                unpersisted_context,
+            ) = await self.event_creation_handler.create_new_client_event(
                 builder=builder
             )
 

synapse/handlers/federation_event.py

@@ -27,6 +27,7 @@
     Sequence,
     Set,
     Tuple,
+    Union,
 )
 
 from prometheus_client import Counter, Histogram
@@ -57,7 +58,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 (
@@ -424,7 +425,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
 

synapse/handlers/message.py

@@ -708,7 +708,6 @@ async def create_event(
 
         builder.internal_metadata.historical = historical
 
-
         event, unpersisted_context = await self.create_new_client_event(
             builder=builder,
             requester=requester,
@@ -1073,7 +1072,7 @@ async def create_new_client_event(
         state_map: Optional[StateMap[str]] = None,
         for_batch: bool = False,
         current_state_group: Optional[int] = None,
-    ) -> Tuple[EventBase, EventContext]:
+    ) -> Tuple[EventBase, UnpersistedEventContext]:
         """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 for
         the event using the parameters state_map and current_state_group, thus these parameters
@@ -1201,7 +1200,7 @@ async def create_new_client_event(
             # after it is created
             if builder.internal_metadata.outlier:
                 event.internal_metadata.outlier = True
-                context = EventContext.for_outlier(self._storage_controllers)
+                context = UnpersistedEventContext.for_outlier(self._storage_controllers)
             elif (
                 event.type == EventTypes.MSC2716_INSERTION
                 and state_event_ids
@@ -1287,7 +1286,7 @@ async def create_new_client_event(
         elif new_content is not None:
             # the third-party rules want to replace the event. We'll need to build a new
             # event.
-            event, context = await self._rebuild_event_after_third_party_rules(
+            event = await self._rebuild_event_after_third_party_rules(
                 new_content, event
             )
 
@@ -2076,9 +2075,9 @@ def _expire_rooms_to_exclude_from_dummy_event_insertion(self) -> None:
 
     async def _rebuild_event_after_third_party_rules(
         self, third_party_result: dict, original_event: EventBase
-    ) -> Tuple[EventBase, EventContext]:
+    ) -> EventBase:
         # the third_party_event_rules want to replace the event.
-        # we do some basic checks, and then return the replacement event and context.
+        # we do some basic checks, and then return the replacement event.
 
         # Construct a new EventBuilder and validate it, which helps with the
         # rest of these checks.
@@ -2130,7 +2129,4 @@ async def _rebuild_event_after_third_party_rules(
             auth_event_ids=None,
         )
 
-        # we rebuild the event context, to be on the safe side. If nothing else,
-        # delta_ids might need an update.
-        context = await self.state.compute_event_context(event)
-        return event, context
+        return event