Replace EventContext fields prev_group and delta_ids with field state_group_deltas

changelog.d/15233.misc

@@ -0,0 +1 @@
+Replace `EventContext` fields `prev_group` and `delta_ids` with field `state_group_deltas`.

synapse/events/snapshot.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, List, Optional, Tuple
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
 import attr
 from immutabledict import immutabledict
@@ -106,21 +106,15 @@ class EventContext(UnpersistedEventContextBase):
         state_delta_due_to_event: If `state_group` and `state_group_before_event` are not None
             then this is the delta of the state between the two groups.
 
-        prev_group: If it is known, ``state_group``'s prev_group. Note that this being
-            None does not necessarily mean that ``state_group`` does not have
-            a prev_group!
-
-            If the event is a state event, this is normally the same as
-            ``state_group_before_event``.
-
-            If ``state_group`` is None (ie, the event is an outlier), ``prev_group``
-            will always also be ``None``.
-
-            Note that this *not* (necessarily) the state group associated with
-            ``_prev_state_ids``.
-
-        delta_ids: If ``prev_group`` is not None, the state delta between ``prev_group``
-            and ``state_group``.
+        state_group_deltas: if it exists, this is a dict collecting a mapping of two possible
+            sets of state groups and their deltas: the state group before the event and
+            after the event and the state map that represents the state group between the
+            two, i.e. (state_group_2, state_group_1) -> (state_key, event_type) -> event_id
+            or/and the state group preceding the state group before the event and the state
+            group before the event and the state map that represents the delta between those
+            two state groups.
+            This information is collected and stored as part of an optimization for persisting
+            events.
 
         partial_state: if True, we may be storing this event with a temporary,
             incomplete state.
@@ -131,8 +125,7 @@ class EventContext(UnpersistedEventContextBase):
     _state_group: Optional[int] = None
     state_group_before_event: Optional[int] = None
     _state_delta_due_to_event: Optional[StateMap[str]] = None
-    prev_group: Optional[int] = None
-    delta_ids: Optional[StateMap[str]] = None
+    state_group_deltas: Optional[Dict[Tuple[int, int], StateMap[str]]] = None
     app_service: Optional[ApplicationService] = None
 
     partial_state: bool = False
@@ -144,16 +137,14 @@ def with_state(
         state_group_before_event: Optional[int],
         state_delta_due_to_event: Optional[StateMap[str]],
         partial_state: bool,
-        prev_group: Optional[int] = None,
-        delta_ids: Optional[StateMap[str]] = None,
+        state_group_deltas: Optional[Dict[Tuple[int, int], StateMap[str]]],
     ) -> "EventContext":
         return EventContext(
             storage=storage,
             state_group=state_group,
             state_group_before_event=state_group_before_event,
             state_delta_due_to_event=state_delta_due_to_event,
-            prev_group=prev_group,
-            delta_ids=delta_ids,
+            state_group_deltas=state_group_deltas,
             partial_state=partial_state,
         )
 
@@ -177,16 +168,14 @@ async def serialize(self, event: EventBase, store: "DataStore") -> JsonDict:
         Returns:
             The serialized event.
         """
-
         return {
             "state_group": self._state_group,
             "state_group_before_event": self.state_group_before_event,
             "rejected": self.rejected,
-            "prev_group": self.prev_group,
+            "state_group_deltas": _encode_state_group_delta(self.state_group_deltas),
             "state_delta_due_to_event": _encode_state_dict(
                 self._state_delta_due_to_event
             ),
-            "delta_ids": _encode_state_dict(self.delta_ids),
             "app_service_id": self.app_service.id if self.app_service else None,
             "partial_state": self.partial_state,
         }
@@ -209,11 +198,10 @@ def deserialize(storage: "StorageControllers", input: JsonDict) -> "EventContext
             storage=storage,
             state_group=input["state_group"],
             state_group_before_event=input["state_group_before_event"],
-            prev_group=input["prev_group"],
+            state_group_deltas=_decode_state_group_delta(input["state_group_deltas"]),
             state_delta_due_to_event=_decode_state_dict(
                 input["state_delta_due_to_event"]
             ),
-            delta_ids=_decode_state_dict(input["delta_ids"]),
             rejected=input["rejected"],
             partial_state=input.get("partial_state", False),
         )
@@ -344,7 +332,7 @@ class UnpersistedEventContext(UnpersistedEventContextBase):
     _storage: "StorageControllers"
     state_group_before_event: Optional[int]
     state_group_after_event: Optional[int]
-    state_delta_due_to_event: Optional[dict]
+    state_delta_due_to_event: Optional[StateMap[str]]
     prev_group_for_state_group_before_event: Optional[int]
     delta_ids_to_state_group_before_event: Optional[StateMap[str]]
     partial_state: bool
@@ -375,26 +363,16 @@ async def batch_persist_unpersisted_contexts(
 
         events_and_persisted_context = []
         for event, unpersisted_context in amended_events_and_context:
-            if event.is_state():
-                context = EventContext(
-                    storage=unpersisted_context._storage,
-                    state_group=unpersisted_context.state_group_after_event,
-                    state_group_before_event=unpersisted_context.state_group_before_event,
-                    state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
-                    partial_state=unpersisted_context.partial_state,
-                    prev_group=unpersisted_context.state_group_before_event,
-                    delta_ids=unpersisted_context.state_delta_due_to_event,
-                )
-            else:
-                context = EventContext(
-                    storage=unpersisted_context._storage,
-                    state_group=unpersisted_context.state_group_after_event,
-                    state_group_before_event=unpersisted_context.state_group_before_event,
-                    state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
-                    partial_state=unpersisted_context.partial_state,
-                    prev_group=unpersisted_context.prev_group_for_state_group_before_event,
-                    delta_ids=unpersisted_context.delta_ids_to_state_group_before_event,
-                )
+            state_group_deltas = _build_state_group_deltas(unpersisted_context)
+
+            context = EventContext(
+                storage=unpersisted_context._storage,
+                state_group=unpersisted_context.state_group_after_event,
+                state_group_before_event=unpersisted_context.state_group_before_event,
+                state_delta_due_to_event=unpersisted_context.state_delta_due_to_event,
+                partial_state=unpersisted_context.partial_state,
+                state_group_deltas=state_group_deltas,
+            )
             events_and_persisted_context.append((event, context))
         return events_and_persisted_context
 
@@ -447,26 +425,27 @@ async def persist(self, event: EventBase) -> EventContext:
 
         # if the event isn't a state event the state group doesn't change
         if not self.state_delta_due_to_event:
-            state_group_after_event = self.state_group_before_event
+            self.state_group_after_event = self.state_group_before_event
 
         # otherwise if it is a state event we need to get a state group for it
         else:
-            state_group_after_event = await self._storage.state.store_state_group(
+            self.state_group_after_event = await self._storage.state.store_state_group(
                 event.event_id,
                 event.room_id,
                 prev_group=self.state_group_before_event,
                 delta_ids=self.state_delta_due_to_event,
                 current_state_ids=None,
             )
 
+        state_group_deltas = _build_state_group_deltas(self)
+
         return EventContext.with_state(
             storage=self._storage,
-            state_group=state_group_after_event,
+            state_group=self.state_group_after_event,
             state_group_before_event=self.state_group_before_event,
             state_delta_due_to_event=self.state_delta_due_to_event,
+            state_group_deltas=state_group_deltas,
             partial_state=self.partial_state,
-            prev_group=self.state_group_before_event,
-            delta_ids=self.state_delta_due_to_event,
         )
 
 
@@ -482,6 +461,34 @@ def _encode_state_dict(
     return [(etype, state_key, v) for (etype, state_key), v in state_dict.items()]
 
 
+def _encode_state_group_delta(
+    state_group_delta: Optional[Dict[Tuple[int, int], StateMap[str]]]
+) -> Optional[List[Tuple[int, int, Optional[List[Tuple[str, str, str]]]]]]:
+    if not state_group_delta:
+        return None
+
+    state_group_delta_encoded = []
+    for key, value in state_group_delta.items():
+        state_group_delta_encoded.append((key[0], key[1], _encode_state_dict(value)))
+
+    return state_group_delta_encoded
+
+
+def _decode_state_group_delta(
+    input: Optional[List[Tuple[int, int, List[Tuple[str, str, str]]]]]
+) -> Optional[Dict[Tuple[int, int], StateMap[str]]]:
+    if not input:
+        return None
+
+    state_group_deltas = {}
+    for element in input:
+        state_map = _decode_state_dict(element[2])
+        assert state_map is not None
+        state_group_deltas[(element[0], element[1])] = state_map
+
+    return state_group_deltas
+
+
 def _decode_state_dict(
     input: Optional[List[Tuple[str, str, str]]]
 ) -> Optional[StateMap[str]]:
@@ -490,3 +497,47 @@ def _decode_state_dict(
         return None
 
     return immutabledict({(etype, state_key): v for etype, state_key, v in input})
+
+
+def _build_state_group_deltas(
+    unpersisted_ctx: UnpersistedEventContext,
+) -> Dict[Tuple[int, int], StateMap]:
+    """
+    Internal function to take an UnpersistedEventContext and collect any known deltas
+    between the state groups associated with the context. This information is used in an
+    optimization when persisting events.
+    """
+    state_group_deltas = {}
+
+    # if we know the state group before the event and after the event, add them and the
+    # state delta between them to state_group_deltas
+    if (
+        unpersisted_ctx.state_group_before_event
+        and unpersisted_ctx.state_group_after_event
+    ):
+        # if we have the state groups we should have the delta
+        assert unpersisted_ctx.state_delta_due_to_event is not None
+        state_group_deltas[
+            (
+                unpersisted_ctx.state_group_before_event,
+                unpersisted_ctx.state_group_after_event,
+            )
+        ] = unpersisted_ctx.state_delta_due_to_event
+
+    # the state group before the event may also have a state group which precedes it, if
+    # we have that and the state group before the event, add them and the state
+    # delta between them to state_group_deltas
+    if (
+        unpersisted_ctx.prev_group_for_state_group_before_event
+        and unpersisted_ctx.state_group_before_event
+    ):
+        # if we have both state groups we should have the delta between them
+        assert unpersisted_ctx.delta_ids_to_state_group_before_event is not None
+        state_group_deltas[
+            (
+                unpersisted_ctx.prev_group_for_state_group_before_event,
+                unpersisted_ctx.state_group_before_event,
+            )
+        ] = unpersisted_ctx.delta_ids_to_state_group_before_event
+
+    return state_group_deltas

synapse/storage/controllers/persist_events.py

@@ -839,9 +839,8 @@ async def _get_new_state_after_events(
                         "group" % (ev.event_id,)
                     )
                 continue
-
-            if ctx.prev_group:
-                state_group_deltas[(ctx.prev_group, ctx.state_group)] = ctx.delta_ids
+            if ctx.state_group_deltas:
+                state_group_deltas.update(ctx.state_group_deltas)
 
         # We need to map the event_ids to their state groups. First, let's
         # check if the event is one we're persisting, in which case we can

tests/events/test_snapshot.py

@@ -101,8 +101,7 @@ def _check_serialize_deserialize(
         self.assertEqual(
             context.state_group_before_event, d_context.state_group_before_event
         )
-        self.assertEqual(context.prev_group, d_context.prev_group)
-        self.assertEqual(context.delta_ids, d_context.delta_ids)
+        self.assertEqual(context.state_group_deltas, d_context.state_group_deltas)
         self.assertEqual(context.app_service, d_context.app_service)
 
         self.assertEqual(

tests/test_state.py

@@ -556,9 +556,16 @@ def test_annotate_with_old_state(
         )
 
         self.assertIsNotNone(context.state_group_before_event)
+        assert context.state_group_deltas is not None
+        assert context.state_group_before_event is not None
+        assert context.state_group is not None
+        self.assertEqual(
+            context.state_group_deltas.get(
+                (context.state_group_before_event, context.state_group)
+            ),
+            {(event.type, event.state_key): event.event_id},
+        )
         self.assertNotEqual(context.state_group_before_event, context.state_group)
-        self.assertEqual(context.state_group_before_event, context.prev_group)
-        self.assertEqual({("state", ""): event.event_id}, context.delta_ids)
 
     @defer.inlineCallbacks
     def test_trivial_annotate_message(