From e8a96bcc399dc84b44251d64334e68749c93844c Mon Sep 17 00:00:00 2001
From: Patrick Cloke <patrickc@matrix.org>
Date: Wed, 25 Jan 2023 15:33:01 -0500
Subject: [PATCH] Batch look-ups to see if rooms are partial stated.

---
 changelog.d/14917.misc                 |  1 +
 synapse/handlers/sync.py               | 12 ++++++------
 synapse/storage/databases/main/room.py | 23 +++++++++++++++++++++--
 3 files changed, 28 insertions(+), 8 deletions(-)
 create mode 100644 changelog.d/14917.misc

diff --git a/changelog.d/14917.misc b/changelog.d/14917.misc
new file mode 100644
index 0000000000..8a248495b2
--- /dev/null
+++ b/changelog.d/14917.misc
@@ -0,0 +1 @@
+Improve performance of looking up partial-state status.
diff --git a/synapse/handlers/sync.py b/synapse/handlers/sync.py
index ee11764567..396deb39ce 100644
--- a/synapse/handlers/sync.py
+++ b/synapse/handlers/sync.py
@@ -1383,9 +1383,9 @@ class SyncHandler:
         if not sync_config.filter_collection.lazy_load_members():
             # Non-lazy syncs should never include partially stated rooms.
             # Exclude all partially stated rooms from this sync.
-            for room_id in mutable_joined_room_ids:
-                if await self.store.is_partial_state_room(room_id):
-                    mutable_rooms_to_exclude.add(room_id)
+            results = await self.store.is_partial_state_rooms(mutable_joined_room_ids)
+            for room_id, result in results.items():
+                mutable_rooms_to_exclude.add(room_id)
 
         # Incremental eager syncs should additionally include rooms that
         # - we are joined to
@@ -1401,9 +1401,9 @@ class SyncHandler:
                     mutable_joined_room_ids,
                 )
             )
-            for room_id in un_partial_stated_rooms:
-                if not await self.store.is_partial_state_room(room_id):
-                    forced_newly_joined_room_ids.add(room_id)
+            results = await self.store.is_partial_state_rooms(un_partial_stated_rooms)
+            for room_id, result in results.items():
+                forced_newly_joined_room_ids.add(room_id)
 
         # Now we have our list of joined room IDs, exclude as configured and freeze
         joined_room_ids = frozenset(
diff --git a/synapse/storage/databases/main/room.py b/synapse/storage/databases/main/room.py
index 3aa7b94560..aa7e2d6e55 100644
--- a/synapse/storage/databases/main/room.py
+++ b/synapse/storage/databases/main/room.py
@@ -60,9 +60,9 @@ from synapse.storage.util.id_generators import (
     MultiWriterIdGenerator,
     StreamIdGenerator,
 )
-from synapse.types import JsonDict, RetentionPolicy, ThirdPartyInstanceID
+from synapse.types import JsonDict, RetentionPolicy, StrCollection, ThirdPartyInstanceID
 from synapse.util import json_encoder
-from synapse.util.caches.descriptors import cached
+from synapse.util.caches.descriptors import cached, cachedList
 from synapse.util.stringutils import MXC_REGEX
 
 if TYPE_CHECKING:
@@ -1274,6 +1274,25 @@ class RoomWorkerStore(CacheInvalidationWorkerStore):
 
         return entry is not None
 
+    @cachedList(cached_method_name="is_partial_State_room", list_name="room_ids")
+    async def is_partial_state_rooms(self, room_ids: StrCollection) -> Mapping[str, bool]:
+        """Checks if this room has partial state.
+
+        Returns true if this is a "partial-state" room, which means that the state
+        at events in the room, and `current_state_events`, may not yet be
+        complete.
+        """
+
+        entries = set(await self.db_pool.simple_select_many_batch(
+            table="partial_state_rooms",
+            column="room_id",
+            iterable=room_ids,
+            retcols=("room_id",),
+            desc="is_partial_state_room",
+        ))
+
+        return {room_id: room_id in entries for room_id in room_ids}
+
     async def get_join_event_id_and_device_lists_stream_id_for_partial_state(
         self, room_id: str
     ) -> Tuple[str, int]: