replit-river 0.1.16.dev2__tar.gz → 0.1.16.dev3__tar.gz

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: replit-river
-Version: 0.1.16.dev2
+Version: 0.1.16.dev3
 Summary: Replit river toolkit for Python
 License: LICENSE
 Keywords: rpc,websockets

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name="replit-river"
-version="0.1.16.dev2"
+version="0.1.16.dev3"
 description="Replit river toolkit for Python"
 authors = ["Replit <eng@replit.com>"]
 license = "LICENSE"

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev3}/replit_river/client_transport.py

@@ -57,35 +57,26 @@ class ClientTransport(Transport):
         self._rate_limiter = LeakyBucketRateLimit(
             transport_options.connection_retry_options
         )
-        # We want to make sure there's only one session creation at a time
-        self._create_session_lock = asyncio.Lock()
-        # Only one retry should happen at a time
-        self._retry_ws_lock = asyncio.Lock()
-
-    async def _on_session_closed(self, session: Session) -> None:
-        logging.info(f"Client session {session.advertised_session_id} closed")
-        await self._delete_session(session)
 
     async def close(self) -> None:
         self._rate_limiter.close()
         await self._close_all_sessions()
 
     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)}"
-                )
+        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,
@@ -112,7 +103,9 @@ class ClientTransport(Transport):
                     if not existing_session
                     else existing_session.session_id
                 )
-                logging.error(f"##### existing session : {existing_session}")
+                logging.error(
+                    f"##### _establish_new_connection: existing session : {existing_session}"
+                )
                 rate_limit.consume_budget(client_id)
                 handshake_request, handshake_response = await self._establish_handshake(
                     self._transport_id, self._server_id, session_id, ws
@@ -130,94 +123,65 @@ class ClientTransport(Transport):
             "Failed to create session after retrying max number of times",
         )
 
-    async def _retry_session_connection(
-        self, session_to_replace_ws: Session
-    ) -> Session:
-        async with self._retry_ws_lock:
-            if await session_to_replace_ws.is_websocket_open():
-                # other retry successfully replaced the websocket,
-                return session_to_replace_ws
-            if not await session_to_replace_ws.is_session_open():
-                # If the session is already closing we don't retry connection
-                return session_to_replace_ws
-            new_ws, hs_request, hs_response = await self._establish_new_connection()
-            # If the server session id different, we create a new session.
-            if (
-                hs_response.status.sessionId
-                != session_to_replace_ws.advertised_session_id
-            ):
-                server_session_id = hs_response.status.sessionId
-                if not server_session_id:
-                    raise RiverException(
-                        ERROR_SESSION,
-                        "Server did not return a sessionId in successful handshake",
-                    )
-                new_session = ClientSession(
-                    transport_id=self._transport_id,
-                    to_id=self._server_id,
-                    session_id=hs_request.sessionId,
-                    advertised_session_id=server_session_id,
-                    websocket=new_ws,
-                    transport_options=self._transport_options,
-                    is_server=False,
-                    handlers={},
-                    close_session_callback=self._on_session_closed,
-                    retry_connection_callback=lambda x: self._retry_session_connection(
-                        x
-                    ),
-                )
-                return new_session
-            else:
-                # If the session is still active and aligns with the server session
-                # we replace the websocket in it.
-                await session_to_replace_ws.replace_with_new_websocket(new_ws)
-            return session_to_replace_ws
+    async def _create_new_session(
+        self,
+    ) -> ClientSession:
+        new_ws, hs_request, hs_response = await self._establish_new_connection()
+        advertised_session_id = hs_response.status.sessionId
+        if not advertised_session_id:
+            raise RiverException(
+                ERROR_SESSION,
+                "Server did not return a sessionId in successful handshake",
+            )
+        new_session = ClientSession(
+            transport_id=self._transport_id,
+            to_id=self._server_id,
+            session_id=hs_request.sessionId,
+            advertised_session_id=advertised_session_id,
+            websocket=new_ws,
+            transport_options=self._transport_options,
+            is_server=False,
+            handlers={},
+            close_session_callback=self._delete_session,
+            retry_connection_callback=lambda x: self._get_or_create_session(),
+        )
 
-    async def _get_or_create_session(self) -> ClientSession:
-        logging.error(f"####### _get_or_create_session")
+        await self._set_session(new_session, acquire_lock=False)
+        await new_session.start_serve_responses()
+        return new_session
 
-        async with self._create_session_lock:
+    async def _get_or_create_session(self) -> ClientSession:
+        async with self._session_lock:
             existing_session = await self._get_existing_session()
-            if existing_session:
-                if await existing_session.is_websocket_open():
-                    logging.error(f"####### use existing session")
+            if not existing_session:
+                logging.error(f"##### No existing session")
+                return await self._create_new_session()
+            is_session_open = await existing_session.is_session_open()
+            if not is_session_open:
+                logging.error(f"##### session open, creating new session")
+                await existing_session.close(
+                    is_unexpected_close=False, acquire_transport_lock=False
+                )
+                return await self._create_new_session()
+            is_ws_open = existing_session.is_websocket_open()
+            if is_ws_open:
+                logging.error(f"##### Reuse existing session")
+                return existing_session
+            else:
+                new_ws, _, hs_response = await self._establish_new_connection()
+                if (
+                    hs_response.status.sessionId
+                    == existing_session.advertised_session_id
+                ):
+                    logging.error(f"##### session open, replacing websocket")
+                    await existing_session.replace_with_new_websocket(new_ws)
                     return existing_session
                 else:
-                    logging.error(f"####### retry session connection")
-                    session = await self._retry_session_connection(existing_session)
-                    # This should never happen, adding here to make mypy happy
-                    if not isinstance(session, ClientSession):
-                        raise RiverException(
-                            ERROR_SESSION,
-                            f"Session type is not ClientSession, got {type(session)}",
-                        )
-                    return session
-            else:
-                logging.error(f"####### establish new connection")
-                new_ws, hs_request, hs_response = await self._establish_new_connection()
-                advertised_session_id = hs_response.status.sessionId
-                if not advertised_session_id:
-                    raise RiverException(
-                        ERROR_SESSION,
-                        "Server did not return a sessionId in successful handshake",
+                    logging.error(f"##### session open, not same session id, reuse")
+                    await existing_session.close(
+                        is_unexpected_close=False, acquire_transport_lock=False
                     )
-                new_session = ClientSession(
-                    transport_id=self._transport_id,
-                    to_id=self._server_id,
-                    session_id=hs_request.sessionId,
-                    advertised_session_id=advertised_session_id,
-                    websocket=new_ws,
-                    transport_options=self._transport_options,
-                    is_server=False,
-                    handlers={},
-                    close_session_callback=self._on_session_closed,
-                    retry_connection_callback=lambda x: self._retry_session_connection(
-                        x
-                    ),
-                )
-                await self._set_session(new_session)
-                await new_session.start_serve_responses()
-                return new_session
+                    return await self._create_new_session()
 
     async def _send_handshake_request(
         self,

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev3}/replit_river/session.py

@@ -52,7 +52,7 @@ class Session(object):
         transport_options: TransportOptions,
         is_server: bool,
         handlers: Dict[Tuple[str, str], Tuple[str, GenericRpcHandler]],
-        close_session_callback: Callable[["Session"], Coroutine[Any, Any, None]],
+        close_session_callback: Callable[["Session", bool], Coroutine[Any, Any, None]],
         retry_connection_callback: Optional[
             Callable[
                 ["Session"],
@@ -104,13 +104,6 @@ class Session(object):
         async with self._ws_lock:
             return await self._ws_wrapper.is_open()
 
-    async def _on_websocket_unexpected_close(self) -> None:
-        """Handle unexpected websocket close."""
-        logging.debug(
-            "Unexpected websocket close from %s to %s", self._transport_id, self._to_id
-        )
-        await self._begin_close_session_countdown()
-
     async def _begin_close_session_countdown(self) -> None:
         """Begin the countdown to close session, this should be called when
         websocket is closed.
@@ -119,6 +112,7 @@ class Session(object):
         if self._close_session_after_time_secs is not None:
             # already in grace period, no need to set again
             return
+        await self._ws_wrapper.close()
         logging.debug(
             "websocket closed from %s to %s begin grace period",
             self._transport_id,
@@ -137,7 +131,7 @@ class Session(object):
                 try:
                     await self._handle_messages_from_ws(tg)
                 except ConnectionClosed as e:
-                    await self._on_websocket_unexpected_close()
+                    await self._begin_close_session_countdown()
                     logging.debug("ConnectionClosed while serving: %r", e)
                 except FailedSendingMessageException as e:
                     # Expected error if the connection is closed.
@@ -292,7 +286,7 @@ class Session(object):
                         "%r closing websocket because of heartbeat misses",
                         self.session_id,
                     )
-                    await self._on_websocket_unexpected_close()
+                    await self._begin_close_session_countdown()
                     await self.close_websocket(
                         self._ws_wrapper, should_retry=not self._is_server
                     )
@@ -327,7 +321,7 @@ class Session(object):
     ) -> None:
         try:
             await send_transport_message(
-                msg, websocket, self._on_websocket_unexpected_close, prefix_bytes
+                msg, websocket, self._begin_close_session_countdown, prefix_bytes
             )
         except WebsocketClosedException as e:
             raise e
@@ -416,6 +410,9 @@ class Session(object):
     ) -> None:
         """Mark the websocket as closed, close the websocket, and retry if needed."""
         async with self._ws_lock:
+            # Already closed.
+            if not await ws_wrapper.is_open():
+                return
             await ws_wrapper.close()
         if should_retry and self._retry_connection_callback:
             await self._retry_connection_callback(self)
@@ -495,7 +492,9 @@ class Session(object):
     async def start_serve_responses(self) -> None:
         await self._task_manager.create_task(self.serve())
 
-    async def close(self, is_unexpected_close: bool) -> None:
+    async def close(
+        self, is_unexpected_close: bool, acquire_transport_lock: bool = True
+    ) -> None:
         """Close the session and all associated streams."""
         logging.info(
             f"{self._transport_id} closing session "
@@ -508,9 +507,10 @@ class Session(object):
                 return
             self._state = SessionState.CLOSING
             self._reset_session_close_countdown()
-            await self.close_websocket(self._ws_wrapper, should_retry=False)
+            async with self._ws_lock:
+                await self._ws_wrapper.close()
             # Clear the session in transports
-            await self._close_session_callback(self)
+            await self._close_session_callback(self, acquire_transport_lock)
             await self._task_manager.cancel_all_tasks()
             # TODO: unexpected_close should close stream differently here to
             # throw exception correctly.

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev3}/replit_river/transport.py

@@ -45,14 +45,24 @@ class Transport:
         await asyncio.gather(*tasks)
         logging.info(f"Transport closed {self._transport_id}")
 
-    async def _delete_session(self, session: Session) -> None:
-        async with self._session_lock:
+    async def _delete_session(
+        self, session: Session, acquire_lock: bool = True
+    ) -> None:
+        if acquire_lock:
+            async with self._session_lock:
+                if session._to_id in self._sessions:
+                    del self._sessions[session._to_id]
+        else:
             if session._to_id in self._sessions:
                 del self._sessions[session._to_id]
 
-    async def _set_session(self, session: Session) -> None:
-        async with self._session_lock:
-            self._sessions[session._to_id] = session
+    async def _set_session(self, session: Session, acquire_lock: bool = True) -> None:
+        if acquire_lock:
+            async with self._session_lock:
+                self._sessions[session._to_id] = session
+        else:
+            if session._to_id in self._sessions:
+                del self._sessions[session._to_id]
 
     def generate_nanoid(self) -> str:
         return str(nanoid.generate())
@@ -87,6 +97,7 @@ class Transport:
     ) -> Session:
         session_to_close: Optional[Session] = None
         new_session: Optional[Session] = None
+        logging.error(f"## get_or_create_session, {to_id}")
         async with self._session_lock:
             if to_id not in self._sessions:
                 logging.debug(
@@ -138,13 +149,13 @@ class Transport:
                         new_session = old_session
                     except FailedSendingMessageException as e:
                         raise e
-        if session_to_close:
-            logging.debug(
-                "Closing stale session %s", session_to_close.advertised_session_id
-            )
-            await session_to_close.close(False)
-            logging.info(
-                f"Closed stale session {session_to_close.advertised_session_id}"
-            )
-        await self._set_session(new_session)
+
+            if session_to_close:
+                logging.debug(
+                    "Closing stale session %s", session_to_close.advertised_session_id
+                )
+                await session_to_close.close(
+                    is_unexpected_close=False, acquire_transport_lock=False
+                )
+            await self._set_session(new_session, acquire_lock=False)
         return new_session