Bundle aggregations outside of the serialization method. (#11612)

This makes the serialization of events synchronous (and it no
longer access the database), but we must manually calculate and
provide the bundled aggregations.

Overall this should cause no change in behavior, but is prep work
for other improvements.

6 files changed, 56 insertions, 42 deletions

diff --git a/synapse/rest/admin/rooms.py b/synapse/rest/admin/rooms.py
index 6030373ebc..2e714ac87b 100644
--- a/synapse/rest/admin/rooms.py
+++ b/synapse/rest/admin/rooms.py
@@ -424,7 +424,7 @@ class RoomStateRestServlet(RestServlet):
         event_ids = await self.store.get_current_state_ids(room_id)
         events = await self.store.get_events(event_ids.values())
         now = self.clock.time_msec()
-        room_state = await self._event_serializer.serialize_events(events.values(), now)
+        room_state = self._event_serializer.serialize_events(events.values(), now)
         ret = {"state": room_state}
 
         return HTTPStatus.OK, ret
@@ -744,22 +744,22 @@ class RoomEventContextServlet(RestServlet):
             )
 
         time_now = self.clock.time_msec()
-        results["events_before"] = await self._event_serializer.serialize_events(
+        results["events_before"] = self._event_serializer.serialize_events(
             results["events_before"],
             time_now,
-            bundle_aggregations=True,
+            bundle_aggregations=results["aggregations"],
         )
-        results["event"] = await self._event_serializer.serialize_event(
+        results["event"] = self._event_serializer.serialize_event(
             results["event"],
             time_now,
-            bundle_aggregations=True,
+            bundle_aggregations=results["aggregations"],
         )
-        results["events_after"] = await self._event_serializer.serialize_events(
+        results["events_after"] = self._event_serializer.serialize_events(
             results["events_after"],
             time_now,
-            bundle_aggregations=True,
+            bundle_aggregations=results["aggregations"],
         )
-        results["state"] = await self._event_serializer.serialize_events(
+        results["state"] = self._event_serializer.serialize_events(
             results["state"], time_now
         )
 
diff --git a/synapse/rest/client/events.py b/synapse/rest/client/events.py
index 13b72a045a..672c821061 100644
--- a/synapse/rest/client/events.py
+++ b/synapse/rest/client/events.py
@@ -91,7 +91,7 @@ class EventRestServlet(RestServlet):
 
         time_now = self.clock.time_msec()
         if event:
-            result = await self._event_serializer.serialize_event(event, time_now)
+            result = self._event_serializer.serialize_event(event, time_now)
             return 200, result
         else:
             return 404, "Event not found."
diff --git a/synapse/rest/client/notifications.py b/synapse/rest/client/notifications.py
index acd0c9e135..8e427a96a3 100644
--- a/synapse/rest/client/notifications.py
+++ b/synapse/rest/client/notifications.py
@@ -72,7 +72,7 @@ class NotificationsServlet(RestServlet):
                 "actions": pa.actions,
                 "ts": pa.received_ts,
                 "event": (
-                    await self._event_serializer.serialize_event(
+                    self._event_serializer.serialize_event(
                         notif_events[pa.event_id],
                         self.clock.time_msec(),
                         event_format=format_event_for_client_v2_without_room_id,
diff --git a/synapse/rest/client/relations.py b/synapse/rest/client/relations.py
index 3823498012..37d949a71e 100644
--- a/synapse/rest/client/relations.py
+++ b/synapse/rest/client/relations.py
@@ -113,13 +113,14 @@ class RelationPaginationServlet(RestServlet):
         now = self.clock.time_msec()
         # Do not bundle aggregations when retrieving the original event because
         # we want the content before relations are applied to it.
-        original_event = await self._event_serializer.serialize_event(
-            event, now, bundle_aggregations=False
+        original_event = self._event_serializer.serialize_event(
+            event, now, bundle_aggregations=None
         )
         # The relations returned for the requested event do include their
         # bundled aggregations.
-        serialized_events = await self._event_serializer.serialize_events(
-            events, now, bundle_aggregations=True
+        aggregations = await self.store.get_bundled_aggregations(events)
+        serialized_events = self._event_serializer.serialize_events(
+            events, now, bundle_aggregations=aggregations
         )
 
         return_value = pagination_chunk.to_dict()
@@ -308,7 +309,7 @@ class RelationAggregationGroupPaginationServlet(RestServlet):
         )
 
         now = self.clock.time_msec()
-        serialized_events = await self._event_serializer.serialize_events(events, now)
+        serialized_events = self._event_serializer.serialize_events(events, now)
 
         return_value = result.to_dict()
         return_value["chunk"] = serialized_events
diff --git a/synapse/rest/client/room.py b/synapse/rest/client/room.py
index 40330749e5..da6014900a 100644
--- a/synapse/rest/client/room.py
+++ b/synapse/rest/client/room.py
@@ -642,6 +642,7 @@ class RoomEventServlet(RestServlet):
     def __init__(self, hs: "HomeServer"):
         super().__init__()
         self.clock = hs.get_clock()
+        self._store = hs.get_datastore()
         self.event_handler = hs.get_event_handler()
         self._event_serializer = hs.get_event_client_serializer()
         self.auth = hs.get_auth()
@@ -660,10 +661,13 @@ class RoomEventServlet(RestServlet):
             # https://matrix.org/docs/spec/client_server/r0.5.0#get-matrix-client-r0-rooms-roomid-event-eventid
             raise SynapseError(404, "Event not found.", errcode=Codes.NOT_FOUND)
 
-        time_now = self.clock.time_msec()
         if event:
-            event_dict = await self._event_serializer.serialize_event(
-                event, time_now, bundle_aggregations=True
+            # Ensure there are bundled aggregations available.
+            aggregations = await self._store.get_bundled_aggregations([event])
+
+            time_now = self.clock.time_msec()
+            event_dict = self._event_serializer.serialize_event(
+                event, time_now, bundle_aggregations=aggregations
             )
             return 200, event_dict
 
@@ -708,16 +712,20 @@ class RoomEventContextServlet(RestServlet):
             raise SynapseError(404, "Event not found.", errcode=Codes.NOT_FOUND)
 
         time_now = self.clock.time_msec()
-        results["events_before"] = await self._event_serializer.serialize_events(
-            results["events_before"], time_now, bundle_aggregations=True
+        results["events_before"] = self._event_serializer.serialize_events(
+            results["events_before"],
+            time_now,
+            bundle_aggregations=results["aggregations"],
         )
-        results["event"] = await self._event_serializer.serialize_event(
-            results["event"], time_now, bundle_aggregations=True
+        results["event"] = self._event_serializer.serialize_event(
+            results["event"], time_now, bundle_aggregations=results["aggregations"]
         )
-        results["events_after"] = await self._event_serializer.serialize_events(
-            results["events_after"], time_now, bundle_aggregations=True
+        results["events_after"] = self._event_serializer.serialize_events(
+            results["events_after"],
+            time_now,
+            bundle_aggregations=results["aggregations"],
         )
-        results["state"] = await self._event_serializer.serialize_events(
+        results["state"] = self._event_serializer.serialize_events(
             results["state"], time_now
         )
 
diff --git a/synapse/rest/client/sync.py b/synapse/rest/client/sync.py
index e99a943d0d..a3e57e4b20 100644
--- a/synapse/rest/client/sync.py
+++ b/synapse/rest/client/sync.py
@@ -17,7 +17,6 @@ from collections import defaultdict
 from typing import (
     TYPE_CHECKING,
     Any,
-    Awaitable,
     Callable,
     Dict,
     Iterable,
@@ -395,7 +394,7 @@ class SyncRestServlet(RestServlet):
         """
         invited = {}
         for room in rooms:
-            invite = await self._event_serializer.serialize_event(
+            invite = self._event_serializer.serialize_event(
                 room.invite,
                 time_now,
                 token_id=token_id,
@@ -432,7 +431,7 @@ class SyncRestServlet(RestServlet):
         """
         knocked = {}
         for room in rooms:
-            knock = await self._event_serializer.serialize_event(
+            knock = self._event_serializer.serialize_event(
                 room.knock,
                 time_now,
                 token_id=token_id,
@@ -525,21 +524,14 @@ class SyncRestServlet(RestServlet):
             The room, encoded in our response format
         """
 
-        def serialize(events: Iterable[EventBase]) -> Awaitable[List[JsonDict]]:
+        def serialize(
+            events: Iterable[EventBase],
+            aggregations: Optional[Dict[str, Dict[str, Any]]] = None,
+        ) -> List[JsonDict]:
             return self._event_serializer.serialize_events(
                 events,
                 time_now=time_now,
-                # Don't bother to bundle aggregations if the timeline is unlimited,
-                # as clients will have all the necessary information.
-                # bundle_aggregations=room.timeline.limited,
-                #
-                # richvdh 2021-12-15: disable this temporarily as it has too high an
-                # overhead for initialsyncs. We need to figure out a way that the
-                # bundling can be done *before* the events are stored in the
-                # SyncResponseCache so that this part can be synchronous.
-                #
-                # Ensure to re-enable the test at tests/rest/client/test_relations.py::RelationsTestCase.test_bundled_aggregations.
-                bundle_aggregations=False,
+                bundle_aggregations=aggregations,
                 token_id=token_id,
                 event_format=event_formatter,
                 only_event_fields=only_fields,
@@ -561,8 +553,21 @@ class SyncRestServlet(RestServlet):
                     event.room_id,
                 )
 
-        serialized_state = await serialize(state_events)
-        serialized_timeline = await serialize(timeline_events)
+        serialized_state = serialize(state_events)
+        # Don't bother to bundle aggregations if the timeline is unlimited,
+        # as clients will have all the necessary information.
+        # bundle_aggregations=room.timeline.limited,
+        #
+        # richvdh 2021-12-15: disable this temporarily as it has too high an
+        # overhead for initialsyncs. We need to figure out a way that the
+        # bundling can be done *before* the events are stored in the
+        # SyncResponseCache so that this part can be synchronous.
+        #
+        # Ensure to re-enable the test at tests/rest/client/test_relations.py::RelationsTestCase.test_bundled_aggregations.
+        # if room.timeline.limited:
+        #    aggregations = await self.store.get_bundled_aggregations(timeline_events)
+        aggregations = None
+        serialized_timeline = serialize(timeline_events, aggregations)
 
         account_data = room.account_data