Don't invalidate all `get_relations_for_event` on history purge (#17083)

This is a tree cache already, so may as well move the room ID to the
front and use that

5 files changed, 40 insertions, 13 deletions

diff --git a/synapse/handlers/relations.py b/synapse/handlers/relations.py
index c5cee8860b..de092f8623 100644
--- a/synapse/handlers/relations.py
+++ b/synapse/handlers/relations.py
@@ -393,9 +393,9 @@ class RelationsHandler:
 
                 # Attempt to find another event to use as the latest event.
                 potential_events, _ = await self._main_store.get_relations_for_event(
+                    room_id,
                     event_id,
                     event,
-                    room_id,
                     RelationTypes.THREAD,
                     direction=Direction.FORWARDS,
                 )
diff --git a/synapse/storage/databases/main/cache.py b/synapse/storage/databases/main/cache.py
index bfd492d95d..c6787faea0 100644
--- a/synapse/storage/databases/main/cache.py
+++ b/synapse/storage/databases/main/cache.py
@@ -318,7 +318,13 @@ class CacheInvalidationWorkerStore(SQLBaseStore):
             self._invalidate_local_get_event_cache(redacts)  # type: ignore[attr-defined]
             # Caches which might leak edits must be invalidated for the event being
             # redacted.
-            self._attempt_to_invalidate_cache("get_relations_for_event", (redacts,))
+            self._attempt_to_invalidate_cache(
+                "get_relations_for_event",
+                (
+                    room_id,
+                    redacts,
+                ),
+            )
             self._attempt_to_invalidate_cache("get_applicable_edit", (redacts,))
             self._attempt_to_invalidate_cache("get_thread_id", (redacts,))
             self._attempt_to_invalidate_cache("get_thread_id_for_receipts", (redacts,))
@@ -345,7 +351,13 @@ class CacheInvalidationWorkerStore(SQLBaseStore):
             )
 
         if relates_to:
-            self._attempt_to_invalidate_cache("get_relations_for_event", (relates_to,))
+            self._attempt_to_invalidate_cache(
+                "get_relations_for_event",
+                (
+                    room_id,
+                    relates_to,
+                ),
+            )
             self._attempt_to_invalidate_cache("get_references_for_event", (relates_to,))
             self._attempt_to_invalidate_cache("get_applicable_edit", (relates_to,))
             self._attempt_to_invalidate_cache("get_thread_summary", (relates_to,))
@@ -380,9 +392,9 @@ class CacheInvalidationWorkerStore(SQLBaseStore):
         self._attempt_to_invalidate_cache(
             "get_unread_event_push_actions_by_room_for_user", (room_id,)
         )
+        self._attempt_to_invalidate_cache("get_relations_for_event", (room_id,))
 
         self._attempt_to_invalidate_cache("_get_membership_from_event_id", None)
-        self._attempt_to_invalidate_cache("get_relations_for_event", None)
         self._attempt_to_invalidate_cache("get_applicable_edit", None)
         self._attempt_to_invalidate_cache("get_thread_id", None)
         self._attempt_to_invalidate_cache("get_thread_id_for_receipts", None)
diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py
index 990698aa5c..fd7167904d 100644
--- a/synapse/storage/databases/main/events.py
+++ b/synapse/storage/databases/main/events.py
@@ -1923,7 +1923,12 @@ class PersistEventsStore:
 
         # Any relation information for the related event must be cleared.
         self.store._invalidate_cache_and_stream(
-            txn, self.store.get_relations_for_event, (redacted_relates_to,)
+            txn,
+            self.store.get_relations_for_event,
+            (
+                room_id,
+                redacted_relates_to,
+            ),
         )
         if rel_type == RelationTypes.REFERENCE:
             self.store._invalidate_cache_and_stream(
diff --git a/synapse/storage/databases/main/events_bg_updates.py b/synapse/storage/databases/main/events_bg_updates.py
index 6c979f9f2c..64d303e330 100644
--- a/synapse/storage/databases/main/events_bg_updates.py
+++ b/synapse/storage/databases/main/events_bg_updates.py
@@ -1181,7 +1181,7 @@ class EventsBackgroundUpdatesStore(SQLBaseStore):
 
             results = list(txn)
             # (event_id, parent_id, rel_type) for each relation
-            relations_to_insert: List[Tuple[str, str, str]] = []
+            relations_to_insert: List[Tuple[str, str, str, str]] = []
             for event_id, event_json_raw in results:
                 try:
                     event_json = db_to_json(event_json_raw)
@@ -1214,7 +1214,8 @@ class EventsBackgroundUpdatesStore(SQLBaseStore):
                 if not isinstance(parent_id, str):
                     continue
 
-                relations_to_insert.append((event_id, parent_id, rel_type))
+                room_id = event_json["room_id"]
+                relations_to_insert.append((room_id, event_id, parent_id, rel_type))
 
             # Insert the missing data, note that we upsert here in case the event
             # has already been processed.
@@ -1223,18 +1224,27 @@ class EventsBackgroundUpdatesStore(SQLBaseStore):
                     txn=txn,
                     table="event_relations",
                     key_names=("event_id",),
-                    key_values=[(r[0],) for r in relations_to_insert],
+                    key_values=[(r[1],) for r in relations_to_insert],
                     value_names=("relates_to_id", "relation_type"),
-                    value_values=[r[1:] for r in relations_to_insert],
+                    value_values=[r[2:] for r in relations_to_insert],
                 )
 
                 # Iterate the parent IDs and invalidate caches.
-                cache_tuples = {(r[1],) for r in relations_to_insert}
                 self._invalidate_cache_and_stream_bulk(  # type: ignore[attr-defined]
-                    txn, self.get_relations_for_event, cache_tuples  # type: ignore[attr-defined]
+                    txn,
+                    self.get_relations_for_event,  # type: ignore[attr-defined]
+                    {
+                        (
+                            r[0],  # room_id
+                            r[2],  # parent_id
+                        )
+                        for r in relations_to_insert
+                    },
                 )
                 self._invalidate_cache_and_stream_bulk(  # type: ignore[attr-defined]
-                    txn, self.get_thread_summary, cache_tuples  # type: ignore[attr-defined]
+                    txn,
+                    self.get_thread_summary,  # type: ignore[attr-defined]
+                    {(r[1],) for r in relations_to_insert},
                 )
 
             if results:
diff --git a/synapse/storage/databases/main/relations.py b/synapse/storage/databases/main/relations.py
index 77f3641525..29a001ff92 100644
--- a/synapse/storage/databases/main/relations.py
+++ b/synapse/storage/databases/main/relations.py
@@ -169,9 +169,9 @@ class RelationsWorkerStore(SQLBaseStore):
     @cached(uncached_args=("event",), tree=True)
     async def get_relations_for_event(
         self,
+        room_id: str,
         event_id: str,
         event: EventBase,
-        room_id: str,
         relation_type: Optional[str] = None,
         event_type: Optional[str] = None,
         limit: int = 5,