From 8cd1196c4bfef07183192e2a8d7efcec5b3b5f30 Mon Sep 17 00:00:00 2001 From: "H. Shay" Date: Tue, 8 Nov 2022 11:28:56 -0800 Subject: update compute_event_context_for_batched to take list and assign state groups --- synapse/state/__init__.py | 78 ++++++++++++++++++++++++++--------------------- 1 file changed, 44 insertions(+), 34 deletions(-) diff --git a/synapse/state/__init__.py b/synapse/state/__init__.py index 0b181a2c5a..30bbc0437a 100644 --- a/synapse/state/__init__.py +++ b/synapse/state/__init__.py @@ -422,49 +422,59 @@ class StateHandler: async def compute_event_context_for_batched( self, - event: EventBase, - state_ids_before_event: StateMap[str], - ) -> EventContext: + events_and_context: List[Tuple[EventBase, EventContext]], + prev_group: int, + state_ids_before_event: StateMap, + ) -> List[Tuple[EventBase, EventContext]]: """ Generate an event context for an event that has not yet been persisted to the database. Intended for use with events that are created to be persisted in a batch. Args: - event: the event the context is being computed for - state_ids_before_event: a state map consisting of the state ids of the events - created prior to this event. - current_state_group: the current state group before the event. + events_and_context: a list of events and their associated contexts + prev_group: the state group of the last event persisted before the batched events + were created + state_ids_before_event: a state map consisting of current state ids """ - state_group_before_event_prev_group = None - deltas_to_state_group_before_event = None - - # if the event is not state, we are set - if not event.is_state(): - return EventContext.without_state_group( - storage=self._storage_controllers, - state_delta_due_to_event={}, - prev_group=state_group_before_event_prev_group, - delta_ids=deltas_to_state_group_before_event, - partial_state=False, - ) + # separate out state and non-state contexts + state_events = [] + for event, context in events_and_context: + if event.is_state(): + state_events.append((event, context)) + + # get state groups for state events + room_id = events_and_context[0][0].room_id + assert self.hs.datastores is not None + await self.hs.datastores.state.store_state_deltas_for_batched( + state_events, room_id, prev_group=prev_group + ) - # otherwise, we'll need to create a new state group for after the event - key = (event.type, event.state_key) + # iterate through all contexts and update everything + current_state_group = prev_group + for event, context in events_and_context: + + # if the event is not state, we need to update it + if not event.is_state(): + context._state_group = current_state_group + context.state_group_before_event = current_state_group + context._state_delta_due_to_event = {} + context.prev_group = None + context.delta_ids = None + context.partial_state = False + + # the context should have been updated when storing the state groups but let's + # be sure - if it does not have a state group there is a problem + if context._state_group is None: + raise RuntimeError( + f"Event {event.event_id} is missing a state group." + ) + current_state_group = context._state_group - if state_ids_before_event is not None: + key = (event.type, event.state_key) replaces = state_ids_before_event.get(key) + if replaces and replaces != event.event_id: + event.unsigned["replaces_state"] = replaces - if replaces and replaces != event.event_id: - event.unsigned["replaces_state"] = replaces - - delta_ids = {key: event.event_id} - - context = EventContext.without_state_group( - storage=self._storage_controllers, - state_delta_due_to_event=delta_ids, - delta_ids=delta_ids, - partial_state=False, - ) - return context + return events_and_context @measure_func() async def resolve_state_groups_for_events( -- cgit 1.4.1