From 384ccf2c86f9190e2867212514ac8202981f57b5 Mon Sep 17 00:00:00 2001
From: Mathieu Velten <mathieuv@matrix.org>
Date: Fri, 2 Dec 2022 18:00:01 +0100
Subject: [PATCH 01/19] Faster joins: don't stale when an user joins during a
 fast join

Signed-off-by: Mathieu Velten <mathieuv@matrix.org>
---
 changelog.d/14606.misc               |  1 +
 synapse/handlers/event_auth.py       |  8 ++---
 synapse/handlers/federation.py       |  6 ++--
 synapse/handlers/federation_event.py |  5 +--
 synapse/handlers/room_member.py      | 52 +++++++++++++++++++++-------
 5 files changed, 50 insertions(+), 22 deletions(-)
 create mode 100644 changelog.d/14606.misc

diff --git a/changelog.d/14606.misc b/changelog.d/14606.misc
new file mode 100644
index 000000000000..52f9a901e915
--- /dev/null
+++ b/changelog.d/14606.misc
@@ -0,0 +1 @@
+Faster joins: don't stale when another user joins during a fast join resync.
diff --git a/synapse/handlers/event_auth.py b/synapse/handlers/event_auth.py
index f91dbbecb79c..a258af6f4473 100644
--- a/synapse/handlers/event_auth.py
+++ b/synapse/handlers/event_auth.py
@@ -202,7 +202,7 @@ async def check_restricted_join_rules(
         state_ids: StateMap[str],
         room_version: RoomVersion,
         user_id: str,
-        prev_member_event: Optional[EventBase],
+        prev_membership: Optional[str],
     ) -> None:
         """
         Check whether a user can join a room without an invite due to restricted join rules.
@@ -214,15 +214,13 @@ async def check_restricted_join_rules(
             state_ids: The state of the room as it currently is.
             room_version: The room version of the room being joined.
             user_id: The user joining the room.
-            prev_member_event: The current membership event for this user.
+            prev_membership: The current membership state for this user.
 
         Raises:
             AuthError if the user cannot join the room.
         """
         # If the member is invited or currently joined, then nothing to do.
-        if prev_member_event and (
-            prev_member_event.membership in (Membership.JOIN, Membership.INVITE)
-        ):
+        if prev_membership in (Membership.JOIN, Membership.INVITE):
             return
 
         # This is not a room with a restricted join rule, so we don't need to do the
diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index eca75f1108d1..0d2e1700b949 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -599,6 +599,8 @@ async def do_invite_join(
             except ValueError:
                 pass
 
+            already_partially_joined = await self.store.is_partial_state_room(room_id)
+
             ret = await self.federation_client.send_join(
                 host_list, event, room_version_obj
             )
@@ -629,7 +631,7 @@ async def do_invite_join(
                 state_events=state,
             )
 
-            if ret.partial_state:
+            if ret.partial_state and not already_partially_joined:
                 # Mark the room as having partial state.
                 # The background process is responsible for unmarking this flag,
                 # even if the join fails.
@@ -676,7 +678,7 @@ async def do_invite_join(
                 # state for the room.
                 # If the join failed, the background process is responsible for
                 # cleaning up — including unmarking the room as a partial state room.
-                if ret.partial_state:
+                if ret.partial_state and not already_partially_joined:
                     # Kick off the process of asynchronously fetching the state for this
                     # room.
                     run_as_background_process(
diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index 31df7f55cc97..bb39cae01f27 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -438,16 +438,17 @@ async def check_join_restrictions(
         # Check if the user is already in the room or invited to the room.
         user_id = event.state_key
         prev_member_event_id = prev_state_ids.get((EventTypes.Member, user_id), None)
-        prev_member_event = None
+        prev_membership = None
         if prev_member_event_id:
             prev_member_event = await self._store.get_event(prev_member_event_id)
+            prev_membership = prev_member_event.membership
 
         # Check if the member should be allowed access via membership in a space.
         await self._event_auth_handler.check_restricted_join_rules(
             prev_state_ids,
             event.room_version,
             user_id,
-            prev_member_event,
+            prev_membership,
         )
 
     @trace
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index d236cc09b526..e9061049c5c1 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -823,8 +823,14 @@ async def update_membership_locked(
 
         latest_event_ids = await self.store.get_prev_events_for_room(room_id)
 
+        # There are 2 reasons that mandate to have the full state:
+        # - to calculate `is_host_in_room`, however we now approximate that we consider the server
+        # in the room if we are still in the middle of a fast join.
+        # - to calculate if another user trying to join is allowed to. In this case we now
+        # make another `make/send_join` to a resident server to validate that. We could just
+        # accept it but it would mean that we would potentially leak the history to a banned user.
         state_before_join = await self.state_handler.compute_state_after_events(
-            room_id, latest_event_ids
+            room_id, latest_event_ids, await_full_state=False
         )
 
         # TODO: Refactor into dictionary of explicitly allowed transitions
@@ -881,7 +887,10 @@ async def update_membership_locked(
             if action == "kick":
                 raise AuthError(403, "The target user is not in the room")
 
-        is_host_in_room = await self._is_host_in_room(state_before_join)
+        is_partial_state_room = await self.store.is_partial_state_room(room_id)
+        is_host_in_room = await self._is_host_in_room(
+            is_partial_state_room, state_before_join
+        )
 
         if effective_membership_state == Membership.JOIN:
             if requester.is_guest:
@@ -927,6 +936,7 @@ async def update_membership_locked(
                 room_id,
                 remote_room_hosts,
                 content,
+                is_partial_state_room,
                 is_host_in_room,
                 state_before_join,
             )
@@ -1073,6 +1083,7 @@ async def _should_perform_remote_join(
         room_id: str,
         remote_room_hosts: List[str],
         content: JsonDict,
+        is_partial_state_room: bool,
         is_host_in_room: bool,
         state_before_join: StateMap[str],
     ) -> Tuple[bool, List[str]]:
@@ -1093,6 +1104,8 @@ async def _should_perform_remote_join(
             remote_room_hosts: A list of remote room hosts.
             content: The content to use as the event body of the join. This may
                 be modified.
+            is_partial_state_room: True if the server currently doesn't hold the fully
+                state of the room.
             is_host_in_room: True if the host is in the room.
             state_before_join: The state before the join event (i.e. the resolution of
                 the states after its parent events).
@@ -1109,6 +1122,20 @@ async def _should_perform_remote_join(
         if not is_host_in_room:
             return True, remote_room_hosts
 
+        prev_member_event_id = state_before_join.get((EventTypes.Member, user_id), None)
+        previous_membership = None
+        if prev_member_event_id:
+            prev_member_event = await self.store.get_event(prev_member_event_id)
+            previous_membership = prev_member_event.membership
+
+        # If we are not fully join yet, and the target is not already in the room,
+        # let's do a remote join so another server with the full state can validate
+        # that the user has not been ban for example.
+        # We could just accept the join and wait for state-res to resolve that later on
+        # but we would then leak room history to this person until then, which is pretty bad.
+        if is_partial_state_room and previous_membership != Membership.JOIN:
+            return True, remote_room_hosts
+
         # If the host is in the room, but not one of the authorised hosts
         # for restricted join rules, a remote join must be used.
         room_version = await self.store.get_room_version(room_id)
@@ -1122,15 +1149,8 @@ async def _should_perform_remote_join(
 
         # If the user is invited to the room or already joined, the join
         # event can always be issued locally.
-        prev_member_event_id = state_before_join.get((EventTypes.Member, user_id), None)
-        prev_member_event = None
-        if prev_member_event_id:
-            prev_member_event = await self.store.get_event(prev_member_event_id)
-            if prev_member_event.membership in (
-                Membership.JOIN,
-                Membership.INVITE,
-            ):
-                return False, []
+        if previous_membership in (Membership.JOIN, Membership.INVITE):
+            return False, []
 
         # If the local host has a user who can issue invites, then a local
         # join can be done.
@@ -1154,7 +1174,7 @@ async def _should_perform_remote_join(
 
         # Ensure the member should be allowed access via membership in a room.
         await self.event_auth_handler.check_restricted_join_rules(
-            state_before_join, room_version, user_id, prev_member_event
+            state_before_join, room_version, user_id, previous_membership
         )
 
         # If this is going to be a local join, additional information must
@@ -1634,7 +1654,13 @@ async def _make_and_store_3pid_invite(
         )
         return event, stream_id
 
-    async def _is_host_in_room(self, current_state_ids: StateMap[str]) -> bool:
+    async def _is_host_in_room(
+        self, is_partial_state_room: bool, current_state_ids: StateMap[str]
+    ) -> bool:
+        # When we only have a partial state, let's assume we are still in the room
+        if is_partial_state_room:
+            return True
+
         # Have we just created the room, and is this about to be the very
         # first member event?
         create_event_id = current_state_ids.get(("m.room.create", ""))

From 7c1e8cf4e3dcc4727a83878abf4834232d1b1a40 Mon Sep 17 00:00:00 2001
From: Mathieu Velten <matmaul@gmail.com>
Date: Thu, 29 Dec 2022 16:42:49 +0100
Subject: [PATCH 02/19] Only delete fw extrems on join if room state is not
 partial

---
 synapse/storage/databases/main/event_federation.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/synapse/storage/databases/main/event_federation.py b/synapse/storage/databases/main/event_federation.py
index bbee02ab18f0..185fecbeacb5 100644
--- a/synapse/storage/databases/main/event_federation.py
+++ b/synapse/storage/databases/main/event_federation.py
@@ -2016,7 +2016,12 @@ async def clean_room_for_join(self, room_id: str) -> None:
         )
 
     def _clean_room_for_join_txn(self, txn: LoggingTransaction, room_id: str) -> None:
-        query = "DELETE FROM event_forward_extremities WHERE room_id = ?"
+        query = """
+                DELETE FROM event_forward_extremities
+                    WHERE room_id = ? AND room_id NOT IN (
+                        SELECT room_id FROM partial_state_rooms
+                    )
+                """
 
         txn.execute(query, (room_id,))
         txn.call_after(self.get_latest_event_ids_in_room.invalidate, (room_id,))

From 0c3bc68a002ca26c79289f9086554fddd5de32fc Mon Sep 17 00:00:00 2001
From: Mathieu Velten <matmaul@gmail.com>
Date: Thu, 29 Dec 2022 16:43:28 +0100
Subject: [PATCH 03/19] Only call _clean_room_for_join if send_join was
 successful

---
 synapse/handlers/federation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index 0d2e1700b949..a53a4d92f493 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -587,8 +587,6 @@ async def do_invite_join(
 
         self._federation_event_handler.room_queues[room_id] = []
 
-        await self._clean_room_for_join(room_id)
-
         try:
             # Try the host we successfully got a response to /make_join/
             # request first.
@@ -616,6 +614,8 @@ async def do_invite_join(
 
             logger.debug("do_invite_join event: %s", event)
 
+            await self._clean_room_for_join(room_id)
+
             # if this is the first time we've joined this room, it's time to add
             # a row to `rooms` with the correct room version. If there's already a
             # row there, we should override it, since it may have been populated

From 4f6c3422f9b01963b9b06f90db03fb16c9357364 Mon Sep 17 00:00:00 2001
From: Mathieu Velten <matmaul@gmail.com>
Date: Thu, 29 Dec 2022 17:11:57 +0100
Subject: [PATCH 04/19] move already_partially_joined closer to its use to
 avoid races

---
 synapse/handlers/federation.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index a53a4d92f493..fdf9f257b68d 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -597,8 +597,6 @@ async def do_invite_join(
             except ValueError:
                 pass
 
-            already_partially_joined = await self.store.is_partial_state_room(room_id)
-
             ret = await self.federation_client.send_join(
                 host_list, event, room_version_obj
             )
@@ -631,6 +629,7 @@ async def do_invite_join(
                 state_events=state,
             )
 
+            already_partially_joined = await self.store.is_partial_state_room(room_id)
             if ret.partial_state and not already_partially_joined:
                 # Mark the room as having partial state.
                 # The background process is responsible for unmarking this flag,

From 6aa72e12b30d77cfb6fbfb130ce37406a676ae7f Mon Sep 17 00:00:00 2001
From: Mathieu Velten <mathieuv@matrix.org>
Date: Mon, 2 Jan 2023 20:33:18 +0100
Subject: [PATCH 05/19] Fallback to local join if a room is unpartial stated
 during remote join

---
 synapse/handlers/federation.py           |  9 ++++---
 synapse/handlers/room_member.py          | 31 ++++++++++++++++++------
 synapse/handlers/room_member_worker.py   | 26 ++++++++++++--------
 synapse/replication/http/membership.py   | 11 ++++++---
 synapse/storage/databases/main/events.py |  6 ++++-
 5 files changed, 59 insertions(+), 24 deletions(-)

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index fdf9f257b68d..cd4b1c2f6b68 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -562,6 +562,10 @@ async def do_invite_join(
             joinee: The User ID of the joining user.
 
             content: The event content to use for the join event.
+
+        Raises:
+            PartialStateConflictError if homeserver was already in the room and it is
+                no longer partial stated.
         """
         # TODO: We should be able to call this on workers, but the upgrading of
         # room stuff after join currently doesn't work on workers.
@@ -655,15 +659,14 @@ async def do_invite_join(
                 )
             except PartialStateConflictError as e:
                 # The homeserver was already in the room and it is no longer partial
-                # stated. We ought to be doing a local join instead. Turn the error into
-                # a 429, as a hint to the client to try again.
+                # stated. We ought to be doing a local join instead.
                 # TODO(faster_joins): `_should_perform_remote_join` suggests that we may
                 #   do a remote join for restricted rooms even if we have full state.
                 logger.error(
                     "Room %s was un-partial stated while processing remote join.",
                     room_id,
                 )
-                raise LimitExceededError(msg=e.msg, errcode=e.errcode, retry_after_ms=0)
+                raise e
             else:
                 # Record the join event id for future use (when we finish the full
                 # join). We have to do this after persisting the event to keep foreign
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index e9061049c5c1..56807ec8a9f1 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -56,6 +56,13 @@
 logger = logging.getLogger(__name__)
 
 
+class NoKnownServersError(SynapseError):
+    """No server already resident to the room was provided to the join/knock operation."""
+
+    def __init__(self, msg: str = "No known servers"):
+        super().__init__(404, msg)
+
+
 class RoomMemberHandler(metaclass=abc.ABCMeta):
     # TODO(paul): This handler currently contains a messy conflation of
     #   low-level API that works on UserID objects and so on, and REST-level
@@ -185,6 +192,12 @@ async def _remote_join(
             room_id: Room that we are trying to join
             user: User who is trying to join
             content: A dict that should be used as the content of the join event.
+
+        Raises:
+            NoKnownServersError: if remote_room_hosts does not contain a server joined to
+                the room.
+            PartialStateConflictError: if the room was un-partial stated in the meantime,
+                a local room join should be done instead.
         """
         raise NotImplementedError()
 
@@ -972,11 +985,16 @@ async def update_membership_locked(
                 if requester.is_guest:
                     content["kind"] = "guest"
 
-                remote_join_response = await self._remote_join(
-                    requester, remote_room_hosts, room_id, target, content
-                )
+                try:
+                    remote_join_response = await self._remote_join(
+                        requester, remote_room_hosts, room_id, target, content
+                    )
 
-                return remote_join_response
+                    return remote_join_response
+                except PartialStateConflictError:
+                    # Room has been un-partial stated in the meantime, let's continue
+                    # the code flow to trigger a local join through _local_membership_update.
+                    pass
 
         elif effective_membership_state == Membership.LEAVE:
             if not is_host_in_room:
@@ -1741,8 +1759,7 @@ async def _remote_join(
         ]
 
         if len(remote_room_hosts) == 0:
-            raise SynapseError(
-                404,
+            raise NoKnownServersError(
                 "Can't join remote room because no servers "
                 "that are in the room have been provided.",
             )
@@ -1973,7 +1990,7 @@ async def remote_knock(
         ]
 
         if len(remote_room_hosts) == 0:
-            raise SynapseError(404, "No known servers")
+            raise NoKnownServersError()
 
         return await self.federation_handler.do_knock(
             remote_room_hosts, room_id, user.to_string(), content=content
diff --git a/synapse/handlers/room_member_worker.py b/synapse/handlers/room_member_worker.py
index 221552a2a64b..9a1f1edb7122 100644
--- a/synapse/handlers/room_member_worker.py
+++ b/synapse/handlers/room_member_worker.py
@@ -13,10 +13,11 @@
 # limitations under the License.
 
 import logging
+from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from synapse.api.errors import SynapseError
-from synapse.handlers.room_member import RoomMemberHandler
+from synapse.handlers.room_member import NoKnownServersError, RoomMemberHandler
 from synapse.replication.http.membership import (
     ReplicationRemoteJoinRestServlet as ReplRemoteJoin,
     ReplicationRemoteKnockRestServlet as ReplRemoteKnock,
@@ -24,6 +25,7 @@
     ReplicationRemoteRescindKnockRestServlet as ReplRescindKnock,
     ReplicationUserJoinedLeftRoomRestServlet as ReplJoinedLeft,
 )
+from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.types import JsonDict, Requester, UserID
 
 if TYPE_CHECKING:
@@ -52,15 +54,19 @@ async def _remote_join(
     ) -> Tuple[str, int]:
         """Implements RoomMemberHandler._remote_join"""
         if len(remote_room_hosts) == 0:
-            raise SynapseError(404, "No known servers")
-
-        ret = await self._remote_join_client(
-            requester=requester,
-            remote_room_hosts=remote_room_hosts,
-            room_id=room_id,
-            user_id=user.to_string(),
-            content=content,
-        )
+            raise NoKnownServersError()
+
+        try:
+            ret = await self._remote_join_client(
+                requester=requester,
+                remote_room_hosts=remote_room_hosts,
+                room_id=room_id,
+                user_id=user.to_string(),
+                content=content,
+            )
+        except SynapseError as e:
+            if e.code == HTTPStatus.CONFLICT:
+                raise PartialStateConflictError()
 
         return ret["event_id"], ret["stream_id"]
 
diff --git a/synapse/replication/http/membership.py b/synapse/replication/http/membership.py
index 663bff573848..de451e4fe5d5 100644
--- a/synapse/replication/http/membership.py
+++ b/synapse/replication/http/membership.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
+from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from twisted.web.server import Request
@@ -20,6 +21,7 @@
 from synapse.http.servlet import parse_json_object_from_request
 from synapse.http.site import SynapseRequest
 from synapse.replication.http._base import ReplicationEndpoint
+from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.types import JsonDict, Requester, UserID
 from synapse.util.distributor import user_left_room
 
@@ -91,9 +93,12 @@ async def _handle_request(  # type: ignore[override]
 
         logger.info("remote_join: %s into room: %s", user_id, room_id)
 
-        event_id, stream_id = await self.federation_handler.do_invite_join(
-            remote_room_hosts, room_id, user_id, event_content
-        )
+        try:
+            event_id, stream_id = await self.federation_handler.do_invite_join(
+                remote_room_hosts, room_id, user_id, event_content
+            )
+        except PartialStateConflictError:
+            return HTTPStatus.CONFLICT, {"error": PartialStateConflictError.message()}
 
         return 200, {"event_id": event_id, "stream_id": stream_id}
 
diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py
index 0f097a2927c1..9e1a2b244251 100644
--- a/synapse/storage/databases/main/events.py
+++ b/synapse/storage/databases/main/events.py
@@ -82,10 +82,14 @@ class PartialStateConflictError(SynapseError):
     It should not be exposed to clients.
     """
 
+    @staticmethod
+    def message() -> str:
+        return "Cannot persist partial state event in un-partial stated room"
+
     def __init__(self) -> None:
         super().__init__(
             HTTPStatus.CONFLICT,
-            msg="Cannot persist partial state event in un-partial stated room",
+            msg=PartialStateConflictError.message(),
             errcode=Codes.UNKNOWN,
         )
 

From d98994c1dd85f7ecd53f292988f93855824d8b57 Mon Sep 17 00:00:00 2001
From: Mathieu Velten <mathieuv@matrix.org>
Date: Mon, 2 Jan 2023 22:06:10 +0100
Subject: [PATCH 06/19] Move PartialStateConflictError to api.errors because of
 circular dep

---
 synapse/api/errors.py                    | 22 ++++++++++++++++++++++
 synapse/federation/federation_server.py  |  2 +-
 synapse/handlers/federation.py           |  2 +-
 synapse/handlers/federation_event.py     |  2 +-
 synapse/handlers/message.py              |  2 +-
 synapse/handlers/room.py                 |  2 +-
 synapse/handlers/room_member.py          |  9 +++++++--
 synapse/handlers/room_member_worker.py   |  3 +--
 synapse/replication/http/membership.py   |  2 +-
 synapse/storage/databases/main/events.py | 24 +-----------------------
 10 files changed, 37 insertions(+), 33 deletions(-)

diff --git a/synapse/api/errors.py b/synapse/api/errors.py
index c2c177fd71d0..9235ce653638 100644
--- a/synapse/api/errors.py
+++ b/synapse/api/errors.py
@@ -751,3 +751,25 @@ class ModuleFailedException(Exception):
     Raised when a module API callback fails, for example because it raised an
     exception.
     """
+
+
+class PartialStateConflictError(SynapseError):
+    """An internal error raised when attempting to persist an event with partial state
+    after the room containing the event has been un-partial stated.
+
+    This error should be handled by recomputing the event context and trying again.
+
+    This error has an HTTP status code so that it can be transported over replication.
+    It should not be exposed to clients.
+    """
+
+    @staticmethod
+    def message() -> str:
+        return "Cannot persist partial state event in un-partial stated room"
+
+    def __init__(self) -> None:
+        super().__init__(
+            HTTPStatus.CONFLICT,
+            msg=PartialStateConflictError.message(),
+            errcode=Codes.UNKNOWN,
+        )
diff --git a/synapse/federation/federation_server.py b/synapse/federation/federation_server.py
index bb20af6e91ed..53c73580f340 100644
--- a/synapse/federation/federation_server.py
+++ b/synapse/federation/federation_server.py
@@ -41,6 +41,7 @@
     FederationError,
     IncompatibleRoomVersionError,
     NotFoundError,
+    PartialStateConflictError,
     SynapseError,
     UnsupportedRoomVersionError,
 )
@@ -72,7 +73,6 @@
     ReplicationFederationSendEduRestServlet,
     ReplicationGetQueryRestServlet,
 )
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.storage.databases.main.lock import Lock
 from synapse.storage.databases.main.roommember import extract_heroes_from_room_summary
 from synapse.storage.roommember import MemberSummary
diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index cd4b1c2f6b68..9558bee9ccbb 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -49,6 +49,7 @@
     HttpResponseException,
     LimitExceededError,
     NotFoundError,
+    PartialStateConflictError,
     RequestSendFailed,
     SynapseError,
 )
@@ -68,7 +69,6 @@
     ReplicationCleanRoomRestServlet,
     ReplicationStoreRoomOnOutlierMembershipRestServlet,
 )
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
 from synapse.types import JsonDict, get_domain_from_id
 from synapse.types.state import StateFilter
diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index bb39cae01f27..8d7c4a24d073 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -47,6 +47,7 @@
     FederationError,
     FederationPullAttemptBackoffError,
     HttpResponseException,
+    PartialStateConflictError,
     RequestSendFailed,
     SynapseError,
 )
@@ -74,7 +75,6 @@
     ReplicationFederationSendEventsRestServlet,
 )
 from synapse.state import StateResolutionStore
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
 from synapse.types import (
     PersistedEventPosition,
diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 88fc51a4c97e..475f8048be37 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -38,6 +38,7 @@
     Codes,
     ConsentNotGivenError,
     NotFoundError,
+    PartialStateConflictError,
     ShadowBanError,
     SynapseError,
     UnstableSpecAuthError,
@@ -57,7 +58,6 @@
 from synapse.metrics.background_process_metrics import run_as_background_process
 from synapse.replication.http.send_event import ReplicationSendEventRestServlet
 from synapse.replication.http.send_events import ReplicationSendEventsRestServlet
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.storage.databases.main.events_worker import EventRedactBehaviour
 from synapse.types import (
     MutableStateMap,
diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index 572c7b4db344..08cc7862f6c7 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -51,6 +51,7 @@
     Codes,
     LimitExceededError,
     NotFoundError,
+    PartialStateConflictError,
     StoreError,
     SynapseError,
 )
@@ -62,7 +63,6 @@
 from synapse.handlers.relations import BundledAggregations
 from synapse.module_api import NOT_SPAM
 from synapse.rest.admin._base import assert_user_is_admin
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.streams import EventSource
 from synapse.types import (
     JsonDict,
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index 56807ec8a9f1..869ab252df8b 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -26,7 +26,13 @@
     GuestAccess,
     Membership,
 )
-from synapse.api.errors import AuthError, Codes, ShadowBanError, SynapseError
+from synapse.api.errors import (
+    AuthError,
+    Codes,
+    PartialStateConflictError,
+    ShadowBanError,
+    SynapseError,
+)
 from synapse.api.ratelimiting import Ratelimiter
 from synapse.event_auth import get_named_level, get_power_level_event
 from synapse.events import EventBase
@@ -34,7 +40,6 @@
 from synapse.handlers.profile import MAX_AVATAR_URL_LEN, MAX_DISPLAYNAME_LEN
 from synapse.logging import opentracing
 from synapse.module_api import NOT_SPAM
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.types import (
     JsonDict,
     Requester,
diff --git a/synapse/handlers/room_member_worker.py b/synapse/handlers/room_member_worker.py
index 9a1f1edb7122..7c931599d06e 100644
--- a/synapse/handlers/room_member_worker.py
+++ b/synapse/handlers/room_member_worker.py
@@ -16,7 +16,7 @@
 from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
-from synapse.api.errors import SynapseError
+from synapse.api.errors import PartialStateConflictError, SynapseError
 from synapse.handlers.room_member import NoKnownServersError, RoomMemberHandler
 from synapse.replication.http.membership import (
     ReplicationRemoteJoinRestServlet as ReplRemoteJoin,
@@ -25,7 +25,6 @@
     ReplicationRemoteRescindKnockRestServlet as ReplRescindKnock,
     ReplicationUserJoinedLeftRoomRestServlet as ReplJoinedLeft,
 )
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.types import JsonDict, Requester, UserID
 
 if TYPE_CHECKING:
diff --git a/synapse/replication/http/membership.py b/synapse/replication/http/membership.py
index de451e4fe5d5..d10a24923835 100644
--- a/synapse/replication/http/membership.py
+++ b/synapse/replication/http/membership.py
@@ -17,11 +17,11 @@
 
 from twisted.web.server import Request
 
+from synapse.api.errors import PartialStateConflictError
 from synapse.http.server import HttpServer
 from synapse.http.servlet import parse_json_object_from_request
 from synapse.http.site import SynapseRequest
 from synapse.replication.http._base import ReplicationEndpoint
-from synapse.storage.databases.main.events import PartialStateConflictError
 from synapse.types import JsonDict, Requester, UserID
 from synapse.util.distributor import user_left_room
 
diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py
index 9e1a2b244251..10746f886701 100644
--- a/synapse/storage/databases/main/events.py
+++ b/synapse/storage/databases/main/events.py
@@ -36,7 +36,7 @@
 
 import synapse.metrics
 from synapse.api.constants import EventContentFields, EventTypes, RelationTypes
-from synapse.api.errors import Codes, SynapseError
+from synapse.api.errors import Codes, PartialStateConflictError, SynapseError
 from synapse.api.room_versions import RoomVersions
 from synapse.events import EventBase, relation_from_event
 from synapse.events.snapshot import EventContext
@@ -72,28 +72,6 @@
 )
 
 
-class PartialStateConflictError(SynapseError):
-    """An internal error raised when attempting to persist an event with partial state
-    after the room containing the event has been un-partial stated.
-
-    This error should be handled by recomputing the event context and trying again.
-
-    This error has an HTTP status code so that it can be transported over replication.
-    It should not be exposed to clients.
-    """
-
-    @staticmethod
-    def message() -> str:
-        return "Cannot persist partial state event in un-partial stated room"
-
-    def __init__(self) -> None:
-        super().__init__(
-            HTTPStatus.CONFLICT,
-            msg=PartialStateConflictError.message(),
-            errcode=Codes.UNKNOWN,
-        )
-
-
 @attr.s(slots=True, auto_attribs=True)
 class DeltaState:
     """Deltas to use to update the `current_state_events` table.

From a24d07adb9db1754d4e2102b4af86df759377190 Mon Sep 17 00:00:00 2001
From: Mathieu Velten <mathieuv@matrix.org>
Date: Mon, 9 Jan 2023 15:35:52 +0100
Subject: [PATCH 07/19] Small mistake

---
 synapse/handlers/room_member_worker.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/synapse/handlers/room_member_worker.py b/synapse/handlers/room_member_worker.py
index 7c931599d06e..bd40f2c50548 100644
--- a/synapse/handlers/room_member_worker.py
+++ b/synapse/handlers/room_member_worker.py
@@ -66,6 +66,8 @@ async def _remote_join(
         except SynapseError as e:
             if e.code == HTTPStatus.CONFLICT:
                 raise PartialStateConflictError()
+            else:
+                raise e
 
         return ret["event_id"], ret["stream_id"]
 

From 9a934c1d1768a0a26ea28a4cccf284e68eb0eb85 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Fri, 27 Jan 2023 15:21:08 +0000
Subject: [PATCH 08/19] fix up for #14882

* the filter for partial state rooms in `_clean_room_for_join_txn` has
  been replaced with a host membership check in `_do_invite_join`.
* `_do_invite_join` can no longer raise a `PartialStateConflictError`.
---
 synapse/handlers/room_member.py               | 23 ++++---------------
 synapse/handlers/room_member_worker.py        | 22 +++++++-----------
 .../databases/main/event_federation.py        |  7 +-----
 3 files changed, 14 insertions(+), 38 deletions(-)

diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index 869ab252df8b..dd4c383cf03c 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -26,13 +26,7 @@
     GuestAccess,
     Membership,
 )
-from synapse.api.errors import (
-    AuthError,
-    Codes,
-    PartialStateConflictError,
-    ShadowBanError,
-    SynapseError,
-)
+from synapse.api.errors import AuthError, Codes, ShadowBanError, SynapseError
 from synapse.api.ratelimiting import Ratelimiter
 from synapse.event_auth import get_named_level, get_power_level_event
 from synapse.events import EventBase
@@ -201,8 +195,6 @@ async def _remote_join(
         Raises:
             NoKnownServersError: if remote_room_hosts does not contain a server joined to
                 the room.
-            PartialStateConflictError: if the room was un-partial stated in the meantime,
-                a local room join should be done instead.
         """
         raise NotImplementedError()
 
@@ -990,16 +982,11 @@ async def update_membership_locked(
                 if requester.is_guest:
                     content["kind"] = "guest"
 
-                try:
-                    remote_join_response = await self._remote_join(
-                        requester, remote_room_hosts, room_id, target, content
-                    )
+                remote_join_response = await self._remote_join(
+                    requester, remote_room_hosts, room_id, target, content
+                )
 
-                    return remote_join_response
-                except PartialStateConflictError:
-                    # Room has been un-partial stated in the meantime, let's continue
-                    # the code flow to trigger a local join through _local_membership_update.
-                    pass
+                return remote_join_response
 
         elif effective_membership_state == Membership.LEAVE:
             if not is_host_in_room:
diff --git a/synapse/handlers/room_member_worker.py b/synapse/handlers/room_member_worker.py
index bd40f2c50548..2f3d9c8ca050 100644
--- a/synapse/handlers/room_member_worker.py
+++ b/synapse/handlers/room_member_worker.py
@@ -16,7 +16,7 @@
 from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
-from synapse.api.errors import PartialStateConflictError, SynapseError
+from synapse.api.errors import SynapseError
 from synapse.handlers.room_member import NoKnownServersError, RoomMemberHandler
 from synapse.replication.http.membership import (
     ReplicationRemoteJoinRestServlet as ReplRemoteJoin,
@@ -55,19 +55,13 @@ async def _remote_join(
         if len(remote_room_hosts) == 0:
             raise NoKnownServersError()
 
-        try:
-            ret = await self._remote_join_client(
-                requester=requester,
-                remote_room_hosts=remote_room_hosts,
-                room_id=room_id,
-                user_id=user.to_string(),
-                content=content,
-            )
-        except SynapseError as e:
-            if e.code == HTTPStatus.CONFLICT:
-                raise PartialStateConflictError()
-            else:
-                raise e
+        ret = await self._remote_join_client(
+            requester=requester,
+            remote_room_hosts=remote_room_hosts,
+            room_id=room_id,
+            user_id=user.to_string(),
+            content=content,
+        )
 
         return ret["event_id"], ret["stream_id"]
 
diff --git a/synapse/storage/databases/main/event_federation.py b/synapse/storage/databases/main/event_federation.py
index 185fecbeacb5..bbee02ab18f0 100644
--- a/synapse/storage/databases/main/event_federation.py
+++ b/synapse/storage/databases/main/event_federation.py
@@ -2016,12 +2016,7 @@ async def clean_room_for_join(self, room_id: str) -> None:
         )
 
     def _clean_room_for_join_txn(self, txn: LoggingTransaction, room_id: str) -> None:
-        query = """
-                DELETE FROM event_forward_extremities
-                    WHERE room_id = ? AND room_id NOT IN (
-                        SELECT room_id FROM partial_state_rooms
-                    )
-                """
+        query = "DELETE FROM event_forward_extremities WHERE room_id = ?"
 
         txn.execute(query, (room_id,))
         txn.call_after(self.get_latest_event_ids_in_room.invalidate, (room_id,))

From 5826574c66f3e83b833b1c4fe73e9eb62bf924a2 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Fri, 27 Jan 2023 15:51:32 +0000
Subject: [PATCH 09/19] fixup imports

---
 synapse/handlers/room_member.py          | 8 +++++++-
 synapse/handlers/room_member_worker.py   | 2 --
 synapse/storage/databases/main/events.py | 3 +--
 3 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index dd4c383cf03c..ab45cb27e930 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -26,7 +26,13 @@
     GuestAccess,
     Membership,
 )
-from synapse.api.errors import AuthError, Codes, ShadowBanError, SynapseError
+from synapse.api.errors import (
+    AuthError,
+    Codes,
+    PartialStateConflictError,
+    ShadowBanError,
+    SynapseError,
+)
 from synapse.api.ratelimiting import Ratelimiter
 from synapse.event_auth import get_named_level, get_power_level_event
 from synapse.events import EventBase
diff --git a/synapse/handlers/room_member_worker.py b/synapse/handlers/room_member_worker.py
index 2f3d9c8ca050..ba261702d4bc 100644
--- a/synapse/handlers/room_member_worker.py
+++ b/synapse/handlers/room_member_worker.py
@@ -13,10 +13,8 @@
 # limitations under the License.
 
 import logging
-from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
-from synapse.api.errors import SynapseError
 from synapse.handlers.room_member import NoKnownServersError, RoomMemberHandler
 from synapse.replication.http.membership import (
     ReplicationRemoteJoinRestServlet as ReplRemoteJoin,
diff --git a/synapse/storage/databases/main/events.py b/synapse/storage/databases/main/events.py
index 10746f886701..ca01501d9f10 100644
--- a/synapse/storage/databases/main/events.py
+++ b/synapse/storage/databases/main/events.py
@@ -16,7 +16,6 @@
 import itertools
 import logging
 from collections import OrderedDict
-from http import HTTPStatus
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -36,7 +35,7 @@
 
 import synapse.metrics
 from synapse.api.constants import EventContentFields, EventTypes, RelationTypes
-from synapse.api.errors import Codes, PartialStateConflictError, SynapseError
+from synapse.api.errors import PartialStateConflictError
 from synapse.api.room_versions import RoomVersions
 from synapse.events import EventBase, relation_from_event
 from synapse.events.snapshot import EventContext

From 1df7060c8d0d3bf7f60a8e1f5b5ecb3cb3a4492d Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Fri, 27 Jan 2023 16:00:41 +0000
Subject: [PATCH 10/19] fixup exception docstring

---
 synapse/handlers/federation.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index b7787a8b4393..ba81bba3b756 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -569,10 +569,6 @@ async def do_invite_join(
             joinee: The User ID of the joining user.
 
             content: The event content to use for the join event.
-
-        Raises:
-            PartialStateConflictError if homeserver was already in the room and it is
-                no longer partial stated.
         """
         # TODO: We should be able to call this on workers, but the upgrading of
         # room stuff after join currently doesn't work on workers.

From bfba7bb386a1785250ddbff8058a48eb8462115b Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Fri, 27 Jan 2023 16:03:02 +0000
Subject: [PATCH 11/19] fixup: remove more defunct PartialStateConflictError
 handling

---
 synapse/replication/http/membership.py | 11 +++--------
 1 file changed, 3 insertions(+), 8 deletions(-)

diff --git a/synapse/replication/http/membership.py b/synapse/replication/http/membership.py
index 38480762c63d..9fa1060d48f6 100644
--- a/synapse/replication/http/membership.py
+++ b/synapse/replication/http/membership.py
@@ -12,12 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from twisted.web.server import Request
 
-from synapse.api.errors import PartialStateConflictError
 from synapse.http.server import HttpServer
 from synapse.http.site import SynapseRequest
 from synapse.replication.http._base import ReplicationEndpoint
@@ -90,12 +88,9 @@ async def _handle_request(  # type: ignore[override]
 
         logger.info("remote_join: %s into room: %s", user_id, room_id)
 
-        try:
-            event_id, stream_id = await self.federation_handler.do_invite_join(
-                remote_room_hosts, room_id, user_id, event_content
-            )
-        except PartialStateConflictError:
-            return HTTPStatus.CONFLICT, {"error": PartialStateConflictError.message()}
+        event_id, stream_id = await self.federation_handler.do_invite_join(
+            remote_room_hosts, room_id, user_id, event_content
+        )
 
         return 200, {"event_id": event_id, "stream_id": stream_id}
 

From 0ecd9ba786bf861c23e41136cdb14489f5cf4824 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Tue, 31 Jan 2023 00:33:27 +0000
Subject: [PATCH 12/19] Avoid clobbering existing local memberships on a second
 partial join

Signed-off-by: Sean Quah <seanq@matrix.org>
---
 synapse/handlers/federation_event.py | 52 ++++++++++++++++++++++++++--
 1 file changed, 49 insertions(+), 3 deletions(-)

diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index 2586252f192d..ccdd8bf6b6f5 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -525,11 +525,57 @@ async def process_remote_join(
             "Peristing join-via-remote %s (partial_state: %s)", event, partial_state
         )
         with nested_logging_context(suffix=event.event_id):
+            if partial_state:
+                # When handling a second partial join into a partial state room, the
+                # returned state will exclude the membership from the first join. To
+                # preserve prior memberships, we try to compute the state before the
+                # event ourselves if we know about any of the prev events.
+                #
+                # When we don't know about any of the prev events, it's fine to just use
+                # the returned state, since the new join will create a new forward
+                # extremity, and leave the forward extremity containing our prior
+                # memberships alone.
+                prev_event_ids = set(event.prev_event_ids())
+                seen_event_ids = await self._store.have_events_in_timeline(
+                    prev_event_ids
+                )
+                missing_event_ids = prev_event_ids - seen_event_ids
+
+                state_maps_to_resolve: List[StateMap[str]] = []
+
+                # Fetch the state after the prev events that we know about.
+                state_maps_to_resolve.extend(
+                    (
+                        await self._state_storage_controller.get_state_groups_ids(
+                            room_id, seen_event_ids, await_full_state=False
+                        )
+                    ).values()
+                )
+
+                # When there are prev events we do not have the state for, we state
+                # resolve with the state returned by the remote homeserver.
+                if missing_event_ids or len(state_maps_to_resolve) == 0:
+                    state_maps_to_resolve.append(
+                        {(e.type, e.state_key): e.event_id for e in state}
+                    )
+
+                state_ids_before_event = (
+                    await self._state_resolution_handler.resolve_events_with_store(
+                        event.room_id,
+                        room_version.identifier,
+                        state_maps_to_resolve,
+                        event_map=None,
+                        state_res_store=StateResolutionStore(self._store),
+                    )
+                )
+            else:
+                state_ids_before_event = {
+                    (e.type, e.state_key): e.event_id for e in state
+                }
+
             context = await self._state_handler.compute_event_context(
                 event,
-                state_ids_before_event={
-                    (e.type, e.state_key): e.event_id for e in state
-                },
+                state_ids_before_event=state_ids_before_event,
                 partial_state=partial_state,
             )
 

From a23a64f303c7a08645852bacac5d02e0fa77038c Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Tue, 31 Jan 2023 01:37:56 +0000
Subject: [PATCH 13/19] fix trial tests

---
 tests/handlers/test_federation.py | 40 +++++++++++++++----------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/tests/handlers/test_federation.py b/tests/handlers/test_federation.py
index c1558c40c370..968871f3a0d7 100644
--- a/tests/handlers/test_federation.py
+++ b/tests/handlers/test_federation.py
@@ -575,26 +575,6 @@ def test_failed_partial_join_is_clean(self) -> None:
         fed_client = fed_handler.federation_client
 
         room_id = "!room:example.com"
-        membership_event = make_event_from_dict(
-            {
-                "room_id": room_id,
-                "type": "m.room.member",
-                "sender": "@alice:test",
-                "state_key": "@alice:test",
-                "content": {"membership": "join"},
-            },
-            RoomVersions.V10,
-        )
-
-        mock_make_membership_event = Mock(
-            return_value=make_awaitable(
-                (
-                    "example.com",
-                    membership_event,
-                    RoomVersions.V10,
-                )
-            )
-        )
 
         EVENT_CREATE = make_event_from_dict(
             {
@@ -640,6 +620,26 @@ def test_failed_partial_join_is_clean(self) -> None:
             },
             room_version=RoomVersions.V10,
         )
+        membership_event = make_event_from_dict(
+            {
+                "room_id": room_id,
+                "type": "m.room.member",
+                "sender": "@alice:test",
+                "state_key": "@alice:test",
+                "content": {"membership": "join"},
+                "prev_events": [EVENT_INVITATION_MEMBERSHIP.event_id],
+            },
+            RoomVersions.V10,
+        )
+        mock_make_membership_event = Mock(
+            return_value=make_awaitable(
+                (
+                    "example.com",
+                    membership_event,
+                    RoomVersions.V10,
+                )
+            )
+        )
         mock_send_join = Mock(
             return_value=make_awaitable(
                 SendJoinResult(

From 95a529c56e42981cbfee5d9d2eaf570e211f6714 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Fri, 3 Feb 2023 23:07:34 +0000
Subject: [PATCH 14/19] fix typo in changelog

---
 changelog.d/14606.misc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/changelog.d/14606.misc b/changelog.d/14606.misc
index 52f9a901e915..e2debc96d8d4 100644
--- a/changelog.d/14606.misc
+++ b/changelog.d/14606.misc
@@ -1 +1 @@
-Faster joins: don't stale when another user joins during a fast join resync.
+Faster joins: don't stall when another user joins during a fast join resync.

From 438d592839b4622d013a043de2f07322b9d73794 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Sat, 4 Feb 2023 22:34:26 +0000
Subject: [PATCH 15/19] fixup docstring to explain what a None prev_membership
 means

---
 synapse/handlers/event_auth.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/synapse/handlers/event_auth.py b/synapse/handlers/event_auth.py
index 31468c177ef6..08709c65be7c 100644
--- a/synapse/handlers/event_auth.py
+++ b/synapse/handlers/event_auth.py
@@ -214,7 +214,8 @@ async def check_restricted_join_rules(
             state_ids: The state of the room as it currently is.
             room_version: The room version of the room being joined.
             user_id: The user joining the room.
-            prev_membership: The current membership state for this user.
+            prev_membership: The current membership state for this user. `None` if the
+                user has never joined the room (equivalent to "leave").
 
         Raises:
             AuthError if the user cannot join the room.

From cfb71f00381d5f5dfa293a2041a85aa93a9a26d0 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Sat, 4 Feb 2023 22:36:31 +0000
Subject: [PATCH 16/19] fixup: is_host_in_room was already accurate given
 partial state

---
 synapse/handlers/room_member.py | 22 ++++++----------------
 1 file changed, 6 insertions(+), 16 deletions(-)

diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index ab45cb27e930..cd9bab24ef42 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -839,12 +839,10 @@ async def update_membership_locked(
 
         latest_event_ids = await self.store.get_prev_events_for_room(room_id)
 
-        # There are 2 reasons that mandate to have the full state:
-        # - to calculate `is_host_in_room`, however we now approximate that we consider the server
-        # in the room if we are still in the middle of a fast join.
-        # - to calculate if another user trying to join is allowed to. In this case we now
-        # make another `make/send_join` to a resident server to validate that. We could just
-        # accept it but it would mean that we would potentially leak the history to a banned user.
+        # The full state is needed to calculate if another user trying to join is
+        # allowed to. In this case we now make another `make/send_join` to a resident
+        # server to validate that. We could just accept it but it would mean that we
+        # would potentially leak the history to a banned user.
         state_before_join = await self.state_handler.compute_state_after_events(
             room_id, latest_event_ids, await_full_state=False
         )
@@ -904,9 +902,7 @@ async def update_membership_locked(
                 raise AuthError(403, "The target user is not in the room")
 
         is_partial_state_room = await self.store.is_partial_state_room(room_id)
-        is_host_in_room = await self._is_host_in_room(
-            is_partial_state_room, state_before_join
-        )
+        is_host_in_room = await self._is_host_in_room(state_before_join)
 
         if effective_membership_state == Membership.JOIN:
             if requester.is_guest:
@@ -1670,13 +1666,7 @@ async def _make_and_store_3pid_invite(
         )
         return event, stream_id
 
-    async def _is_host_in_room(
-        self, is_partial_state_room: bool, current_state_ids: StateMap[str]
-    ) -> bool:
-        # When we only have a partial state, let's assume we are still in the room
-        if is_partial_state_room:
-            return True
-
+    async def _is_host_in_room(self, current_state_ids: StateMap[str]) -> bool:
         # Have we just created the room, and is this about to be the very
         # first member event?
         create_event_id = current_state_ids.get(("m.room.create", ""))

From 4fcf84ef2dc4ac540f73669799102a47264eaf71 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Sat, 4 Feb 2023 23:12:13 +0000
Subject: [PATCH 17/19] make partial statedness of state maps explicit

---
 synapse/handlers/event_auth.py  |  7 ++--
 synapse/handlers/room_member.py | 64 +++++++++++++++++++++------------
 2 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/synapse/handlers/event_auth.py b/synapse/handlers/event_auth.py
index 08709c65be7c..46dd63c3f00d 100644
--- a/synapse/handlers/event_auth.py
+++ b/synapse/handlers/event_auth.py
@@ -254,13 +254,14 @@ async def check_restricted_join_rules(
             )
 
     async def has_restricted_join_rules(
-        self, state_ids: StateMap[str], room_version: RoomVersion
+        self, partial_state_ids: StateMap[str], room_version: RoomVersion
     ) -> bool:
         """
         Return if the room has the proper join rules set for access via rooms.
 
         Args:
-            state_ids: The state of the room as it currently is.
+            state_ids: The state of the room as it currently is. May be full or partial
+                state.
             room_version: The room version of the room to query.
 
         Returns:
@@ -271,7 +272,7 @@ async def has_restricted_join_rules(
             return False
 
         # If there's no join rule, then it defaults to invite (so this doesn't apply).
-        join_rules_event_id = state_ids.get((EventTypes.JoinRules, ""), None)
+        join_rules_event_id = partial_state_ids.get((EventTypes.JoinRules, ""), None)
         if not join_rules_event_id:
             return False
 
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index cd9bab24ef42..258a266d8331 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -839,18 +839,19 @@ async def update_membership_locked(
 
         latest_event_ids = await self.store.get_prev_events_for_room(room_id)
 
-        # The full state is needed to calculate if another user trying to join is
-        # allowed to. In this case we now make another `make/send_join` to a resident
-        # server to validate that. We could just accept it but it would mean that we
-        # would potentially leak the history to a banned user.
-        state_before_join = await self.state_handler.compute_state_after_events(
+        is_partial_state_room = await self.store.is_partial_state_room(room_id)
+        partial_state_before_join = await self.state_handler.compute_state_after_events(
             room_id, latest_event_ids, await_full_state=False
         )
+        # `is_partial_state_room` also indicates whether `partial_state_before_join` is
+        # partial.
 
         # TODO: Refactor into dictionary of explicitly allowed transitions
         # between old and new state, with specific error messages for some
         # transitions and generic otherwise
-        old_state_id = state_before_join.get((EventTypes.Member, target.to_string()))
+        old_state_id = partial_state_before_join.get(
+            (EventTypes.Member, target.to_string())
+        )
         if old_state_id:
             old_state = await self.store.get_event(old_state_id, allow_none=True)
             old_membership = old_state.content.get("membership") if old_state else None
@@ -901,12 +902,11 @@ async def update_membership_locked(
             if action == "kick":
                 raise AuthError(403, "The target user is not in the room")
 
-        is_partial_state_room = await self.store.is_partial_state_room(room_id)
-        is_host_in_room = await self._is_host_in_room(state_before_join)
+        is_host_in_room = await self._is_host_in_room(partial_state_before_join)
 
         if effective_membership_state == Membership.JOIN:
             if requester.is_guest:
-                guest_can_join = await self._can_guest_join(state_before_join)
+                guest_can_join = await self._can_guest_join(partial_state_before_join)
                 if not guest_can_join:
                     # This should be an auth check, but guests are a local concept,
                     # so don't really fit into the general auth process.
@@ -950,7 +950,7 @@ async def update_membership_locked(
                 content,
                 is_partial_state_room,
                 is_host_in_room,
-                state_before_join,
+                partial_state_before_join,
             )
             if remote_join:
                 if ratelimit:
@@ -1097,7 +1097,7 @@ async def _should_perform_remote_join(
         content: JsonDict,
         is_partial_state_room: bool,
         is_host_in_room: bool,
-        state_before_join: StateMap[str],
+        partial_state_before_join: StateMap[str],
     ) -> Tuple[bool, List[str]]:
         """
         Check whether the server should do a remote join (as opposed to a local
@@ -1119,8 +1119,9 @@ async def _should_perform_remote_join(
             is_partial_state_room: True if the server currently doesn't hold the fully
                 state of the room.
             is_host_in_room: True if the host is in the room.
-            state_before_join: The state before the join event (i.e. the resolution of
-                the states after its parent events).
+            partial_state_before_join: The state before the join event (i.e. the
+                resolution of the states after its parent events). May be full or
+                partial state, depending on `is_partial_state_room`.
 
         Returns:
             A tuple of:
@@ -1134,7 +1135,9 @@ async def _should_perform_remote_join(
         if not is_host_in_room:
             return True, remote_room_hosts
 
-        prev_member_event_id = state_before_join.get((EventTypes.Member, user_id), None)
+        prev_member_event_id = partial_state_before_join.get(
+            (EventTypes.Member, user_id), None
+        )
         previous_membership = None
         if prev_member_event_id:
             prev_member_event = await self.store.get_event(prev_member_event_id)
@@ -1155,7 +1158,7 @@ async def _should_perform_remote_join(
         # If restricted join rules are not being used, a local join can always
         # be used.
         if not await self.event_auth_handler.has_restricted_join_rules(
-            state_before_join, room_version
+            partial_state_before_join, room_version
         ):
             return False, []
 
@@ -1164,6 +1167,11 @@ async def _should_perform_remote_join(
         if previous_membership in (Membership.JOIN, Membership.INVITE):
             return False, []
 
+        # All the partial state cases are covered above. We have been given the full
+        # state of the room.
+        assert not is_partial_state_room
+        state_before_join = partial_state_before_join
+
         # If the local host has a user who can issue invites, then a local
         # join can be done.
         #
@@ -1336,11 +1344,17 @@ async def send_membership_event(
                 if prev_member_event.membership == Membership.JOIN:
                     await self._user_left_room(target_user, room_id)
 
-    async def _can_guest_join(self, current_state_ids: StateMap[str]) -> bool:
+    async def _can_guest_join(self, partial_current_state_ids: StateMap[str]) -> bool:
         """
         Returns whether a guest can join a room based on its current state.
+
+        Args:
+            partial_current_state_ids: The current state of the room. May be full or
+                partial state.
         """
-        guest_access_id = current_state_ids.get((EventTypes.GuestAccess, ""), None)
+        guest_access_id = partial_current_state_ids.get(
+            (EventTypes.GuestAccess, ""), None
+        )
         if not guest_access_id:
             return False
 
@@ -1666,19 +1680,25 @@ async def _make_and_store_3pid_invite(
         )
         return event, stream_id
 
-    async def _is_host_in_room(self, current_state_ids: StateMap[str]) -> bool:
+    async def _is_host_in_room(self, partial_current_state_ids: StateMap[str]) -> bool:
+        """Returns whether the homeserver is in the room based on its current state.
+
+        Args:
+            partial_current_state_ids: The current state of the room. May be full or
+                partial state.
+        """
         # Have we just created the room, and is this about to be the very
         # first member event?
-        create_event_id = current_state_ids.get(("m.room.create", ""))
-        if len(current_state_ids) == 1 and create_event_id:
+        create_event_id = partial_current_state_ids.get(("m.room.create", ""))
+        if len(partial_current_state_ids) == 1 and create_event_id:
             # We can only get here if we're in the process of creating the room
             return True
 
-        for etype, state_key in current_state_ids:
+        for etype, state_key in partial_current_state_ids:
             if etype != EventTypes.Member or not self.hs.is_mine_id(state_key):
                 continue
 
-            event_id = current_state_ids[(etype, state_key)]
+            event_id = partial_current_state_ids[(etype, state_key)]
             event = await self.store.get_event(event_id, allow_none=True)
             if not event:
                 continue

From 4da4edad381716f5d786a3a56546dbda37a27780 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Sat, 4 Feb 2023 23:16:38 +0000
Subject: [PATCH 18/19] fix typos and line wrapping

---
 synapse/handlers/federation_event.py |  4 ++--
 synapse/handlers/room_member.py      | 13 +++++++------
 2 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index ccdd8bf6b6f5..8bad1a45e525 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -526,8 +526,8 @@ async def process_remote_join(
         )
         with nested_logging_context(suffix=event.event_id):
             if partial_state:
-                # When handling a second partial join into a partial state room, the
-                # returned state will exclude the membership from the first join. To
+                # When handling a second partial state join into a partial state room,
+                # the returned state will exclude the membership from the first join. To
                 # preserve prior memberships, we try to compute the state before the
                 # event ourselves if we know about any of the prev events.
                 #
diff --git a/synapse/handlers/room_member.py b/synapse/handlers/room_member.py
index 258a266d8331..f8d739752525 100644
--- a/synapse/handlers/room_member.py
+++ b/synapse/handlers/room_member.py
@@ -1116,9 +1116,9 @@ async def _should_perform_remote_join(
             remote_room_hosts: A list of remote room hosts.
             content: The content to use as the event body of the join. This may
                 be modified.
-            is_partial_state_room: True if the server currently doesn't hold the fully
+            is_partial_state_room: `True` if the server currently doesn't hold the full
                 state of the room.
-            is_host_in_room: True if the host is in the room.
+            is_host_in_room: `True` if the host is in the room.
             partial_state_before_join: The state before the join event (i.e. the
                 resolution of the states after its parent events). May be full or
                 partial state, depending on `is_partial_state_room`.
@@ -1143,11 +1143,12 @@ async def _should_perform_remote_join(
             prev_member_event = await self.store.get_event(prev_member_event_id)
             previous_membership = prev_member_event.membership
 
-        # If we are not fully join yet, and the target is not already in the room,
+        # If we are not fully joined yet, and the target is not already in the room,
         # let's do a remote join so another server with the full state can validate
-        # that the user has not been ban for example.
-        # We could just accept the join and wait for state-res to resolve that later on
-        # but we would then leak room history to this person until then, which is pretty bad.
+        # that the user has not been banned for example.
+        # We could just accept the join and wait for state res to resolve that later on
+        # but we would then leak room history to this person until then, which is pretty
+        # bad.
         if is_partial_state_room and previous_membership != Membership.JOIN:
             return True, remote_room_hosts
 

From e97b44d064ccc23d1d582bf5d9e3fc94a59405b0 Mon Sep 17 00:00:00 2001
From: Sean Quah <seanq@matrix.org>
Date: Fri, 10 Feb 2023 01:37:49 +0000
Subject: [PATCH 19/19] fixup: comment wording

---
 synapse/handlers/federation_event.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/synapse/handlers/federation_event.py b/synapse/handlers/federation_event.py
index 8bad1a45e525..a2d29da02421 100644
--- a/synapse/handlers/federation_event.py
+++ b/synapse/handlers/federation_event.py
@@ -528,8 +528,8 @@ async def process_remote_join(
             if partial_state:
                 # When handling a second partial state join into a partial state room,
                 # the returned state will exclude the membership from the first join. To
-                # preserve prior memberships, we try to compute the state before the
-                # event ourselves if we know about any of the prev events.
+                # preserve prior memberships, we try to compute the partial state before
+                # the event ourselves if we know about any of the prev events.
                 #
                 # When we don't know about any of the prev events, it's fine to just use
                 # the returned state, since the new join will create a new forward