Block ensure_connected until connected

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -41,10 +41,11 @@ class SessionState(enum.Enum):
     - CLOSING -> {CLOSED}
     """
 
-    CONNECTING = 0
-    ACTIVE = 1
-    CLOSING = 2
-    CLOSED = 3
+    PENDING = 0
+    CONNECTING = 1
+    ACTIVE = 2
+    CLOSING = 3
+    CLOSED = 4
 
 
 TerminalStates = set([SessionState.CLOSING, SessionState.CLOSED])

src/replit_river/v2/session.py

@@ -102,6 +102,7 @@ class Session:
 
     # ws state
     _ws_unwrapped: ClientConnection | None
+    _ensure_connected_condition: asyncio.Condition
     _heartbeat_misses: int
     _retry_connection_callback: RetryConnectionCallback | None
 
@@ -129,13 +130,14 @@ def __init__(
         self.session_id = session_id
         self._transport_options = transport_options
 
-        # session state, only modified during closing
-        self._state = SessionState.CONNECTING
+        # session state
+        self._state = SessionState.PENDING
         self._close_session_callback = close_session_callback
         self._close_session_after_time_secs: float | None = None
 
         # ws state
         self._ws_unwrapped = None
+        self._ensure_connected_condition = asyncio.Condition()
         self._heartbeat_misses = 0
         self._retry_connection_callback = retry_connection_callback
 
@@ -236,11 +238,22 @@ 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._state == SessionState.ACTIVE:
-            return
         max_retry = self._transport_options.connection_retry_options.max_retry
         logger.info("Attempting to establish new ws connection")
 
+        if self.is_connected():
+            return
+
+        while True:
+            await self._ensure_connected_condition.acquire()
+            if self._state == SessionState.ACTIVE:
+                return
+            elif self._state == SessionState.PENDING:
+                self._state = SessionState.CONNECTING
+                break
+            elif self._state in TerminalStates:
+                raise RiverException("SESSION_CLOSING", "Going away")
+
         last_error: Exception | None = None
         i = 0
         while rate_limiter.has_budget_or_throw(client_id, ERROR_HANDSHAKE, last_error):
@@ -349,6 +362,7 @@ async def websocket_closed_callback() -> None:
 
                     rate_limiter.start_restoring_budget(client_id)
                     self._state = SessionState.ACTIVE
+                    self._ensure_connected_condition.notify_all()
                 except RiverException as e:
                     await ws.close()
                     raise e