This is what Semaphores are for

Author: blast-hardcheese

src/replit_river/common_session.py

@@ -127,16 +127,17 @@ async def check_to_close_session(
 
 
 async def buffered_message_sender(
+    message_enqueued: asyncio.Semaphore,
     get_ws: Callable[[], WebSocketCommonProtocol | ClientConnection | None],
     websocket_closed_callback: Callable[[], Coroutine[Any, Any, None]],
     get_next_pending: Callable[[], TransportMessage | None],
     commit: Callable[[TransportMessage], None],
 ) -> None:
     logger.debug("Entering buffered_message_sender")
     while True:
-        logger.debug("buffered_message_sender: LOOP")
-        sent = False
-        while msg := get_next_pending():
+        await message_enqueued.acquire()
+        logger.debug("buffered_message_sender: acquired")
+        if msg := get_next_pending():
             ws = get_ws()
             logger.debug(
                 "buffered_message_sender: Dequeued %r to send over %r",
@@ -150,7 +151,6 @@ async def buffered_message_sender(
                 await send_transport_message(msg, ws, websocket_closed_callback)
                 logger.debug("buffered_message_sender: committing!")
                 commit(msg)
-                sent = True
             except WebsocketClosedException as e:
                 logger.debug(
                     "Connection closed while sending message %r, waiting for "
@@ -168,8 +168,6 @@ async def buffered_message_sender(
             except Exception:
                 logger.exception("Error attempting to send buffered messages")
                 break
-        if not sent:
-            await asyncio.sleep(0.25)
 
 
 async def add_msg_to_stream(

src/replit_river/v2/session.py

@@ -129,6 +129,7 @@ def __init__(
         self._retry_connection_callback = retry_connection_callback
 
         # message state
+        self._message_enqueued = asyncio.Semaphore()
         self._space_available_cond = asyncio.Condition()
         self._queue_full_lock = asyncio.Lock()
 
@@ -174,6 +175,7 @@ def get_next_pending() -> TransportMessage | None:
 
         self._task_manager.create_task(
             buffered_message_sender(
+                self._message_enqueued,
                 get_ws=lambda: (
                     cast(WebSocketCommonProtocol | ClientConnection, self._ws_unwrapped)
                     if self.is_websocket_open()
@@ -309,6 +311,8 @@ async def send_message(
             await self._queue_full_lock.acquire()
             logger.warning("LOCK RELEASED %r", repr(payload))
         self._send_buffer.append(msg)
+        # Wake up buffered_message_sender
+        self._message_enqueued.release()
         self.seq += 1
 
     async def close(self) -> None: