From d963c69ba56ea45276ec3d11d191a20e8a38881d Mon Sep 17 00:00:00 2001
From: Andrew Morgan <1342360+anoadragon453@users.noreply.github.com>
Date: Wed, 25 Nov 2020 20:06:13 +0000
Subject: Speed up remote invite rejection database call (#8815)

This is another PR that grew out of #6739.

The existing code for checking whether a user is currently invited to a room when they want to leave the room looks like the following:

https://github.com/matrix-org/synapse/blob/f737368a26bb9eea401fcc3a5bdd7e0b59e91f09/synapse/handlers/room_member.py#L518-L540

It calls `get_invite_for_local_user_in_room`, which will actually query *all* rooms the user has been invited to, before iterating over them and matching via the room ID. It will then return a tuple of a lot of information which we pull the event ID out of.

I need to do a similar check for knocking, but this code wasn't very efficient. I then tried to write a different implementation using `StateHandler.get_current_state` but this actually didn't work as we haven't *joined* the room yet - we've only been invited to it. That means that only certain tables in Synapse have our desired `invite` membership state. One of those tables is `local_current_membership`.

So I wrote a store method that just queries that table instead
---
 synapse/storage/databases/main/roommember.py | 34 +++++++++++++++++++++++++++-
 1 file changed, 33 insertions(+), 1 deletion(-)

(limited to 'synapse/storage')

diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index 01d9dbb36f..dcdaf09682 100644
--- a/synapse/storage/databases/main/roommember.py
+++ b/synapse/storage/databases/main/roommember.py
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import TYPE_CHECKING, Dict, FrozenSet, Iterable, List, Optional, Set
+from typing import TYPE_CHECKING, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple
 
 from synapse.api.constants import EventTypes, Membership
 from synapse.events import EventBase
@@ -350,6 +350,38 @@ class RoomMemberWorkerStore(EventsWorkerStore):
 
         return results
 
+    async def get_local_current_membership_for_user_in_room(
+        self, user_id: str, room_id: str
+    ) -> Tuple[Optional[str], Optional[str]]:
+        """Retrieve the current local membership state and event ID for a user in a room.
+
+        Args:
+            user_id: The ID of the user.
+            room_id: The ID of the room.
+
+        Returns:
+            A tuple of (membership_type, event_id). Both will be None if a
+                room_id/user_id pair is not found.
+        """
+        # Paranoia check.
+        if not self.hs.is_mine_id(user_id):
+            raise Exception(
+                "Cannot call 'get_local_current_membership_for_user_in_room' on "
+                "non-local user %s" % (user_id,),
+            )
+
+        results_dict = await self.db_pool.simple_select_one(
+            "local_current_membership",
+            {"room_id": room_id, "user_id": user_id},
+            ("membership", "event_id"),
+            allow_none=True,
+            desc="get_local_current_membership_for_user_in_room",
+        )
+        if not results_dict:
+            return None, None
+
+        return results_dict.get("membership"), results_dict.get("event_id")
+
     @cached(max_entries=500000, iterable=True)
     async def get_rooms_for_user_with_stream_ordering(
         self, user_id: str
-- 
cgit 1.4.1