Permit sensible message_buffer state transitions

Author: blast-hardcheese

src/replit_river/client_session.py

@@ -140,7 +140,7 @@ async def _handle_messages_from_ws(self) -> None:
                         case other:
                             assert_never(other)
 
-                    self._buffer.remove_old_messages(
+                    await self._buffer.remove_old_messages(
                         self._seq_manager.receiver_ack,
                     )
                     self._reset_session_close_countdown()

src/replit_river/message_buffer.py

@@ -17,6 +17,7 @@ class MessageBuffer:
     def __init__(self, max_num_messages: int = MAX_MESSAGE_BUFFER_SIZE):
         self.max_size = max_num_messages
         self.buffer: list[TransportMessage] = []
+        self._has_messages = asyncio.Event()
         self._space_available_cond = asyncio.Condition()
         self._closed = False
 
@@ -33,23 +34,33 @@ async def put(self, message: TransportMessage) -> None:
             if self._closed:
                 raise MessageBufferClosedError("message buffer is closed")
             self.buffer.append(message)
+            self._has_messages.set()
 
     def peek(self) -> TransportMessage | None:
         """Peek the first message in the buffer, returns None if the buffer is empty."""
         if len(self.buffer) == 0:
             return None
         return self.buffer[0]
 
-    def remove_old_messages(self, min_seq: int) -> None:
+    async def remove_old_messages(self, min_seq: int) -> None:
         """Remove messages in the buffer with a seq number less than min_seq."""
         self.buffer = [msg for msg in self.buffer if msg.seq >= min_seq]
+        if self.buffer:
+            self._has_messages.set()
+        else:
+            self._has_messages.clear()
         async with self._space_available_cond:
             self._space_available_cond.notify_all()
 
-    def close(self) -> None:
+    async def block_until_message_available(self) -> None:
+        """Allow consumers to avoid spinning unnecessarily"""
+        await self._has_messages.wait()
+
+    async def close(self) -> None:
         """
         Closes the message buffer and rejects any pending put operations.
         """
         self._closed = True
+        self._has_messages.set()
         async with self._space_available_cond:
             self._space_available_cond.notify_all()

src/replit_river/server_session.py

@@ -136,7 +136,7 @@ async def _handle_messages_from_ws(self, tg: asyncio.TaskGroup) -> None:
                             pass
                         case other:
                             assert_never(other)
-                    self._buffer.remove_old_messages(
+                    await self._buffer.remove_old_messages(
                         self._seq_manager.receiver_ack,
                     )
                     self._reset_session_close_countdown()

src/replit_river/session.py

@@ -325,7 +325,7 @@ async def close(self) -> None:
 
             await self.close_websocket(self._ws_wrapper, should_retry=False)
 
-            self._buffer.close()
+            await self._buffer.close()
 
             # Clear the session in transports
             await self._close_session_callback(self)

tests/test_message_buffer.py

@@ -44,7 +44,7 @@ async def put_messages() -> None:
         # Wait for the put call to return.
         await sync_events.get()
         assert len(buffer.buffer) == 1
-        buffer.remove_old_messages(i)
+        await buffer.remove_old_messages(i)
 
     await background_puts
 
@@ -57,7 +57,7 @@ async def test_message_buffer_close() -> None:
     buffer = MessageBuffer(max_num_messages=1)
     await buffer.put(mock_transport_message(seq=1))
     background_put = asyncio.create_task(buffer.put(mock_transport_message(seq=1)))
-    buffer.close()
+    await buffer.close()
     with pytest.raises(MessageBufferClosedError):
         await background_put
     with pytest.raises(MessageBufferClosedError):