User directory: use calculated room membership state instead (#9821)

Fixes: #9797.

Should help reduce CPU usage on the user directory, especially when memberships change in rooms with lots of state history.

2 files changed, 35 insertions, 7 deletions

diff --git a/synapse/handlers/user_directory.py b/synapse/handlers/user_directory.py
index 9b1e6d5c18..dacc4f3076 100644
--- a/synapse/handlers/user_directory.py
+++ b/synapse/handlers/user_directory.py
@@ -44,7 +44,6 @@ class UserDirectoryHandler(StateDeltasHandler):
         super().__init__(hs)
 
         self.store = hs.get_datastore()
-        self.state = hs.get_state_handler()
         self.server_name = hs.hostname
         self.clock = hs.get_clock()
         self.notifier = hs.get_notifier()
@@ -302,10 +301,12 @@ class UserDirectoryHandler(StateDeltasHandler):
             # ignore the change
             return
 
-        users_with_profile = await self.state.get_current_users_in_room(room_id)
+        other_users_in_room_with_profiles = (
+            await self.store.get_users_in_room_with_profiles(room_id)
+        )
 
         # Remove every user from the sharing tables for that room.
-        for user_id in users_with_profile.keys():
+        for user_id in other_users_in_room_with_profiles.keys():
             await self.store.remove_user_who_share_room(user_id, room_id)
 
         # Then, re-add them to the tables.
@@ -314,7 +315,7 @@ class UserDirectoryHandler(StateDeltasHandler):
         # which when ran over an entire room, will result in the same values
         # being added multiple times. The batching upserts shouldn't make this
         # too bad, though.
-        for user_id, profile in users_with_profile.items():
+        for user_id, profile in other_users_in_room_with_profiles.items():
             await self._handle_new_user(room_id, user_id, profile)
 
     async def _handle_new_user(
@@ -336,7 +337,7 @@ class UserDirectoryHandler(StateDeltasHandler):
             room_id
         )
         # Now we update users who share rooms with users.
-        users_with_profile = await self.state.get_current_users_in_room(room_id)
+        other_users_in_room = await self.store.get_users_in_room(room_id)
 
         if is_public:
             await self.store.add_users_in_public_rooms(room_id, (user_id,))
@@ -352,14 +353,14 @@ class UserDirectoryHandler(StateDeltasHandler):
 
                 # We don't care about appservice users.
                 if not is_appservice:
-                    for other_user_id in users_with_profile:
+                    for other_user_id in other_users_in_room:
                         if user_id == other_user_id:
                             continue
 
                         to_insert.add((user_id, other_user_id))
 
             # Next we need to update for every local user in the room
-            for other_user_id in users_with_profile:
+            for other_user_id in other_users_in_room:
                 if user_id == other_user_id:
                     continue
 
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index ef5587f87a..fd525dce65 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -173,6 +173,33 @@ class RoomMemberWorkerStore(EventsWorkerStore):
         txn.execute(sql, (room_id, Membership.JOIN))
         return [r[0] for r in txn]
 
+    @cached(max_entries=100000, iterable=True)
+    async def get_users_in_room_with_profiles(
+        self, room_id: str
+    ) -> Dict[str, ProfileInfo]:
+        """Get a mapping from user ID to profile information for all users in a given room.
+
+        Args:
+            room_id: The ID of the room to retrieve the users of.
+
+        Returns:
+            A mapping from user ID to ProfileInfo.
+        """
+
+        def _get_users_in_room_with_profiles(txn) -> Dict[str, ProfileInfo]:
+            sql = """
+                SELECT user_id, display_name, avatar_url FROM room_memberships
+                WHERE room_id = ? AND membership = ?
+            """
+            txn.execute(sql, (room_id, Membership.JOIN))
+
+            return {r[0]: ProfileInfo(display_name=r[1], avatar_url=r[2]) for r in txn}
+
+        return await self.db_pool.runInteraction(
+            "get_users_in_room_with_profiles",
+            _get_users_in_room_with_profiles,
+        )
+
     @cached(max_entries=100000)
     async def get_room_summary(self, room_id: str) -> Dict[str, MemberSummary]:
         """Get the details of a room roughly suitable for use by the room