Move stitched_ordering property to EventBase

... so that we don't need to have an EventContext when we do the assignment

synapse/events/__init__.py

@@ -208,6 +208,8 @@ class EventBase(metaclass=abc.ABCMeta):
 
         self.internal_metadata = EventInternalMetadata(internal_metadata_dict)
 
+        self._stitched_ordering = None
+
     depth: DictProperty[int] = DictProperty("depth")
     content: DictProperty[JsonDict] = DictProperty("content")
     hashes: DictProperty[Dict[str, str]] = DictProperty("hashes")
@@ -323,6 +325,20 @@ class EventBase(metaclass=abc.ABCMeta):
         # this will be a no-op if the event dict is already frozen.
         self._dict = freeze(self._dict)
 
+    def assign_stitched_ordering(self, stitched_ordering: int) -> None:
+        """Assign a stitched ordering to this event, if one has not already been assigned.
+
+        TODO: figure out a way to only expose this on events that have not yet been persisted.
+        """
+        if self._stitched_ordering is not None:
+            raise RuntimeError("Attempt to assign stitched ordering twice")
+        self._stitched_ordering = stitched_ordering
+
+    @property
+    def stitched_ordering(self) -> Optional[int]:
+        """Return the stitched ordering for this event. If one has not (yet) been assigned, returns `None`."""
+        return self._stitched_ordering
+
     def __str__(self) -> str:
         return self.__repr__()
 

synapse/events/snapshot.py

@@ -130,9 +130,6 @@ class EventContext(UnpersistedEventContextBase):
 
         partial_state: if True, we may be storing this event with a temporary,
             incomplete state.
-
-        stitched_ordering: the assigned stitched ordering for this event, or None if we
-            have not yet assigned one.
     """
 
     _storage: "StorageControllers"
@@ -145,8 +142,6 @@ class EventContext(UnpersistedEventContextBase):
 
     partial_state: bool = False
 
-    stitched_ordering: Optional[int] = None
-
     @staticmethod
     def with_state(
         storage: "StorageControllers",

synapse/storage/controllers/persist_events.py

@@ -745,7 +745,7 @@ class EventsPersistenceStorageController:
 
                 # TODO we may need to reorder existing events
                 previous_event_stitched_order += 1
-                context.stitched_ordering = previous_event_stitched_order
+                event.assign_stitched_ordering(previous_event_stitched_order)
                 logger.debug(
                     "Persisting inserted events with stitched_order=%i",
                     previous_event_stitched_order,
@@ -765,9 +765,9 @@ class EventsPersistenceStorageController:
             await self.main_store.get_room_max_stitched_ordering(room_id) or 0
         )
 
-        for _event, context in remaining_batch:
+        for event, _ in remaining_batch:
             current_max_stream_ordering += 2**16
-            context.stitched_ordering = current_max_stream_ordering
+            event.assign_stitched_ordering(current_max_stream_ordering)
 
     async def _calculate_new_forward_extremities_and_state_delta(
         self, room_id: str, ev_ctx_rm: List[EventPersistencePair]

synapse/storage/databases/main/events.py

@@ -2706,7 +2706,7 @@ class PersistEventsStore:
                     "url" in event.content and isinstance(event.content["url"], str),
                     event.get_state_key(),
                     context.rejected,
-                    context.stitched_ordering,
+                    event.stitched_ordering,
                 )
                 for event, context in events_and_contexts
             ],
@@ -3567,7 +3567,7 @@ class PersistEventsStore:
                 lowest_referring_ordering = potential_backwards_extremities.get(
                     "prev_event"
                 )
-                persisted_event_stitched_ordering = ctx.stitched_ordering
+                persisted_event_stitched_ordering = ev.stitched_ordering
 
                 # If any of the events we persisted did not get assigned a stitched order,
                 # we cannot yet assign a stitched order to the backwards extremity either.

tests/storage/test_events.py

@@ -565,4 +565,4 @@ class AssignStitchedOrderingTestCase(HomeserverTestCase):
             self._persistence._assign_stitched_orders(room_id, [(test_event, context)])
         )
 
-        self.assertEqual(context.stitched_ordering, 6 * 2**16)
+        self.assertEqual(test_event.stitched_ordering, 6 * 2**16)