fix deadlock

Author: jackyzha0

src/replit_river/client_transport.py

@@ -17,6 +17,7 @@
     ERROR_SESSION,
     RiverException,
 )
+from replit_river.lock import AcquiredLock, assert_correct_lock
 from replit_river.messages import (
     PROTOCOL_VERSION,
     FailedSendingMessageException,
@@ -70,51 +71,66 @@ def __init__(
 
     async def close(self) -> None:
         self._rate_limiter.close()
-        await self._close_all_sessions()
+        async with self._session_lock() as session_lock:
+            await self._close_all_sessions(session_lock)
 
+    # get a session, ensuring its in connected state
     async def get_or_create_session(self) -> ClientSession:
-        async with self._create_session_lock:
-            existing_session = await self._get_existing_session()
+        async with self._session_lock() as session_lock:
+            existing_session = await self._get_existing_session(session_lock)
             if not existing_session:
-                return await self._create_new_session()
+                # create a new session
+                return await self._create_new_session(session_lock)
             is_session_open = await existing_session.is_session_open()
             if not is_session_open:
-                return await self._create_new_session()
-            is_ws_open = await existing_session.is_websocket_open()
-            if is_ws_open:
-                return existing_session
-            new_ws, _, hs_response = await self._establish_new_connection(
-                existing_session
-            )
-            if hs_response.status.sessionId == existing_session.session_id:
-                logger.info(
-                    "Replacing ws connection in session id %s",
-                    existing_session.session_id,
-                )
-                await existing_session.replace_with_new_websocket(new_ws)
-                return existing_session
-            else:
-                logger.info("Closing stale session %s", existing_session.session_id)
+                # replace the session with a new one
                 await existing_session.close()
-                return await self._create_new_session()
+                return await self._create_new_session(session_lock)
 
-    async def _get_existing_session(self) -> Optional[ClientSession]:
-        async with self._session_lock:
-            if not self._sessions:
-                return None
-            if len(self._sessions) > 1:
-                raise RiverException(
-                    "session_error",
-                    "More than one session found in client, should only be one",
-                )
-            session = list(self._sessions.values())[0]
-            if isinstance(session, ClientSession):
-                return session
-            else:
-                raise RiverException(
-                    "session_error", f"Client session type wrong, got {type(session)}"
+            # we have a websocket, check if its open
+            is_ws_open = await existing_session.is_websocket_open()
+            if not is_ws_open:
+                # dial a new connection
+                new_ws, _, hs_response = await self._establish_new_connection(
+                    existing_session
                 )
 
+                # session id is the same, it was a transparent reconnect
+                if hs_response.status.sessionId == existing_session.session_id:
+                    logger.info(
+                        "Replacing ws connection in session id %s",
+                        existing_session.session_id,
+                    )
+                    await existing_session.replace_with_new_websocket(new_ws)
+                    return existing_session
+                else:
+                    # session id is different, we need to close the old session
+                    logger.info("Closing stale session %s", existing_session.session_id)
+                    await existing_session.close()
+                    return await self._create_new_session(session_lock)
+
+            # happy path, we have a session and its open
+            return existing_session
+
+    async def _get_existing_session(
+        self, session_lock: AcquiredLock
+    ) -> Optional[ClientSession]:
+        assert_correct_lock(session_lock, self._session_lock)
+        if not self._sessions:
+            return None
+        if len(self._sessions) > 1:
+            raise RiverException(
+                "session_error",
+                "More than one session found in client, should only be one",
+            )
+        session = list(self._sessions.values())[0]
+        if isinstance(session, ClientSession):
+            return session
+        else:
+            raise RiverException(
+                "session_error", f"Client session type wrong, got {type(session)}"
+            )
+
     async def _establish_new_connection(
         self,
         old_session: Optional[ClientSession] = None,
@@ -186,35 +202,39 @@ async def _establish_new_connection(
 
     async def _create_new_session(
         self,
+        session_lock: AcquiredLock,
     ) -> ClientSession:
-        logger.info("Creating new session")
-        new_ws, hs_request, hs_response = await self._establish_new_connection()
-        if not hs_response.status.ok:
-            message = hs_response.status.reason
-            raise RiverException(
-                ERROR_SESSION,
-                f"Server did not return OK status on handshake response: {message}",
+        async with self._create_session_lock:
+            logger.info("Creating new session")
+            new_ws, hs_request, hs_response = await self._establish_new_connection()
+            if not hs_response.status.ok:
+                message = hs_response.status.reason
+                raise RiverException(
+                    ERROR_SESSION,
+                    f"Server did not return OK status on handshake response: {message}",
+                )
+            new_session = ClientSession(
+                transport_id=self._transport_id,
+                to_id=self._server_id,
+                session_id=hs_request.sessionId,
+                websocket=new_ws,
+                transport_options=self._transport_options,
+                is_server=False,
+                close_session_callback=self.lock_and_delete_session,
+                retry_connection_callback=self._retry_connection,
+                handlers={},
             )
-        new_session = ClientSession(
-            transport_id=self._transport_id,
-            to_id=self._server_id,
-            session_id=hs_request.sessionId,
-            websocket=new_ws,
-            transport_options=self._transport_options,
-            is_server=False,
-            close_session_callback=self._delete_session,
-            retry_connection_callback=self._retry_connection,
-            handlers={},
-        )
 
-        self._set_session(new_session)
-        await new_session.start_serve_responses()
-        return new_session
+            self._set_session(session_lock, new_session)
+            await new_session.start_serve_responses()
+            return new_session
 
     async def _retry_connection(self) -> ClientSession:
-        if not self._transport_options.transparent_reconnect:
-            await self._close_all_sessions()
-        return await self.get_or_create_session()
+        if self._transport_options.transparent_reconnect:
+            return await self.get_or_create_session()
+
+        async with self._session_lock() as session_lock:
+            return await self._create_new_session(session_lock)
 
     async def _send_handshake_request(
         self,
@@ -344,7 +364,6 @@ async def _establish_handshake(
                 # If the session status is mismatched, we should close the old session
                 # and let the retry logic to create a new session.
                 await old_session.close()
-                await self._delete_session(old_session)
 
             raise RiverException(
                 ERROR_HANDSHAKE,

src/replit_river/lock.py

@@ -0,0 +1,40 @@
+import asyncio
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator, ParamSpec, TypeVar
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class AcquiredLock:
+    def __init__(self, lock: asyncio.Lock):
+        self._lock = lock
+
+
+class TransferableLock:
+    """A lock that can be transferred between coroutines."""
+
+    def __init__(self) -> None:
+        self._lock = asyncio.Lock()
+
+        class BoundLock(AcquiredLock):
+            pass
+
+        self.BoundLock = BoundLock
+
+    @asynccontextmanager
+    async def __call__(self) -> AsyncGenerator[AcquiredLock]:
+        await self._lock.acquire()
+        yield self.BoundLock(self._lock)
+        self._lock.release()
+
+
+def assert_correct_lock(
+    acquired_lock: AcquiredLock, expected_lock_class: TransferableLock
+) -> None:
+    """Assert that the acquired lock is the correct class."""
+    if not isinstance(acquired_lock, expected_lock_class.BoundLock):
+        raise ValueError(
+            f"Expected {expected_lock_class.BoundLock.__name__}, "
+            f"got {type(acquired_lock).__name__}"
+        )

src/replit_river/seq_manager.py

@@ -87,9 +87,4 @@ async def check_seq_and_update(self, msg: TransportMessage) -> None:
                         f"{self.ack}"
                     )
             self.receiver_ack = msg.ack
-        await self._set_ack(msg.seq + 1)
-
-    async def _set_ack(self, new_ack: int) -> int:
-        async with self._ack_lock:
-            self.ack = new_ack
-            return self.ack
+            self.ack = msg.seq + 1

src/replit_river/server_transport.py

@@ -88,7 +88,8 @@ async def handshake_to_get_session(
         raise WebsocketClosedException("No handshake message received")
 
     async def close(self) -> None:
-        await self._close_all_sessions()
+        async with self._session_lock() as session_lock:
+            await self._close_all_sessions(session_lock)
 
     async def _get_or_create_session(
         self,
@@ -97,7 +98,7 @@ async def _get_or_create_session(
         session_id: str,
         websocket: WebSocketCommonProtocol,
     ) -> Session:
-        async with self._session_lock:
+        async with self._session_lock() as session_lock:
             session_to_close: Optional[Session] = None
             new_session: Optional[Session] = None
             if to_id not in self._sessions:
@@ -112,7 +113,7 @@ async def _get_or_create_session(
                     self._transport_options,
                     self._is_server,
                     self._handlers,
-                    close_session_callback=self._delete_session,
+                    close_session_callback=self.lock_and_delete_session,
                 )
             else:
                 old_session = self._sessions[to_id]
@@ -133,7 +134,7 @@ async def _get_or_create_session(
                         self._transport_options,
                         self._is_server,
                         self._handlers,
-                        close_session_callback=self._delete_session,
+                        close_session_callback=self.lock_and_delete_session,
                     )
                 else:
                     # If the instance id is the same, we reuse the session and assign
@@ -152,7 +153,7 @@ async def _get_or_create_session(
             if session_to_close:
                 logger.info("Closing stale session %s", session_to_close.session_id)
                 await session_to_close.close()
-            self._set_session(new_session)
+            self._set_session(session_lock, new_session)
         return new_session
 
     async def _send_handshake_response(
@@ -228,68 +229,64 @@ async def _establish_handshake(
             )
             raise InvalidMessageException("handshake request to wrong server")
 
-        async with self._session_lock:
-            old_session = self._sessions.get(request_message.from_, None)
-            client_next_expected_seq = (
-                handshake_request.expectedSessionState.nextExpectedSeq
-            )
-            client_next_sent_seq = (
-                handshake_request.expectedSessionState.nextSentSeq or 0
-            )
-            if old_session and old_session.session_id == handshake_request.sessionId:
-                # check invariants
-                # ordering must be correct
-                our_next_seq = await old_session.get_next_sent_seq()
-                our_ack = await old_session.get_next_expected_seq()
-
-                if client_next_sent_seq > our_ack:
-                    message = (
-                        "client is in the future: "
-                        f"server wanted {our_ack} but client has {client_next_sent_seq}"
-                    )
-                    await self._send_handshake_response(
-                        request_message,
-                        HandShakeStatus(ok=False, reason=message),
-                        websocket,
-                    )
-                    raise SessionStateMismatchException(message)
+        old_session = self._sessions.get(request_message.from_, None)
+        client_next_expected_seq = (
+            handshake_request.expectedSessionState.nextExpectedSeq
+        )
+        client_next_sent_seq = handshake_request.expectedSessionState.nextSentSeq or 0
+        if old_session and old_session.session_id == handshake_request.sessionId:
+            # check invariants
+            # ordering must be correct
+            our_next_seq = await old_session.get_next_sent_seq()
+            our_ack = await old_session.get_next_expected_seq()
 
-                if our_next_seq > client_next_expected_seq:
-                    message = (
-                        "server is in the future: "
-                        f"client wanted {client_next_expected_seq} "
-                        f"but server has {our_next_seq}"
-                    )
-                    await self._send_handshake_response(
-                        request_message,
-                        HandShakeStatus(ok=False, reason=message),
-                        websocket,
-                    )
-                    raise SessionStateMismatchException(message)
-            elif old_session:
-                # we have an old session but the session id is different
-                # just delete the old session
-                await old_session.close()
-                await self._delete_session(old_session)
-                old_session = None
+            if client_next_sent_seq > our_ack:
+                message = (
+                    "client is in the future: "
+                    f"server wanted {our_ack} but client has {client_next_sent_seq}"
+                )
+                await self._send_handshake_response(
+                    request_message,
+                    HandShakeStatus(ok=False, reason=message),
+                    websocket,
+                )
+                raise SessionStateMismatchException(message)
 
-            if not old_session and (
-                client_next_sent_seq > 0 or client_next_expected_seq > 0
-            ):
-                message = "client is trying to resume a session but we don't have it"
+            if our_next_seq > client_next_expected_seq:
+                message = (
+                    "server is in the future: "
+                    f"client wanted {client_next_expected_seq} "
+                    f"but server has {our_next_seq}"
+                )
                 await self._send_handshake_response(
                     request_message,
                     HandShakeStatus(ok=False, reason=message),
                     websocket,
                 )
                 raise SessionStateMismatchException(message)
+        elif old_session:
+            # we have an old session but the session id is different
+            # just delete the old session
+            await old_session.close()
+            old_session = None
 
-            # from this point on, we're committed to connecting
-            session_id = handshake_request.sessionId
-            handshake_response = await self._send_handshake_response(
+        if not old_session and (
+            client_next_sent_seq > 0 or client_next_expected_seq > 0
+        ):
+            message = "client is trying to resume a session but we don't have it"
+            await self._send_handshake_response(
                 request_message,
-                HandShakeStatus(ok=True, sessionId=session_id),
+                HandShakeStatus(ok=False, reason=message),
                 websocket,
             )
+            raise SessionStateMismatchException(message)
+
+        # from this point on, we're committed to connecting
+        session_id = handshake_request.sessionId
+        handshake_response = await self._send_handshake_response(
+            request_message,
+            HandShakeStatus(ok=True, sessionId=session_id),
+            websocket,
+        )
 
-            return handshake_request, handshake_response
+        return handshake_request, handshake_response