From a5c50548b07fef80e00a5b7d56f0aed5a55c262c Mon Sep 17 00:00:00 2001 From: Mathieu Velten Date: Wed, 31 May 2023 12:07:04 +0200 Subject: [PATCH] Simplify code --- synapse/handlers/room_list.py | 106 ++++++++++--------------- synapse/storage/databases/main/room.py | 14 ++-- 2 files changed, 47 insertions(+), 73 deletions(-) diff --git a/synapse/handlers/room_list.py b/synapse/handlers/room_list.py index a0742cbc86..664bfbdacf 100644 --- a/synapse/handlers/room_list.py +++ b/synapse/handlers/room_list.py @@ -13,7 +13,7 @@ # limitations under the License. import logging -from typing import TYPE_CHECKING, Any, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import attr import msgpack @@ -170,66 +170,8 @@ class RoomListHandler: # we request one more than wanted to see if there are more pages to come probing_limit = limit + 1 if limit is not None else None - results: List[PublicRoom] = [] - - # print(f"{forwards} {last_joined_members} {last_room_id} {last_module_index}") - - def insert_into_result( - new_room: PublicRoom, module_index: Optional[int] - ) -> None: - # print(f"insert {new_room.room_id} {module_index}") - if new_room.num_joined_members == last_joined_members: - if last_module_index is not None and last_room_id is not None: - if module_index is not None and module_index > last_module_index: - return - inserted = False - for i, r in enumerate(results): - r = results[i] - if ( - forwards and new_room.num_joined_members >= r.num_joined_members - ) or ( - not forwards and new_room.num_joined_members <= r.num_joined_members - ): - results.insert(i, new_room) - inserted = True - return - if not inserted: - if forwards: - results.append(new_room) - else: - results.insert(0, new_room) - - room_ids_to_module_index = {} - - for module_index, fetch_public_rooms in enumerate( - self._module_api_callbacks.fetch_public_rooms_callbacks - ): - # Ask each module for a list of public rooms given the last_joined_members - # value from the since token and the probing limit. - module_last_joined_members = None - if last_joined_members is not None: - module_last_joined_members = last_joined_members - if last_module_index is not None and last_module_index < module_index: - module_last_joined_members = module_last_joined_members - 1 - module_public_rooms = await fetch_public_rooms( - network_tuple, - search_filter, - probing_limit, - ( - module_last_joined_members, - last_room_id if last_module_index == module_index else None, - ), - forwards, - ) - - # We reverse for iteration to keep the order in the final list - # since we preprend when inserting - module_public_rooms.reverse() - - # Insert the module's reported public rooms into the list - for new_room in module_public_rooms: - room_ids_to_module_index[new_room.room_id] = module_index - insert_into_result(new_room, module_index) + num_joined_members_buckets: Dict[int, List[PublicRoom]] = {} + room_ids_to_module_index: Dict[str, int] = {} local_public_rooms = await self.store.get_largest_public_rooms( network_tuple, @@ -243,8 +185,44 @@ class RoomListHandler: ignore_non_federatable=bool(from_remote_server_name), ) - for r in local_public_rooms: - insert_into_result(r, None) + for room in local_public_rooms: + num_joined_members_buckets.setdefault(room.num_joined_members, []).append( + room + ) + + for module_index, fetch_public_rooms in enumerate( + self._module_api_callbacks.fetch_public_rooms_callbacks + ): + # Ask each module for a list of public rooms given the last_joined_members + # value from the since token and the probing limit. + module_last_joined_members = None + if last_joined_members is not None: + module_last_joined_members = last_joined_members + if last_module_index is not None and module_index < last_module_index: + module_last_joined_members = module_last_joined_members - 1 + module_public_rooms = await fetch_public_rooms( + network_tuple, + search_filter, + probing_limit, + ( + module_last_joined_members, + last_room_id if last_module_index == module_index else None, + ), + forwards, + ) + + for room in module_public_rooms: + num_joined_members_buckets.setdefault( + room.num_joined_members, [] + ).append(room) + room_ids_to_module_index[room.room_id] = module_index + + nums_joined_members = list(num_joined_members_buckets.keys()) + nums_joined_members.sort(reverse=forwards) + + results = [] + for num_joined_members in nums_joined_members: + results += num_joined_members_buckets[num_joined_members] response: JsonDict = {} num_results = len(results) @@ -258,8 +236,6 @@ class RoomListHandler: if not forwards: results.reverse() - # print("final ", [(r.room_id, r.num_joined_members) for r in results]) - if num_results > 0: final_entry = results[-1] initial_entry = results[0] diff --git a/synapse/storage/databases/main/room.py b/synapse/storage/databases/main/room.py index 0dcce4fc18..932293ef2e 100644 --- a/synapse/storage/databases/main/room.py +++ b/synapse/storage/databases/main/room.py @@ -386,7 +386,7 @@ class RoomWorkerStore(CacheInvalidationWorkerStore): bounds: An uppoer or lower bound to apply to result set if given, consists of a joined member count and room_id (these are excluded from result set). - forwards: true iff going forwards, going backwards otherwise + forwards: true if going forwards, going backwards otherwise ignore_non_federatable: If true filters out non-federatable rooms. Returns: @@ -424,14 +424,12 @@ class RoomWorkerStore(CacheInvalidationWorkerStore): if last_joined_members is not None: comp = "<" if forwards else ">" - clause = f"joined_members {comp}= ? AND (joined_members {comp} ?" - query_args += [last_joined_members, last_joined_members] + clause = f"joined_members {comp} ?" + query_args += [last_joined_members] - if last_room_id is None: - clause += ")" - else: - clause += f"OR room_id {comp} ?)" - query_args.append(last_room_id) + if last_room_id is not None: + clause += f" AND (joined_members {comp} ? OR room_id {comp} ?)" + query_args += [last_joined_members, last_room_id] where_clauses.append(clause)