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: [PATCH] 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
---
 changelog.d/8815.misc                        |  1 +
 synapse/handlers/room_member.py              | 16 ++++++---
 synapse/storage/databases/main/roommember.py | 34 +++++++++++++++++++-
 3 files changed, 45 insertions(+), 6 deletions(-)
 create mode 100644 changelog.d/8815.misc

diff --git a/changelog.d/8815.misc b/changelog.d/8815.misc
new file mode 100644
index 000000000000..647edeb56851
--- /dev/null
+++ b/changelog.d/8815.misc
@@ -0,0 +1 @@
+Optimise the lookup for an invite from another homeserver when trying to reject it.
\ No newline at end of file
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index 70f896626753..13a793c05a5a 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -31,7 +31,6 @@
 from synapse.api.ratelimiting import Ratelimiter
 from synapse.events import EventBase
 from synapse.events.snapshot import EventContext
-from synapse.storage.roommember import RoomsForUser
 from synapse.types import JsonDict, Requester, RoomAlias, RoomID, StateMap, UserID
 from synapse.util.async_helpers import Linearizer
 from synapse.util.distributor import user_left_room
@@ -515,10 +514,16 @@ async def update_membership_locked(
         elif effective_membership_state == Membership.LEAVE:
             if not is_host_in_room:
                 # perhaps we've been invited
-                invite = await self.store.get_invite_for_local_user_in_room(
-                    user_id=target.to_string(), room_id=room_id
-                )  # type: Optional[RoomsForUser]
-                if not invite:
+                (
+                    current_membership_type,
+                    current_membership_event_id,
+                ) = await self.store.get_local_current_membership_for_user_in_room(
+                    target.to_string(), room_id
+                )
+                if (
+                    current_membership_type != Membership.INVITE
+                    or not current_membership_event_id
+                ):
                     logger.info(
                         "%s sent a leave request to %s, but that is not an active room "
                         "on this server, and there is no pending invite",
@@ -528,6 +533,7 @@ async def update_membership_locked(
 
                     raise SynapseError(404, "Not a known room")
 
+                invite = await self.store.get_event(current_membership_event_id)
                 logger.info(
                     "%s rejects invite to %s from %s", target, room_id, invite.sender
                 )
diff --git a/synapse/storage/databases/main/roommember.py b/synapse/storage/databases/main/roommember.py
index 01d9dbb36f44..dcdaf09682b6 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 @@ def _get_rooms_for_local_user_where_membership_is_txn(
 
         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