Address review comments

Author: richvdh

synapse/handlers/federation.py

@@ -576,7 +576,9 @@ async def _get_state_for_room(
 
         # Fetch the state events from the DB, and check we have the auth events.
         event_map = await self.store.get_events(state_event_ids, allow_rejected=True)
-        auth_events_in_store = await self.store.have_seen_events(auth_event_ids)
+        auth_events_in_store = await self.store.have_seen_events(
+            room_id, auth_event_ids
+        )
 
         # Check for missing events. We handle state and auth event seperately,
         # as we want to pull the state from the DB, but we don't for the auth
@@ -609,7 +611,7 @@ async def _get_state_for_room(
 
             if missing_auth_events:
                 auth_events_in_store = await self.store.have_seen_events(
-                    missing_auth_events
+                    room_id, missing_auth_events
                 )
                 missing_auth_events.difference_update(auth_events_in_store)
 
@@ -709,7 +711,7 @@ async def _get_state_after_missing_prev_event(
 
         missing_auth_events = set(auth_event_ids) - fetched_events.keys()
         missing_auth_events.difference_update(
-            await self.store.have_seen_events(missing_auth_events)
+            await self.store.have_seen_events(room_id, missing_auth_events)
         )
         logger.debug("We are also missing %i auth events", len(missing_auth_events))
 
@@ -2474,7 +2476,7 @@ async def _update_auth_events_and_context_for_auth(
         #
         # we start by checking if they are in the store, and then try calling /event_auth/.
         if missing_auth:
-            have_events = await self.store.have_seen_events(missing_auth)
+            have_events = await self.store.have_seen_events(event.room_id, missing_auth)
             logger.debug("Events %s are in the store", have_events)
             missing_auth.difference_update(have_events)
 
@@ -2493,7 +2495,7 @@ async def _update_auth_events_and_context_for_auth(
                     return context
 
                 seen_remotes = await self.store.have_seen_events(
-                    [e.event_id for e in remote_auth_chain]
+                    event.room_id, [e.event_id for e in remote_auth_chain]
                 )
 
                 for e in remote_auth_chain:

synapse/storage/databases/main/cache.py

@@ -168,7 +168,7 @@ def _invalidate_caches_for_event(
         backfilled,
     ):
         self._invalidate_get_event_cache(event_id)
-        self.have_seen_event.prefill((event_id,), True)
+        self.have_seen_event.invalidate((room_id, event_id))
 
         self.get_latest_event_ids_in_room.invalidate((room_id,))
 

synapse/storage/databases/main/events_worker.py

@@ -1046,51 +1046,70 @@ async def have_events_in_timeline(self, event_ids):
 
         return {r["event_id"] for r in rows}
 
-    async def have_seen_events(self, event_ids: Collection[str]) -> Set[str]:
+    async def have_seen_events(
+        self, room_id: str, event_ids: Iterable[str]
+    ) -> Set[str]:
         """Given a list of event ids, check if we have already processed them.
 
+        The room_id is only used to structure the cache (so that it can later be
+        invalidated by room_id) - there is no guarantee that the events are actually
+        in the room in question.
+
         Args:
-            event_ids (iterable[str]):
+            room_id: Room we are polling
+            event_ids: events we are looking for
 
         Returns:
             set[str]: The events we have already seen.
         """
-        res = await self._have_seen_events_dict(event_ids)
-        return {x for (x, y) in res.items() if y}
+        res = await self._have_seen_events_dict(
+            (room_id, event_id) for event_id in event_ids
+        )
+        return {eid for ((_rid, eid), have_event) in res.items() if have_event}
 
-    @cachedList("have_seen_event", "event_ids")
+    @cachedList("have_seen_event", "keys")
     async def _have_seen_events_dict(
-        self, event_ids: Collection[str]
-    ) -> Dict[str, bool]:
+        self, keys: Iterable[Tuple[str, str]]
+    ) -> Dict[Tuple[str, str], bool]:
         """Helper for have_seen_events
 
-        Returns a dict, which is the right format for @cachedList
+        Returns:
+             a dict {(room_id, event_id)-> bool}
         """
         # if the event cache contains the event, obviously we've seen it.
-        cache_results = {x for x in event_ids if self._get_event_cache.contains(x)}
+        cache_results = {
+            eid for (_rid, eid) in keys if self._get_event_cache.contains(eid)
+        }
         results = {x: True for x in cache_results}
 
-        def have_seen_events_txn(txn, chunk):
+        def have_seen_events_txn(txn, chunk: Tuple[str, ...]):
             # assume everything in this chunk is not found initially
             results.update({x: False for x in chunk})
 
             # check the db and update the results for any row that is found
-            sql = "SELECT event_id FROM events as e WHERE "
+            # NB this deliberately does *not* query on room_id, to make this an
+            #   index-only lookup on `events_event_id_key`.
+            sql = "SELECT event_id FROM events AS e WHERE"
             clause, args = make_in_list_sql_clause(
                 txn.database_engine, "e.event_id", chunk
             )
             txn.execute(sql + clause, args)
             results.update({row[0]: True for row in txn})
 
-        for chunk in batch_iter((x for x in event_ids if x not in cache_results), 100):
+        # each batch requires its own index scan, so we make the batches as big as
+        # possible.
+        for chunk in batch_iter(
+            (eid for (_rid, eid) in keys if eid not in cache_results),
+            500,
+        ):
             await self.db_pool.runInteraction(
                 "have_seen_events", have_seen_events_txn, chunk
             )
 
         return results
 
-    @cached(max_entries=100000)
-    async def have_seen_event(self, event_id):
+    @cached(max_entries=100000, tree=True)
+    async def have_seen_event(self, room_id: str, event_id: str):
         # this only exists for the benefit of the @cachedList descriptor on
         # _have_seen_events_dict
         raise NotImplementedError()

synapse/storage/databases/main/purge_events.py

@@ -286,11 +286,13 @@ def _purge_history_txn(
                 txn, self._get_state_group_for_event, (event_id,)
             )
 
-            # FIXME: this is racy - what if have_seen_event gets called between the
-            #    transaction completing and the invalidation running?
+            # XXX: This is racy, since have_seen_events could be called between the
+            #    transaction completing and the invalidation running. On the other hand,
+            #    that's no different to calling `have_seen_events` just before the
+            #    event is deleted from the database.
             if should_delete:
                 self._invalidate_cache_and_stream(
-                    txn, self.have_seen_event, (event_id,)
+                    txn, self.have_seen_event, (room_id, event_id)
                 )
 
         logger.info("[purge] done")
@@ -434,12 +436,9 @@ def _purge_room_txn(self, txn, room_id: str) -> List[int]:
 
         # TODO: we could probably usefully do a bunch more cache invalidation here
 
-        # we have no way to know which events to clear out of have_seen_event
-        # so just have to drop the whole cache
-        #
-        # FIXME: this is racy - what if have_seen_event gets called between the
-        #    DELETE completing and the invalidation running?
-        self._invalidate_all_cache_and_stream(txn, self.have_seen_event)
+        # XXX: as with purge_history, this is racy, but no worse than other races
+        #   that already exist.
+        self._invalidate_cache_and_stream(txn, self.have_seen_event, (room_id,))
 
         logger.info("[purge] done")