Assign stitched orders to pulled events before persisting

synapse/handlers/federation_event.py

@@ -86,6 +86,7 @@ from synapse.replication.http.federation import (
     ReplicationFederationSendEventsRestServlet,
 )
 from synapse.state import StateResolutionStore
+from synapse.storage.controllers.persist_events import assign_stitched_orders
 from synapse.storage.database import LoggingTransaction
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
 from synapse.types import (
@@ -895,7 +896,10 @@ class FederationEventHandler:
         )
 
         @trace
-        async def _process_new_pulled_events(new_events: Collection[EventBase]) -> None:
+        async def _process_new_pulled_events(new_events: List[EventBase]) -> None:
+            room_id = new_events[0].room_id
+            await assign_stitched_orders(room_id, new_events, self._store)
+
             # We want to sort these by depth so we process them and tell clients about
             # them in order. It's also more efficient to backfill this way (`depth`
             # ascending) because one backfill event is likely to be the `prev_event` of

synapse/storage/controllers/persist_events.py

@@ -617,7 +617,11 @@ class EventsPersistenceStorageController:
             if not events_and_contexts:
                 return replaced_events
 
-        await assign_stitched_orders(room_id, [ev for (ev, _) in events_and_contexts], self.main_store)
+        # TODO massive hack
+        if events_and_contexts[0][0].stitched_ordering is None:
+            await assign_stitched_orders(
+                room_id, [ev for (ev, _) in events_and_contexts], self.main_store
+            )
 
         chunks = [
             events_and_contexts[x : x + 100]
@@ -1347,9 +1351,7 @@ async def assign_stitched_orders(
             )
 
         remaining_batch = still_remaining_batch
-        logger.debug(
-            "Remaining events: %s", [ev.event_id for ev in remaining_batch]
-        )
+        logger.debug("Remaining events: %s", [ev.event_id for ev in remaining_batch])
 
     logger.debug(
         "Remaining events after processing gap matches: %s",

tests/storage/test_events.py

@@ -561,8 +561,6 @@ class AssignStitchedOrderingTestCase(HomeserverTestCase):
             RoomVersions.V12,
         )
 
-        self.get_success(
-            assign_stitched_orders(room_id, [test_event], self.store)
-        )
+        self.get_success(assign_stitched_orders(room_id, [test_event], self.store))
 
         self.assertEqual(test_event.stitched_ordering, 6 * 2**16)