deadlock patch (#148)

Why
===

deadlock in the server-side reconnect case:
1. establish handshake acquires session lock:
https://github.com/replit/river-python/blob/6e2bc1ab1392f5a8407a08f6c8f049ec9e4b308d/src/replit_river/server_transport.py#L249
2. if there is a session mismatch, we close the old session:
https://github.com/replit/river-python/blob/6e2bc1ab1392f5a8407a08f6c8f049ec9e4b308d/src/replit_river/server_transport.py#L290
3. session.close calls _close_session_callback
https://github.com/replit/river-python/blob/6e2bc1ab1392f5a8407a08f6c8f049ec9e4b308d/src/replit_river/session.py#L299
4. _delete_session also tries to acquire the session lock
https://github.com/replit/river-python/blob/6e2bc1ab1392f5a8407a08f6c8f049ec9e4b308d/src/replit_river/server_transport.py#L316

What changed
============

1. dont need to call _delete_session as .close will already do that
2. lift out session close outside of session lock
3. in the handshake case, let the final call to get_or_create_session
replace the session and close the old one

Test plan
=========

added a test

Notes
=========

I have a draft of a more in-depth approach on [this
branch](https://github.com/replit/river-python/tree/jackyzha0/fix-deadlock)
which uses a lock-ownership-transfer based approach that should catch it
more generically but ran into ownership problems lol

Seeing as we are planning on migrating chat service to Node anyways so
we only have one River server implementation, we hopefully don't have to
maintain this surface for much longer 🤞

---------

Co-authored-by: Devon Stewart <[email protected]>

Author: jackyzha0

src/replit_river/client_transport.py

@@ -366,7 +366,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/server_transport.py

@@ -117,10 +117,11 @@ async def _get_or_create_session(
         session_id: str,
         websocket: WebSocketCommonProtocol,
     ) -> ServerSession:
+        new_session: ServerSession | None = None
+        old_session: ServerSession | None = None
         async with self._session_lock:
-            session_to_close: Session | None = None
-            new_session: ServerSession | None = None
-            if to_id not in self._sessions:
+            old_session = self._sessions.get(to_id)
+            if not old_session:
                 logger.info(
                     'Creating new session with "%s" using ws: %s', to_id, websocket.id
                 )
@@ -134,7 +135,6 @@ async def _get_or_create_session(
                     close_session_callback=self._delete_session,
                 )
             else:
-                old_session = self._sessions[to_id]
                 if old_session.session_id != session_id:
                     logger.info(
                         'Create new session with "%s" for session id %s'
@@ -143,7 +143,6 @@ async def _get_or_create_session(
                         session_id,
                         old_session.session_id,
                     )
-                    session_to_close = old_session
                     new_session = ServerSession(
                         transport_id,
                         to_id,
@@ -167,10 +166,12 @@ async def _get_or_create_session(
                     except FailedSendingMessageException as e:
                         raise e
 
-            if session_to_close:
-                logger.info("Closing stale session %s", session_to_close.session_id)
-                await session_to_close.close()
             self._sessions[new_session._to_id] = new_session
+
+        if old_session and new_session != old_session:
+            logger.info("Closing stale session %s", old_session.session_id)
+            await old_session.close()
+
         return new_session
 
     async def _send_handshake_response(
@@ -247,7 +248,7 @@ 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)
+            old_session = self._sessions.get(request_message.from_)
             client_next_expected_seq = (
                 handshake_request.expectedSessionState.nextExpectedSeq
             )
@@ -285,10 +286,6 @@ async def _establish_handshake(
                     )
                     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 not old_session and (

tests/river_fixtures/clientserver.py

@@ -1,4 +1,3 @@
-import asyncio
 import logging
 from typing import AsyncGenerator, Literal
 
@@ -64,7 +63,6 @@ async def websocket_uri_factory() -> UriAndMetadata[None]:
             await client.close()
 
     finally:
-        await asyncio.sleep(1)
         logging.debug("Start closing test server")
         if binding:
             binding.close()

tests/test_communication.py

@@ -268,3 +268,32 @@ async def test_ignore_flood_subscription(client: Client) -> None:
         timedelta(seconds=20),
     )
     assert response == "Hello, Alice!"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("handlers", [{**basic_rpc_method}])
+async def test_rpc_method_reconnect(client: Client) -> None:
+    response = await client.send_rpc(
+        "test_service",
+        "rpc_method",
+        "Alice",
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+        timedelta(seconds=20),
+    )
+    assert response == "Hello, Alice!"
+
+    await client._transport._close_all_sessions()
+
+    response = await client.send_rpc(
+        "test_service",
+        "rpc_method",
+        "Bob",
+        serialize_request,
+        deserialize_response,
+        deserialize_error,
+        timedelta(seconds=20),
+    )
+
+    assert response == "Hello, Bob!"