Whoops

Author: blast-hardcheese

src/replit_river/v2/session.py

@@ -214,7 +214,6 @@ async def _do_ensure_connected[HandshakeMetadata](
 
         last_error: Exception | None = None
         i = 0
-        await self._connection_condition.acquire()
         while rate_limiter.has_budget_or_throw(client_id, ERROR_HANDSHAKE, last_error):
             if i > 0:
                 logger.info(f"Retrying build handshake number {i} times")
@@ -342,7 +341,10 @@ async def websocket_closed_callback() -> None:
                 rate_limiter.start_restoring_budget(client_id)
                 self._state = SessionState.ACTIVE
                 self._ws_unwrapped = ws
-                self._connection_condition.notify_all()
+
+                # We're connected, wake everybody up
+                async with self._connection_condition:
+                    self._connection_condition.notify_all()
                 break
             except Exception as e:
                 if ws:
@@ -370,10 +372,6 @@ async def websocket_closed_callback() -> None:
         ):
             self._connecting_task = None
 
-        # Release the lock we took earlier so we can use it again in the next
-        # connection attempt
-        self._connection_condition.release()
-
         if last_error is not None:
             raise RiverException(
                 ERROR_HANDSHAKE,
@@ -493,10 +491,8 @@ async def close(self) -> None:
         self._state = SessionState.CLOSING
 
         # We need to wake up all tasks waiting for connection to be established
-        if not self._connection_condition.locked():
-            await self._connection_condition.acquire()
-        self._connection_condition.notify_all()
-        self._connection_condition.release()
+        async with self._connection_condition:
+            self._connection_condition.notify_all()
 
         await self._task_manager.cancel_all_tasks()