possible perf boost to calculating device list update recipients

2 files changed, 9 insertions, 9 deletions

diff --git a/synapse/handlers/device.py b/synapse/handlers/device.py
index 82ee11e921..2c07d31dfd 100644
--- a/synapse/handlers/device.py
+++ b/synapse/handlers/device.py
@@ -495,13 +495,11 @@ class DeviceHandler(DeviceWorkerHandler):
                 "Notifying about update %r/%r, ID: %r", user_id, device_id, position
             )
 
-        room_ids = await self.store.get_rooms_for_user(user_id)
-
         # specify the user ID too since the user should always get their own device list
         # updates, even if they aren't in any rooms.
-        self.notifier.on_new_event(
-            "device_list_key", position, users=[user_id], rooms=room_ids
-        )
+        users_to_notify = users_who_share_room.union(user_id)
+
+        self.notifier.on_new_event("device_list_key", position, users=users_to_notify)
 
         if hosts:
             logger.info(
diff --git a/synapse/replication/tcp/client.py b/synapse/replication/tcp/client.py
index e29ae1e375..679df5602f 100644
--- a/synapse/replication/tcp/client.py
+++ b/synapse/replication/tcp/client.py
@@ -173,12 +173,14 @@ class ReplicationDataHandler:
             if entities:
                 self.notifier.on_new_event("to_device_key", token, users=entities)
         elif stream_name == DeviceListsStream.NAME:
-            all_room_ids: Set[str] = set()
+            users_to_notify: Set[str] = set()
             for row in rows:
                 if row.entity.startswith("@"):
-                    room_ids = await self.store.get_rooms_for_user(row.entity)
-                    all_room_ids.update(room_ids)
-            self.notifier.on_new_event("device_list_key", token, rooms=all_room_ids)
+                    user_ids = await self.store.get_users_who_share_room_with_user(
+                        row.entity
+                    )
+                    users_to_notify.update(user_ids)
+            self.notifier.on_new_event("device_list_key", token, users=users_to_notify)
         elif stream_name == GroupServerStream.NAME:
             self.notifier.on_new_event(
                 "groups_key", token, users=[row.user_id for row in rows]