From fd1200344112eb28486ee6f82ee341ada8bb4f06 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erikj@element.io>
Date: Thu, 16 May 2024 16:07:54 +0100
Subject: [PATCH] Revert "Improve perf of sync device lists" (#17207)

Reverts element-hq/synapse#17191
---
 changelog.d/17191.misc                    |  1 -
 synapse/handlers/sync.py                  | 37 +++++++++++++++++++----
 synapse/storage/databases/main/devices.py | 17 +++++++++--
 3 files changed, 46 insertions(+), 9 deletions(-)
 delete mode 100644 changelog.d/17191.misc

diff --git a/changelog.d/17191.misc b/changelog.d/17191.misc
deleted file mode 100644
index bd55eeaa33..0000000000
--- a/changelog.d/17191.misc
+++ /dev/null
@@ -1 +0,0 @@
-Improve performance of calculating device lists changes in `/sync`.
diff --git a/synapse/handlers/sync.py b/synapse/handlers/sync.py
index 659499af75..2bd1b8de88 100644
--- a/synapse/handlers/sync.py
+++ b/synapse/handlers/sync.py
@@ -1854,13 +1854,38 @@ class SyncHandler:
 
         # Step 1a, check for changes in devices of users we share a room
         # with
-        users_that_have_changed = (
-            await self._device_handler.get_device_changes_in_shared_rooms(
-                user_id,
-                sync_result_builder.joined_room_ids,
-                from_token=since_token,
-            )
+        #
+        # We do this in two different ways depending on what we have cached.
+        # If we already have a list of all the user that have changed since
+        # the last sync then it's likely more efficient to compare the rooms
+        # they're in with the rooms the syncing user is in.
+        #
+        # If we don't have that info cached then we get all the users that
+        # share a room with our user and check if those users have changed.
+        cache_result = self.store.get_cached_device_list_changes(
+            since_token.device_list_key
         )
+        if cache_result.hit:
+            changed_users = cache_result.entities
+
+            result = await self.store.get_rooms_for_users(changed_users)
+
+            for changed_user_id, entries in result.items():
+                # Check if the changed user shares any rooms with the user,
+                # or if the changed user is the syncing user (as we always
+                # want to include device list updates of their own devices).
+                if user_id == changed_user_id or any(
+                    rid in joined_rooms for rid in entries
+                ):
+                    users_that_have_changed.add(changed_user_id)
+        else:
+            users_that_have_changed = (
+                await self._device_handler.get_device_changes_in_shared_rooms(
+                    user_id,
+                    sync_result_builder.joined_room_ids,
+                    from_token=since_token,
+                )
+            )
 
         # Step 1b, check for newly joined rooms
         for room_id in newly_joined_rooms:
diff --git a/synapse/storage/databases/main/devices.py b/synapse/storage/databases/main/devices.py
index d98f0593bc..8dbcb3f5a0 100644
--- a/synapse/storage/databases/main/devices.py
+++ b/synapse/storage/databases/main/devices.py
@@ -70,7 +70,10 @@ from synapse.types import (
 from synapse.util import json_decoder, json_encoder
 from synapse.util.caches.descriptors import cached, cachedList
 from synapse.util.caches.lrucache import LruCache
-from synapse.util.caches.stream_change_cache import StreamChangeCache
+from synapse.util.caches.stream_change_cache import (
+    AllEntitiesChangedResult,
+    StreamChangeCache,
+)
 from synapse.util.cancellation import cancellable
 from synapse.util.iterutils import batch_iter
 from synapse.util.stringutils import shortstr
@@ -829,6 +832,16 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
         )
         return {device[0]: db_to_json(device[1]) for device in devices}
 
+    def get_cached_device_list_changes(
+        self,
+        from_key: int,
+    ) -> AllEntitiesChangedResult:
+        """Get set of users whose devices have changed since `from_key`, or None
+        if that information is not in our cache.
+        """
+
+        return self._device_list_stream_cache.get_all_entities_changed(from_key)
+
     @cancellable
     async def get_all_devices_changed(
         self,
@@ -1462,7 +1475,7 @@ class DeviceWorkerStore(RoomMemberWorkerStore, EndToEndKeyWorkerStore):
 
         sql = """
             SELECT DISTINCT user_id FROM device_lists_changes_in_room
-            WHERE {clause} AND stream_id > ?
+            WHERE {clause} AND stream_id >= ?
         """
 
         def _get_device_list_changes_in_rooms_txn(