diff --git a/synapse/events/__init__.py b/synapse/events/__init__.py index db38754280..7f1fb3c394 100644 --- a/synapse/events/__init__.py +++ b/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__() diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py index 762d6319d5..63551143d8 100644 --- a/synapse/events/snapshot.py +++ b/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", diff --git a/synapse/storage/controllers/persist_events.py b/synapse/storage/controllers/persist_events.py index 3fbf789a60..588bc8cb1e 100644 --- a/synapse/storage/controllers/persist_events.py +++ b/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] diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py index 2f38859ea2..63ed255cad 100644 --- a/synapse/storage/databases/main/events.py +++ b/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. diff --git a/tests/storage/test_events.py b/tests/storage/test_events.py index cf249311de..792bbbae2b 100644 --- a/tests/storage/test_events.py +++ b/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)