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.
This commit is contained in:
Andrew Morgan 2021-04-16 18:17:18 +01:00 committed by GitHub
parent 601b893352
commit c571736c6c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 36 additions and 7 deletions

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

@ -0,0 +1 @@
Reduce CPU usage of the user directory by reusing existing calculated room membership.

View File

@ -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

View File

@ -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