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

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

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

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev4}/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.dev4"
 description="Replit river toolkit for Python"
 authors = ["Replit <eng@replit.com>"]
 license = "LICENSE"

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

@@ -33,7 +33,6 @@ from replit_river.seq_manager import (
     IgnoreMessageException,
     InvalidMessageException,
 )
-from replit_river.session import Session
 from replit_river.transport import Transport
 from replit_river.transport_options import TransportOptions
 
@@ -59,12 +58,6 @@ class ClientTransport(Transport):
         )
         # 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()
@@ -89,6 +82,7 @@ class ClientTransport(Transport):
 
     async def _establish_new_connection(
         self,
+        old_session: Optional[ClientSession] = None,
     ) -> Tuple[
         WebSocketCommonProtocol,
         ControlMessageHandshakeRequest,
@@ -101,18 +95,25 @@ class ClientTransport(Transport):
         for i in range(max_retry):
             if i > 0:
                 logging.info(f"Retrying build handshake number {i} times")
+                logging.info(
+                    f"old_session: {old_session}, old_session.is_session_open(): {await old_session.is_session_open() if old_session else None}"
+                )
             if not rate_limit.has_budget(client_id):
                 logging.debug("No retry budget for %s.", client_id)
                 break
             try:
+                logging.error(
+                    f"##### _establish_new_connection: old session : {old_session}"
+                )
                 ws = await websockets.connect(self._websocket_uri)
-                existing_session = await self._get_existing_session()
                 session_id = (
                     self.generate_session_id()
-                    if not existing_session
-                    else existing_session.session_id
+                    if not old_session
+                    else old_session.session_id
+                )
+                logging.error(
+                    f"##### _establish_new_connection: existing session : {old_session}"
                 )
-                logging.error(f"##### 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
@@ -127,97 +128,70 @@ class ClientTransport(Transport):
                 await asyncio.sleep(backoff_time / 1000)
         raise RiverException(
             ERROR_HANDSHAKE,
-            "Failed to create session after retrying max number of times",
+            "Failed to create ws 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 def _get_or_create_session(self) -> ClientSession:
+        logging.error(f"####### start get or create session")
         async with self._create_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")
-                    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
+            if not existing_session:
+                logging.error(f"##### _get_or_create_session 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"##### _get_or_create_session session open, creating new session"
+                )
+                return await self._create_new_session()
+            is_ws_open = await existing_session.is_websocket_open()
+            if is_ws_open:
+                logging.error(f"##### _get_or_create_session Reuse existing session")
+                return existing_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",
-                    )
-                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
-                    ),
+                new_ws, _, hs_response = await self._establish_new_connection(
+                    existing_session
                 )
-                await self._set_session(new_session)
-                await new_session.start_serve_responses()
-                return new_session
+                if (
+                    hs_response.status.sessionId
+                    == existing_session.advertised_session_id
+                ):
+                    logging.error(
+                        f"##### _get_or_create_session session open, replacing websocket"
+                    )
+                    await existing_session.replace_with_new_websocket(new_ws)
+                    return existing_session
+                else:
+                    logging.error(f"##### session open, not same session id, reuse")
+                    await existing_session.close(is_unexpected_close=False)
+                    return await self._create_new_session()
 
     async def _send_handshake_request(
         self,

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev4}/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,18 +104,10 @@ 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.
         """
-        logging.debug("begin_close_session_countdown")
         if self._close_session_after_time_secs is not None:
             # already in grace period, no need to set again
             return
@@ -128,6 +120,7 @@ class Session(object):
         self._close_session_after_time_secs = (
             await self._get_current_time() + grace_period_ms / 1000
         )
+        await self.close_websocket(self._ws_wrapper, should_retry=not self._is_server)
 
     async def serve(self) -> None:
         """Serve messages from the websocket."""
@@ -137,7 +130,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.
@@ -170,8 +163,15 @@ class Session(object):
             self._ws_wrapper.id,
         )
         try:
-            async for message in self._ws_wrapper.ws:
+            ws_wrapper = self._ws_wrapper
+            async for message in ws_wrapper.ws:
                 try:
+                    logging.error(
+                        f" await ws_wrapper.is_open(): : { await ws_wrapper.is_open()}"
+                    )
+                    if not await ws_wrapper.is_open():
+                        # We should not process messages if the websocket is closed.
+                        break
                     msg = parse_transport_msg(message, self._transport_options)
 
                     logging.debug(f"{self._transport_id} got a message %r", msg)
@@ -241,9 +241,13 @@ class Session(object):
             await asyncio.sleep(
                 self._transport_options.close_session_check_interval_ms / 1000
             )
+            logging.error("#### _check_to_close_session")
             if not self._close_session_after_time_secs:
                 continue
             current_time = await self._get_current_time()
+            logging.error(
+                f"#### _check_to_close_session : current_time: {current_time} self._close_session_after_time_secs: {self._close_session_after_time_secs}, {current_time > self._close_session_after_time_secs}"
+            )
             if current_time > self._close_session_after_time_secs:
                 logging.debug(
                     "Grace period ended for %s, closing session", self._transport_id
@@ -288,14 +292,14 @@ class Session(object):
                     self._heartbeat_misses
                     >= self._transport_options.heartbeats_until_dead
                 ):
+                    if self._close_session_after_time_secs is not None:
+                        # already in grace period, no need to set again
+                        continue
                     logging.debug(
                         "%r closing websocket because of heartbeat misses",
                         self.session_id,
                     )
-                    await self._on_websocket_unexpected_close()
-                    await self.close_websocket(
-                        self._ws_wrapper, should_retry=not self._is_server
-                    )
+                    await self._begin_close_session_countdown()
                     continue
             except FailedSendingMessageException:
                 # this is expected during websocket closed period
@@ -327,7 +331,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,9 +420,13 @@ 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)
+            logging.error("### running retry_connection_callback")
+            await self._task_manager.create_task(self._retry_connection_callback(self))
 
     async def _open_stream_and_call_handler(
         self,
@@ -495,7 +503,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 +518,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.dev4}/replit_river/transport.py

@@ -45,13 +45,22 @@ 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:
+    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:
             self._sessions[session._to_id] = session
 
     def generate_nanoid(self) -> str:
@@ -87,6 +96,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 +148,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

{replit_river-0.1.16.dev2 → replit_river-0.1.16.dev4}/replit_river/transport_options.py

@@ -13,7 +13,7 @@ class ConnectionRetryOptions(BaseModel):
     max_backoff_ms: float = 32_000
     attempt_budget_capacity: float = 5
     budget_restore_interval_ms: float = 200
-    max_retry: int = 10
+    max_retry: int = 100
 
 
 # setup in replit web can be found at