Speed up remote invite rejection database call

changelog.d/8815.misc

@@ -0,0 +1 @@
+Optimise the lookup for an invite from another homeserver when trying to reject it.

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,17 @@ 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 (
+                    not current_membership_type
+                    or 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 +534,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
                 )

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,51 @@ def _get_rooms_for_local_user_where_membership_is_txn(
 
         return results
 
+    @cached()
+    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.
+
+        """
+        return await self.db_pool.runInteraction(
+            "get_local_current_membership_for_user_in_room",
+            self._get_local_current_membership_for_user_in_room_txn,
+            user_id,
+            room_id,
+        )
+
+    def _get_local_current_membership_for_user_in_room_txn(
+        self, txn, user_id: str, room_id: str
+    ) -> Tuple[Optional[str], Optional[str]]:
+        # 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,),
+            )
+
+        sql = """
+            SELECT membership, event_id
+            FROM local_current_membership
+            WHERE
+                room_id = ? AND user_id = ?
+        """
+
+        txn.execute(sql, (room_id, user_id))
+        row = txn.fetchone()
+        if row:
+            return row[0], row[1]
+        return None, None
+
     @cached(max_entries=500000, iterable=True)
     async def get_rooms_for_user_with_stream_ordering(
         self, user_id: str