Fix performance regression in `get_users_in_room` (#13972)

Fixes #13942. Introduced in #13575.

Basically, let's only get the ordered set of hosts out of the DB if we need an ordered set of hosts. Since we split the function up the caching won't be as good, but I think it will still be fine as e.g. multiple backfill requests for the same room will hit the cache.

1 files changed, 18 insertions, 12 deletions

diff --git a/synapse/storage/controllers/state.py b/synapse/storage/controllers/state.py
index bb60130afe..2b31ce54bb 100644
--- a/synapse/storage/controllers/state.py
+++ b/synapse/storage/controllers/state.py
@@ -23,7 +23,7 @@ from typing import (
     List,
     Mapping,
     Optional,
-    Sequence,
+    Set,
     Tuple,
 )
 
@@ -529,7 +529,18 @@ class StateStorageController:
         )
         return state_map.get(key)
 
-    async def get_current_hosts_in_room(self, room_id: str) -> List[str]:
+    async def get_current_hosts_in_room(self, room_id: str) -> Set[str]:
+        """Get current hosts in room based on current state.
+
+        Blocks until we have full state for the given room. This only happens for rooms
+        with partial state.
+        """
+
+        await self._partial_state_room_tracker.await_full_state(room_id)
+
+        return await self.stores.main.get_current_hosts_in_room(room_id)
+
+    async def get_current_hosts_in_room_ordered(self, room_id: str) -> List[str]:
         """Get current hosts in room based on current state.
 
         Blocks until we have full state for the given room. This only happens for rooms
@@ -542,11 +553,11 @@ class StateStorageController:
 
         await self._partial_state_room_tracker.await_full_state(room_id)
 
-        return await self.stores.main.get_current_hosts_in_room(room_id)
+        return await self.stores.main.get_current_hosts_in_room_ordered(room_id)
 
     async def get_current_hosts_in_room_or_partial_state_approximation(
         self, room_id: str
-    ) -> Sequence[str]:
+    ) -> Collection[str]:
         """Get approximation of current hosts in room based on current state.
 
         For rooms with full state, this is equivalent to `get_current_hosts_in_room`,
@@ -566,14 +577,9 @@ class StateStorageController:
         )
 
         hosts_from_state = await self.stores.main.get_current_hosts_in_room(room_id)
-        hosts_from_state_set = set(hosts_from_state)
-
-        # First take the list of hosts based on the current state.
-        # For rooms with partial state, this will be missing most hosts.
-        hosts = list(hosts_from_state)
-        # Then add in the list of hosts in the room at the time we joined.
-        # This will be an empty list for rooms with full state.
-        hosts.extend(host for host in hosts_at_join if host not in hosts_from_state_set)
+
+        hosts = set(hosts_at_join)
+        hosts.update(hosts_from_state)
 
         return hosts