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

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

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

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

{replit_river-0.1.16.dev3 → 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
 
@@ -57,29 +56,33 @@ 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()
 
     async def close(self) -> None:
         self._rate_limiter.close()
         await self._close_all_sessions()
 
     async def _get_existing_session(self) -> Optional[ClientSession]:
-        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 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)}"
+                )
 
     async def _establish_new_connection(
         self,
+        old_session: Optional[ClientSession] = None,
     ) -> Tuple[
         WebSocketCommonProtocol,
         ControlMessageHandshakeRequest,
@@ -92,19 +95,24 @@ 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 : {existing_session}"
+                    f"##### _establish_new_connection: existing session : {old_session}"
                 )
                 rate_limit.consume_budget(client_id)
                 handshake_request, handshake_response = await self._establish_handshake(
@@ -120,7 +128,7 @@ 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 _create_new_session(
@@ -151,36 +159,38 @@ class ClientTransport(Transport):
         return new_session
 
     async def _get_or_create_session(self) -> ClientSession:
-        async with self._session_lock:
+        logging.error(f"####### start get or create session")
+        async with self._create_session_lock:
             existing_session = await self._get_existing_session()
             if not existing_session:
-                logging.error(f"##### No 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"##### session open, creating new session")
-                await existing_session.close(
-                    is_unexpected_close=False, acquire_transport_lock=False
+                logging.error(
+                    f"##### _get_or_create_session session open, creating new session"
                 )
                 return await self._create_new_session()
-            is_ws_open = existing_session.is_websocket_open()
+            is_ws_open = await existing_session.is_websocket_open()
             if is_ws_open:
-                logging.error(f"##### Reuse existing session")
+                logging.error(f"##### _get_or_create_session Reuse existing session")
                 return existing_session
             else:
-                new_ws, _, hs_response = await self._establish_new_connection()
+                new_ws, _, hs_response = await self._establish_new_connection(
+                    existing_session
+                )
                 if (
                     hs_response.status.sessionId
                     == existing_session.advertised_session_id
                 ):
-                    logging.error(f"##### session open, replacing websocket")
+                    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, acquire_transport_lock=False
-                    )
+                    await existing_session.close(is_unexpected_close=False)
                     return await self._create_new_session()
 
     async def _send_handshake_request(

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

@@ -108,11 +108,9 @@ class Session(object):
         """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
-        await self._ws_wrapper.close()
         logging.debug(
             "websocket closed from %s to %s begin grace period",
             self._transport_id,
@@ -122,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."""
@@ -164,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)
@@ -235,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
@@ -282,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._begin_close_session_countdown()
-                    await self.close_websocket(
-                        self._ws_wrapper, should_retry=not self._is_server
-                    )
                     continue
             except FailedSendingMessageException:
                 # this is expected during websocket closed period
@@ -415,7 +425,8 @@ class Session(object):
                 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,

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

@@ -61,8 +61,7 @@ class Transport:
             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]
+            self._sessions[session._to_id] = session
 
     def generate_nanoid(self) -> str:
         return str(nanoid.generate())

{replit_river-0.1.16.dev3 → 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