Experimental support to include bundled aggregations in search results (MSC3666) (#11837)

1 files changed, 11 insertions, 2 deletions

diff --git a/synapse/storage/databases/main/relations.py b/synapse/storage/databases/main/relations.py
index 6180b17296..7718acbf1c 100644
--- a/synapse/storage/databases/main/relations.py
+++ b/synapse/storage/databases/main/relations.py
@@ -715,6 +715,9 @@ class RelationsWorkerStore(SQLBaseStore):
             A map of event ID to the bundled aggregation for the event. Not all
             events may have bundled aggregations in the results.
         """
+        # The already processed event IDs. Tracked separately from the result
+        # since the result omits events which do not have bundled aggregations.
+        seen_event_ids = set()
 
         # State events and redacted events do not get bundled aggregations.
         events = [
@@ -728,13 +731,19 @@ class RelationsWorkerStore(SQLBaseStore):
 
         # Fetch other relations per event.
         for event in events:
+            # De-duplicate events by ID to handle the same event requested multiple
+            # times. The caches that _get_bundled_aggregation_for_event use should
+            # capture this, but best to reduce work.
+            if event.event_id in seen_event_ids:
+                continue
+            seen_event_ids.add(event.event_id)
+
             event_result = await self._get_bundled_aggregation_for_event(event, user_id)
             if event_result:
                 results[event.event_id] = event_result
 
         # Fetch any edits.
-        event_ids = [event.event_id for event in events]
-        edits = await self._get_applicable_edits(event_ids)
+        edits = await self._get_applicable_edits(seen_event_ids)
         for event_id, edit in edits.items():
             results.setdefault(event_id, BundledAggregations()).replace = edit