Batch up calls to `get_rooms_for_users` (#14109)

pull/14155/head
Nick Mills-Barrett 2022-10-12 11:36:22 +01:00 committed by GitHub
parent 09be8ab5f9
commit f9bc5428c4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 17 additions and 1 deletions

1
changelog.d/14109.misc Normal file
View File

@ -0,0 +1 @@
Break up calls to fetch rooms for many users. Contributed by Nick @ Beeper (@fizzadar).

View File

@ -666,7 +666,7 @@ class RoomMemberWorkerStore(EventsWorkerStore):
cached_method_name="get_rooms_for_user", cached_method_name="get_rooms_for_user",
list_name="user_ids", list_name="user_ids",
) )
async def get_rooms_for_users( async def _get_rooms_for_users(
self, user_ids: Collection[str] self, user_ids: Collection[str]
) -> Dict[str, FrozenSet[str]]: ) -> Dict[str, FrozenSet[str]]:
"""A batched version of `get_rooms_for_user`. """A batched version of `get_rooms_for_user`.
@ -697,6 +697,21 @@ class RoomMemberWorkerStore(EventsWorkerStore):
return {key: frozenset(rooms) for key, rooms in user_rooms.items()} return {key: frozenset(rooms) for key, rooms in user_rooms.items()}
async def get_rooms_for_users(
self, user_ids: Collection[str]
) -> Dict[str, FrozenSet[str]]:
"""A batched wrapper around `_get_rooms_for_users`, to prevent locking
other calls to `get_rooms_for_user` for large user lists.
"""
all_user_rooms: Dict[str, FrozenSet[str]] = {}
# 250 users is pretty arbitrary but the data can be quite large if users
# are in many rooms.
for user_ids in batch_iter(user_ids, 250):
all_user_rooms.update(await self._get_rooms_for_users(user_ids))
return all_user_rooms
@cached(max_entries=10000) @cached(max_entries=10000)
async def does_pair_of_users_share_a_room( async def does_pair_of_users_share_a_room(
self, user_id: str, other_user_id: str self, user_id: str, other_user_id: str