Just use _state instead of having two

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -35,20 +35,26 @@ async def __call__(
 class SessionState(enum.Enum):
     """The state a session can be in.
 
-    Can only transition from ACTIVE to CLOSING to CLOSED.
+    Valid transitions:
+    - CONNECTING -> {ACTIVE, CLOSING}
+    - ACTIVE -> {CONNECTING, CLOSING}
+    - CLOSING -> {CLOSED}
     """
 
-    ACTIVE = 0
-    CLOSING = 1
-    CLOSED = 2
+    CONNECTING = 0
+    ACTIVE = 1
+    CLOSING = 2
+    CLOSED = 3
+
+
+TerminalStates = set([SessionState.CLOSING, SessionState.CLOSED])
 
 
 async def setup_heartbeat(
     session_id: str,
     heartbeat_ms: float,
     heartbeats_until_dead: int,
     get_state: Callable[[], SessionState],
-    get_connected: Callable[[], bool],
     get_closing_grace_period: Callable[[], float | None],
     close_websocket: Callable[[], Awaitable[None]],
     send_message: SendMessage,
@@ -58,10 +64,10 @@ async def setup_heartbeat(
     while True:
         await asyncio.sleep(heartbeat_ms / 1000)
         state = get_state()
-        if not get_connected():
+        if state == SessionState.CONNECTING:
             logger.debug("Websocket is not connected, not sending heartbeat")
             continue
-        if state != SessionState.ACTIVE:
+        if state in TerminalStates:
             logger.debug(
                 "Session is closed, no need to send heartbeat, state : "
                 "%r close_session_after_this: %r",
@@ -110,7 +116,7 @@ async def check_to_close_session(
 ) -> None:
     while True:
         await asyncio.sleep(close_session_check_interval_ms / 1000)
-        if get_state() != SessionState.ACTIVE:
+        if get_state() in TerminalStates:
             # already closing
             return
         # calculate the value now before comparing it so that there are no

src/replit_river/session.py

@@ -51,7 +51,7 @@ class Session:
     session_id: str
     _transport_options: TransportOptions
 
-    # session state, only modified during closing
+    # session state
     _state: SessionState
     _state_lock: asyncio.Lock
     _close_session_callback: CloseSessionCallback
@@ -88,7 +88,7 @@ def __init__(
         self.session_id = session_id
         self._transport_options = transport_options
 
-        # session state, only modified during closing
+        # session state
         self._state = SessionState.ACTIVE
         self._state_lock = asyncio.Lock()
         self._close_session_callback = close_session_callback
@@ -123,8 +123,11 @@ def increment_and_get_heartbeat_misses() -> int:
                 self.session_id,
                 self._transport_options.heartbeat_ms,
                 self._transport_options.heartbeats_until_dead,
-                lambda: self._state,
-                lambda: self._ws_wrapper.ws_state == WsState.OPEN,
+                lambda: (
+                    self._state
+                    if self._ws_wrapper.ws_state == WsState.OPEN
+                    else SessionState.CONNECTING
+                ),
                 lambda: self._close_session_after_time_secs,
                 close_websocket=do_close_websocket,
                 send_message=self.send_message,

src/replit_river/v2/client_transport.py

@@ -66,7 +66,7 @@ async def get_or_create_session(self) -> Session:
         call ensure_connected on whatever session is active.
         """
         existing_session = self._session
-        if not existing_session or not existing_session.is_session_open():
+        if not existing_session or existing_session.is_closed():
             logger.info("Creating new session")
             new_session = Session(
                 transport_id=self._transport_id,

src/replit_river/v2/session.py

@@ -28,6 +28,7 @@
 
 from replit_river.common_session import (
     SessionState,
+    TerminalStates,
     buffered_message_sender,
     check_to_close_session,
     setup_heartbeat,
@@ -103,7 +104,6 @@ class Session:
     _close_session_after_time_secs: float | None
 
     # ws state
-    _ws_connected: bool
     _ws_unwrapped: ClientConnection | None
     _heartbeat_misses: int
     _retry_connection_callback: RetryConnectionCallback | None
@@ -133,12 +133,11 @@ def __init__(
         self._transport_options = transport_options
 
         # session state, only modified during closing
-        self._state = SessionState.ACTIVE
+        self._state = SessionState.CONNECTING
         self._close_session_callback = close_session_callback
         self._close_session_after_time_secs: float | None = None
 
         # ws state
-        self._ws_connected = False
         self._ws_unwrapped = None
         self._heartbeat_misses = 0
         self._retry_connection_callback = retry_connection_callback
@@ -160,18 +159,43 @@ def __init__(
 
         async def do_close_websocket() -> None:
             logger.debug(
-                "do_close called, _ws_connected=%r, _ws_unwrapped=%r",
-                self._ws_connected,
+                "do_close called, _state=%r, _ws_unwrapped=%r",
+                self._state,
                 self._ws_unwrapped,
             )
-            self._ws_connected = False
+            self._state = SessionState.CLOSING
             if self._ws_unwrapped:
                 self._task_manager.create_task(self._ws_unwrapped.close())
                 if self._retry_connection_callback:
                     self._task_manager.create_task(self._retry_connection_callback())
             await self._begin_close_session_countdown()
 
-        self._setup_heartbeats_task(do_close_websocket)
+        def increment_and_get_heartbeat_misses() -> int:
+            self._heartbeat_misses += 1
+            return self._heartbeat_misses
+
+        self._task_manager.create_task(
+            setup_heartbeat(
+                self.session_id,
+                self._transport_options.heartbeat_ms,
+                self._transport_options.heartbeats_until_dead,
+                lambda: self._state,
+                lambda: self._close_session_after_time_secs,
+                close_websocket=do_close_websocket,
+                send_message=self.send_message,
+                increment_and_get_heartbeat_misses=increment_and_get_heartbeat_misses,
+            )
+        )
+        self._task_manager.create_task(
+            check_to_close_session(
+                self._transport_id,
+                self._transport_options.close_session_check_interval_ms,
+                lambda: self._state,
+                self._get_current_time,
+                lambda: self._close_session_after_time_secs,
+                self.close,
+            )
+        )
 
         def commit(msg: TransportMessage) -> None:
             pending = self._send_buffer.popleft()
@@ -193,7 +217,7 @@ def get_next_pending() -> TransportMessage | None:
                 self._message_enqueued,
                 get_ws=lambda: (
                     cast(WebSocketCommonProtocol | ClientConnection, self._ws_unwrapped)
-                    if self.is_websocket_open()
+                    if self.is_connected()
                     else None
                 ),
                 websocket_closed_callback=self._begin_close_session_countdown,
@@ -214,7 +238,7 @@ async def ensure_connected[HandshakeMetadata](
         Either return immediately or establish a websocket connection and return
         once we can accept messages
         """
-        if self._ws_unwrapped and self._ws_connected:
+        if self._ws_unwrapped and self._state == SessionState.ACTIVE:
             return
         max_retry = self._transport_options.connection_retry_options.max_retry
         logger.info("Attempting to establish new ws connection")
@@ -326,6 +350,7 @@ async def websocket_closed_callback() -> None:
                         )
 
                     rate_limiter.start_restoring_budget(client_id)
+                    self._state = SessionState.ACTIVE
                 except RiverException as e:
                     await ws.close()
                     raise e
@@ -342,44 +367,17 @@ async def websocket_closed_callback() -> None:
             f"Failed to create ws after retrying {max_retry} number of times",
         ) from last_error
 
-    def _setup_heartbeats_task(
-        self,
-        do_close_websocket: Callable[[], Awaitable[None]],
-    ) -> None:
-        def increment_and_get_heartbeat_misses() -> int:
-            self._heartbeat_misses += 1
-            return self._heartbeat_misses
-
-        self._task_manager.create_task(
-            setup_heartbeat(
-                self.session_id,
-                self._transport_options.heartbeat_ms,
-                self._transport_options.heartbeats_until_dead,
-                lambda: self._state,
-                lambda: self._ws_connected,
-                lambda: self._close_session_after_time_secs,
-                close_websocket=do_close_websocket,
-                send_message=self.send_message,
-                increment_and_get_heartbeat_misses=increment_and_get_heartbeat_misses,
-            )
-        )
-        self._task_manager.create_task(
-            check_to_close_session(
-                self._transport_id,
-                self._transport_options.close_session_check_interval_ms,
-                lambda: self._state,
-                self._get_current_time,
-                lambda: self._close_session_after_time_secs,
-                self.close,
-            )
-        )
+    def is_closed(self) -> bool:
+        """
+        If the session is in a terminal state.
+        Do not send messages, do not expect any more messages to be emitted,
+        the state is expected to be stale.
+        """
+        return self._state not in TerminalStates
 
-    def is_session_open(self) -> bool:
+    def is_connected(self) -> bool:
         return self._state == SessionState.ACTIVE
 
-    def is_websocket_open(self) -> bool:
-        return self._ws_connected
-
     async def _begin_close_session_countdown(self) -> None:
         """Begin the countdown to close session, this should be called when
         websocket is closed.
@@ -400,17 +398,6 @@ async def _begin_close_session_countdown(self) -> None:
             self._to_id,
         )
         self._close_session_after_time_secs = close_session_after_time_secs
-        self._ws_connected = False
-
-    async def replace_with_new_websocket(self, new_ws: ClientConnection) -> None:
-        if self._ws_unwrapped and new_ws.id != self._ws_unwrapped.id:
-            self._task_manager.create_task(
-                self._ws_unwrapped.close(
-                    CloseCode.PROTOCOL_ERROR, "Transparent reconnect"
-                )
-            )
-        self._ws_unwrapped = new_ws
-        self._ws_connected = True
 
     async def _get_current_time(self) -> float:
         return asyncio.get_event_loop().time()
@@ -430,7 +417,7 @@ async def send_message(
     ) -> None:
         """Send serialized messages to the websockets."""
         # if the session is not active, we should not do anything
-        if self._state != SessionState.ACTIVE:
+        if self._state in TerminalStates:
             return
         msg = TransportMessage(
             streamId=stream_id,
@@ -476,7 +463,7 @@ async def close(self) -> None:
             f"{self._transport_id} closing session "
             f"to {self._to_id}, ws: {self._ws_unwrapped}"
         )
-        if self._state != SessionState.ACTIVE:
+        if self._state in TerminalStates:
             # already closing
             return
         self._state = SessionState.CLOSING
@@ -510,6 +497,8 @@ async def _serve(self) -> None:
             try:
                 await self._handle_messages_from_ws()
             except ConnectionClosed:
+                # Set ourselves to closed as soon as we get the signal
+                self._state = SessionState.CONNECTING
                 if self._retry_connection_callback:
                     self._task_manager.create_task(self._retry_connection_callback())
 
@@ -530,7 +519,7 @@ async def _serve(self) -> None:
                 )
 
     async def _handle_messages_from_ws(self) -> None:
-        while self._ws_unwrapped is None or not self._ws_connected:
+        while self._ws_unwrapped is None or self._state == SessionState.CONNECTING:
             await asyncio.sleep(1)
         logger.debug(
             "%s start handling messages from ws %s",
@@ -628,8 +617,10 @@ async def _handle_messages_from_ws(self) -> None:
                     await self.close()
                     return
         except ConnectionClosedOK:
-            pass  # Exited normally
+            # Exited normally
+            self._state = SessionState.CONNECTING
         except ConnectionClosed as e:
+            self._state = SessionState.CONNECTING
             raise e
 
     async def send_rpc[R, A](