Break out has_capacity vs put to avoid race

Author: blast-hardcheese

src/replit_river/message_buffer.py

@@ -21,20 +21,22 @@ def __init__(self, max_num_messages: int = MAX_MESSAGE_BUFFER_SIZE):
         self._space_available_cond = asyncio.Condition()
         self._closed = False
 
-    async def put(self, message: TransportMessage) -> None:
+    async def has_capacity(self) -> None:
+        async with self._space_available_cond:
+            await self._space_available_cond.wait_for(
+                lambda: len(self.buffer) < self.max_size or self._closed
+            )
+
+    def put(self, message: TransportMessage) -> None:
         """Add a message to the buffer. Blocks until there is space in the buffer.
 
         Raises:
             MessageBufferClosedError: if the buffer is closed.
         """
-        async with self._space_available_cond:
-            await self._space_available_cond.wait_for(
-                lambda: len(self.buffer) < self.max_size or self._closed
-            )
-            if self._closed:
-                raise MessageBufferClosedError("message buffer is closed")
-            self.buffer.append(message)
-            self._has_messages.set()
+        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."""

src/replit_river/session.py

@@ -255,8 +255,9 @@ async def send_message(
             # We need this lock to ensure the buffer order and message sending order
             # are the same.
             async with self._msg_lock:
+                await self._buffer.has_capacity()
                 try:
-                    await self._buffer.put(msg)
+                    self._buffer.put(msg)
                 except MessageBufferClosedError:
                     # The session is closed and is no longer accepting new messages.
                     return

tests/test_message_buffer.py

@@ -35,7 +35,8 @@ async def test_message_buffer_backpressure() -> None:
 
     async def put_messages() -> None:
         for i in range(0, iterations):
-            await buffer.put(mock_transport_message(seq=i))
+            await buffer.has_capacity()
+            buffer.put(mock_transport_message(seq=i))
             await sync_events.put(None)
 
     background_puts = asyncio.create_task(put_messages())
@@ -55,10 +56,17 @@ async def test_message_buffer_close() -> None:
     is closed while the put operation is waiting for space in the buffer.
     """
     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)))
+    await buffer.has_capacity()
+    buffer.put(mock_transport_message(seq=1))
+
+    async def bg_put(msg: TransportMessage) -> None:
+        await buffer.has_capacity()
+        buffer.put(msg)
+
+    background_put = asyncio.create_task(bg_put(mock_transport_message(seq=1)))
     await buffer.close()
     with pytest.raises(MessageBufferClosedError):
         await background_put
     with pytest.raises(MessageBufferClosedError):
-        await buffer.put(mock_transport_message(seq=1))
+        await buffer.has_capacity()
+        buffer.put(mock_transport_message(seq=1))