Fix get_users_who_share_room_with_user SQL

Three fixes here:

* Using UNION only selects DISTINCT rows from each query. However, if we DISTINCT the rows during the query, and then use UNION ALL which doesn't attempt to select distinct rows, we get a 1/3rd speedup in query time!
* We should select other_user_id instead of user_id from users_who_share_private_rooms, as user_id will always be the requesting user.
* Added p1.user_id != p2.user_id to filter out the same entries between each table in the users_in_public_rooms query.

changelog.d/9399.misc

@@ -0,0 +1 @@
+Optimise some code behind the user presence feature.

synapse/federation/sender/__init__.py

@@ -472,7 +472,7 @@ class FederationSender:
             self._processing_pending_presence = False
 
     def send_presence_to_destinations(
-        self, states: List[UserPresenceState], destinations: List[str]
+        self, states: List[UserPresenceState], destinations: Iterable[str]
     ) -> None:
         """Send the given presence states to the given destinations.
             destinations (list[str])

synapse/handlers/presence.py

@@ -909,9 +909,6 @@ class PresenceHandler(BasePresenceHandler):
             state = await self.current_state_for_user(user_id)
             hosts = await self.state.get_current_hosts_in_room(room_id)
 
-            # Filter out ourselves.
-            hosts = {host for host in hosts if host != self.server_name}
-
             self.federation.send_presence_to_destinations(
                 states=[state], destinations=hosts
             )
@@ -1115,13 +1112,11 @@ class PresenceEventSource:
         updates for
         """
         user_id = user.to_string()
-        users_interested_in = set()
-        users_interested_in.add(user_id)  # So that we receive our own presence
 
-        users_who_share_room = await self.store.get_users_who_share_room_with_user(
+        users_interested_in = await self.store.get_users_who_share_room_with_user(
             user_id, on_invalidate=cache_context.invalidate
         )
-        users_interested_in.update(users_who_share_room)
+        users_interested_in.add(user_id)  # So that we receive our own presence
 
         if explicit_room_id:
             user_ids = await self.store.get_users_in_room(

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,36 @@ 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
+
+        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 != p2.user_id
+                    AND p1.user_id = ?
+                UNION ALL
+                SELECT DISTINCT other_user_id
+                FROM users_who_share_private_rooms
+                WHERE
+                    user_id = ?
+                """,
+                (user_id, user_id),
+            )
+            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
         )
 
-        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
-            )
-            user_who_share_room.update(user_ids)
-
-        return user_who_share_room
+        return {row.get("user_id") or row.get("other_user_id") for row in rows}
 
     async def get_joined_users_from_context(
         self, event: EventBase, context: EventContext

tests/handlers/test_presence.py

@@ -585,7 +585,11 @@ class PresenceJoinTestCase(unittest.HomeserverTestCase):
         )
         self.assertEqual(expected_state.state, PresenceState.ONLINE)
         self.federation_sender.send_presence_to_destinations.assert_called_once_with(
-            destinations={"server2", "server3"}, states=[expected_state]
+            # "server" is included here as it appears when getting the current hosts for
+            # the room. send_presence_to_destinations will remove the host server before
+            # sending out presence.
+            destinations=frozenset({"server", "server2", "server3"}),
+            states=[expected_state],
         )
 
     def _add_new_user(self, room_id, user_id):