From ce641b0f91737c8b2fad46f2299f9022e67dfd51 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 4 Apr 2023 14:11:34 +0100
Subject: [PATCH] Speed up membership queries for users with forgotten rooms
 (#15385)

---
 changelog.d/15385.misc                        |  1 +
 synapse/storage/databases/main/roommember.py  | 12 +++++++++++-
 .../delta/74/03_room_membership_index.sql     | 19 +++++++++++++++++++
 3 files changed, 31 insertions(+), 1 deletion(-)
 create mode 100644 changelog.d/15385.misc
 create mode 100644 synapse/storage/schema/main/delta/74/03_room_membership_index.sql

diff --git a/changelog.d/15385.misc b/changelog.d/15385.misc
new file mode 100644
index 000000000000..76350c398d6f
--- /dev/null
+++ b/changelog.d/15385.misc
@@ -0,0 +1 @@
+Speed up membership queries for users with forgotten rooms.
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index 694a5b802c7c..daad58291a8b 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -419,7 +419,11 @@ async def get_rooms_for_local_user_where_membership_is(
         )
 
         # Now we filter out forgotten and excluded rooms
-        rooms_to_exclude = await self.get_forgotten_rooms_for_user(user_id)
+        rooms_to_exclude: AbstractSet[str] = set()
+
+        # Users can't forget joined/invited rooms, so we skip the check for such look ups.
+        if any(m not in (Membership.JOIN, Membership.INVITE) for m in membership_list):
+            rooms_to_exclude = await self.get_forgotten_rooms_for_user(user_id)
 
         if excluded_rooms is not None:
             # Take a copy to avoid mutating the in-cache set
@@ -1391,6 +1395,12 @@ def __init__(
             columns=["user_id", "room_id"],
             where_clause="forgotten = 1",
         )
+        self.db_pool.updates.register_background_index_update(
+            "room_membership_user_room_index",
+            index_name="room_membership_user_room_idx",
+            table="room_memberships",
+            columns=["user_id", "room_id"],
+        )
 
     async def _background_add_membership_profile(
         self, progress: JsonDict, batch_size: int
diff --git a/synapse/storage/schema/main/delta/74/03_room_membership_index.sql b/synapse/storage/schema/main/delta/74/03_room_membership_index.sql
new file mode 100644
index 000000000000..81a7d9ff9cad
--- /dev/null
+++ b/synapse/storage/schema/main/delta/74/03_room_membership_index.sql
@@ -0,0 +1,19 @@
+/* Copyright 2023 The Matrix.org Foundation C.I.C
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+-- Add an index to `room_membership(user_id, room_id)` to make querying for
+-- forgotten rooms faster.
+INSERT INTO background_updates (ordering, update_name, progress_json) VALUES
+    (7403, 'room_membership_user_room_index', '{}');