Speed up get_users_who_share_room_with_user with a more efficient query

The old code was pulling all rooms for a given user, *then* looping over each
one in Python and pulling all users from those rooms.

This commit replaces that with a single query which makes use of existing
tables which keep track of user relationships.

2 files changed, 29 insertions, 20 deletions

diff --git a/synapse/handlers/presence.py b/synapse/handlers/presence.py
index 22d1e9d35c..15900eb844 100644
--- a/synapse/handlers/presence.py
+++ b/synapse/handlers/presence.py
@@ -1114,14 +1114,12 @@ class PresenceEventSource:
         """Returns the set of users that the given user should see presence
         updates for
         """
-        user_id = user.to_string()
-        users_interested_in = set()
-        users_interested_in.add(user_id)  # So that we receive our own presence
+        user_id = user.to_string
 
-        users_who_share_room = await self.store.get_users_who_share_room_with_user(
-            user_id, on_invalidate=cache_context.invalidate
+        users_interested_in = await self.store.get_users_who_share_room_with_user(
+            user_id,
         )
-        users_interested_in.update(users_who_share_room)
+        users_interested_in.update(user_id)  # So that we receive our own presence
 
         if explicit_room_id:
             user_ids = await self.store.get_users_in_room(
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index 92382bed28..fbedc18191 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -37,7 +37,7 @@ from synapse.storage.roommember import (
 from synapse.types import Collection, PersistedEventPosition, get_domain_from_id
 from synapse.util.async_helpers import Linearizer
 from synapse.util.caches import intern_string
-from synapse.util.caches.descriptors import _CacheContext, cached, cachedList
+from synapse.util.caches.descriptors import cached, cachedList
 from synapse.util.metrics import Measure
 
 if TYPE_CHECKING:
@@ -484,24 +484,35 @@ class RoomMemberWorkerStore(EventsWorkerStore):
         )
         return frozenset(r.room_id for r in rooms)
 
-    @cached(max_entries=500000, cache_context=True, iterable=True)
-    async def get_users_who_share_room_with_user(
-        self, user_id: str, cache_context: _CacheContext
-    ) -> Set[str]:
+    @cached(max_entries=500000, iterable=True)
+    async def get_users_who_share_room_with_user(self, user_id: str) -> Set[str]:
         """Returns the set of users who share a room with `user_id`
         """
-        room_ids = await self.get_rooms_for_user(
-            user_id, on_invalidate=cache_context.invalidate
-        )
 
-        user_who_share_room = set()
-        for room_id in room_ids:
-            user_ids = await self.get_users_in_room(
-                room_id, on_invalidate=cache_context.invalidate
+        def _get_users_who_share_room_with_user(txn):
+            txn.execute(
+                """
+                SELECT DISTINCT p2.user_id
+                FROM users_in_public_rooms as p1
+                INNER JOIN users_in_public_rooms as p2
+                    ON p1.room_id = p2.room_id
+                    AND p1.user_id = ?
+                UNION
+                SELECT DISTINCT user_id
+                FROM users_who_share_private_rooms
+                WHERE
+                    user_id = ?
+                """,
+                (user_id, user_id),
             )
-            user_who_share_room.update(user_ids)
+            rows = self.db_pool.cursor_to_dict(txn)
+            return rows
+
+        rows = await self.db_pool.runInteraction(
+            "get_users_who_share_room_with_user", _get_users_who_share_room_with_user
+        )
 
-        return user_who_share_room
+        return {row["user_id"] for row in rows}
 
     async def get_joined_users_from_context(
         self, event: EventBase, context: EventContext