Flow controller perf and close fix

pyproject.toml

@@ -28,5 +28,11 @@ build-backend = "uv_build"
 dev = [
     "dotenv>=0.9.9",
     "pipecat-ai[google,silero]>=1.1.0",
+    "pytest>=8.0",
+    "pytest-asyncio>=0.24",
     "uvicorn>=0.46.0",
 ]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]

src/pipecat_asterisk/transport/flow_controller.py

@@ -5,8 +5,10 @@
 #
 
 import asyncio
-from loguru import logger
 import time
+from collections import deque
+
+from loguru import logger
 from pipecat.transports.websocket.fastapi import (
     FastAPIWebsocketClient,
 )
@@ -48,7 +50,11 @@ def __init__(
         self._ptime = ptime  # Audio chunk duration. In milliseconds
         self._psize = psize  # Audio chunk size. In bytes
         self._websocket_client = websocket_client
-        self._local_buffer = bytearray()
+        # Deque of bytes chunks (head == oldest). Total queued bytes is tracked
+        # separately so the working-range check stays O(1) without walking the
+        # deque. See `_pop_bytes` for how sends drain it.
+        self._local_buffer: deque[bytes] = deque()
+        self._local_buffer_size: int = 0
         self._remote_buffer_low_water = (
             self.REMOTE_BUFFER_LOW_WATER * self.REMOTE_BUFFER_SIZE * self._psize
         )
@@ -70,9 +76,12 @@ def __call__(self, chunk: bytes) -> None:
             chunk: The audio chunk to add to the local buffer.
 
         """
-        self._local_buffer.extend(chunk)
+        if not chunk:
+            return
+        self._local_buffer.append(chunk)
+        self._local_buffer_size += len(chunk)
         logger.trace(
-            f"Buffered {len(chunk)} bytes to local buffer. Local buffer size: {len(self._local_buffer)} bytes."
+            f"Buffered {len(chunk)} bytes to local buffer. Local buffer size: {self._local_buffer_size} bytes."
         )
 
     async def flow_control(self):
@@ -100,7 +109,7 @@ async def flow_control(self):
 
             # Flow control logic
             # First check if we have something in the local buffer
-            if len(self._local_buffer) > 0:
+            if self._local_buffer_size > 0:
                 # If the remote buffer is under the low water mark we send whatever we have in the local buffer
                 if self._remote_buffer_utilization < self._remote_buffer_low_water:
                     await self.send_chunks()
@@ -110,10 +119,40 @@ async def flow_control(self):
                 # and we have at least twice as much free space in the remote buffer as the minimum batch size to avoid overfilling the remote buffer and causing audio dropouts on the Asterisk side.
                 elif (
                     self._remote_buffer_utilization < self._remote_buffer_high_water - self._min_batch * 2
-                ) and (len(self._local_buffer) >= self._min_batch):
+                ) and (self._local_buffer_size >= self._min_batch):
                     await self.send_chunks()
                 # If the remote buffer is above the high water mark we don't send anything and wait for the next tick to see if the remote buffer utilization has decreased enough to send more audio
 
+    def _pop_bytes(self, max_bytes: int) -> bytes:
+        """Pop up to ``max_bytes`` from the head of the local buffer.
+
+        Walks the chunk deque, popping whole chunks while they fit and slicing
+        the head chunk if the remaining budget falls below its size. The
+        returned bytes object is allocated once via ``b"".join`` regardless of
+        how many deque entries we drain. ``_local_buffer_size`` is updated in
+        lockstep so the working-range check in ``flow_control`` stays O(1).
+        """
+        if max_bytes <= 0 or self._local_buffer_size == 0:
+            return b""
+
+        parts: list[bytes] = []
+        remaining = max_bytes
+        while remaining > 0 and self._local_buffer:
+            head = self._local_buffer[0]
+            if len(head) <= remaining:
+                parts.append(self._local_buffer.popleft())
+                remaining -= len(head)
+            else:
+                parts.append(head[:remaining])
+                self._local_buffer[0] = head[remaining:]
+                remaining = 0
+
+        # `b"".join` allocates once; in the single-part case skip the join to
+        # avoid the trivial copy and reuse the existing bytes object.
+        chunk = parts[0] if len(parts) == 1 else b"".join(parts)
+        self._local_buffer_size -= len(chunk)
+        return chunk
+
     async def send_chunks(self):
         """Send audio chunks from the local buffer to websocket (effectively to the remote buffer on the Asterisk side).
 
@@ -124,42 +163,48 @@ async def send_chunks(self):
 
         # Calculate the number of bytes to send
         bytes_to_send = min(
-            len(self._local_buffer), self.MAX_WS_SEND
+            self._local_buffer_size, self.MAX_WS_SEND
         )  # Ensure we don't exceed the websocket maximum message size
         if bytes_to_send > 0:
-            # Take the bytes to send from the local buffer
-            chunk = bytes(self._local_buffer[:bytes_to_send])
-            del self._local_buffer[:bytes_to_send]
+            # Take the bytes to send from the head of the local buffer
+            chunk = self._pop_bytes(bytes_to_send)
 
             # Send the chunk to the websocket
             await self._websocket_client.send(chunk)
             # Update the remote buffer utilization
             self._remote_buffer_utilization += len(chunk)
-            logger.debug(
+            logger.trace(
                 f"Sent {len(chunk)} bytes to websocket. Remote buffer utilization: {self._remote_buffer_utilization:.0f} bytes, {self._remote_buffer_utilization / (self._psize * self.REMOTE_BUFFER_SIZE) * 100:.1f}%."
             )
 
-    def close(self, gracefully: bool = False):
-        """Cancel the flow control task and optionally wait for the local buffer to be sent before cancelling.
+    async def close(self, gracefully: bool = False) -> None:
+        """Cancel the flow control task, optionally draining the local buffer first.
 
         Args:
-            gracefully: If True, wait for the local buffer to be sent before cancelling the flow control
+            gracefully: If True, wait until the local buffer is empty (so the
+                ``flow_control`` task has had a chance to send everything),
+                then cancel. If False, cancel immediately and drop any pending
+                audio in the local buffer.
         """
-        if self._flow_control:
-            if gracefully:
-                logger.info(
-                    f"Gracefully closing flow controller. Waiting for local buffer to be sent..."
-                )
-                while len(self._local_buffer) > 0:
-                    time.sleep(
-                        self._ptime / 1000
-                    )  # Sleep for the duration of one audio chunk to give the flow control loop time to send the remaining audio in the local buffer
-            self._flow_control.cancel()
+        if self._flow_control is None:
+            return
+        if gracefully:
+            logger.info(
+                "Gracefully closing flow controller. Waiting for local buffer to be sent..."
+            )
+            # Sleep one chunk-duration at a time on the event loop, allowing
+            # the flow_control task to keep draining `_local_buffer`. Using
+            # `time.sleep` here would block the event loop and prevent the
+            # very draining we're waiting on.
+            while self._local_buffer_size > 0:
+                await asyncio.sleep(self._ptime / 1000)
+        self._flow_control.cancel()
 
     def drop_buffer(self):
         """Drop any buffered audio in the local buffer and reset remote buffer utilization to zero.
 
         This is used when an interruption/stop/cancel frame is processed to avoid replaying stale audio.
         """
         self._local_buffer.clear()
+        self._local_buffer_size = 0
         self._remote_buffer_utilization = 0.0

tests/test_flow_controller.py

@@ -0,0 +1,194 @@
+"""Unit tests for FlowController.
+
+Covers the deque-buffer drain semantics, drop_buffer state reset,
+and that the async close(gracefully=True) yields the event loop so the
+flow_control task can actually drain.
+"""
+
+import asyncio
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from pipecat_asterisk.transport.flow_controller import FlowController
+
+
+PTIME_MS = 20
+PSIZE_BYTES = 640  # slin16 @ 20ms == 16000 Hz * 2 bytes/sample * 0.02 s
+
+
+def _make_controller() -> FlowController:
+    """Build a controller with a mock websocket client.
+
+    The controller starts its `flow_control` task in `__init__`, so this
+    must be called from inside a running event loop.
+    """
+    client = MagicMock()
+    client.send = AsyncMock()
+    return FlowController(ptime=PTIME_MS, psize=PSIZE_BYTES, websocket_client=client)
+
+
+# ---------------------------------------------------------------------------
+# __call__ / buffer accounting
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_call_appends_chunk_and_tracks_size():
+    fc = _make_controller()
+    try:
+        fc(b"\x00" * 640)
+        fc(b"\x01" * 320)
+        assert fc._local_buffer_size == 960
+        assert list(fc._local_buffer) == [b"\x00" * 640, b"\x01" * 320]
+    finally:
+        await fc.close()
+
+
+@pytest.mark.asyncio
+async def test_call_ignores_empty_chunk():
+    fc = _make_controller()
+    try:
+        fc(b"")
+        assert fc._local_buffer_size == 0
+        assert len(fc._local_buffer) == 0
+    finally:
+        await fc.close()
+
+
+# ---------------------------------------------------------------------------
+# _pop_bytes drain semantics
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_pop_bytes_drains_whole_chunks():
+    """Popping a budget that lines up on chunk boundaries pops chunks whole."""
+    fc = _make_controller()
+    try:
+        fc(b"a" * 100)
+        fc(b"b" * 100)
+        fc(b"c" * 100)
+        out = fc._pop_bytes(200)
+        assert out == b"a" * 100 + b"b" * 100
+        assert fc._local_buffer_size == 100
+        assert list(fc._local_buffer) == [b"c" * 100]
+    finally:
+        await fc.close()
+
+
+@pytest.mark.asyncio
+async def test_pop_bytes_slices_head_chunk():
+    """A budget smaller than the head chunk slices the head and leaves the tail."""
+    fc = _make_controller()
+    try:
+        fc(b"abcdef")
+        out = fc._pop_bytes(3)
+        assert out == b"abc"
+        assert fc._local_buffer_size == 3
+        assert list(fc._local_buffer) == [b"def"]
+    finally:
+        await fc.close()
+
+
+@pytest.mark.asyncio
+async def test_pop_bytes_handles_mix_of_whole_and_partial():
+    """A budget that straddles a chunk boundary pops one whole chunk + slices the next."""
+    fc = _make_controller()
+    try:
+        fc(b"a" * 100)
+        fc(b"b" * 100)
+        out = fc._pop_bytes(150)
+        assert out == b"a" * 100 + b"b" * 50
+        assert fc._local_buffer_size == 50
+        assert list(fc._local_buffer) == [b"b" * 50]
+    finally:
+        await fc.close()
+
+
+@pytest.mark.asyncio
+async def test_pop_bytes_caps_at_buffer_size():
+    """Asking for more bytes than are buffered returns everything available."""
+    fc = _make_controller()
+    try:
+        fc(b"hello")
+        out = fc._pop_bytes(1_000_000)
+        assert out == b"hello"
+        assert fc._local_buffer_size == 0
+        assert len(fc._local_buffer) == 0
+    finally:
+        await fc.close()
+
+
+@pytest.mark.asyncio
+async def test_pop_bytes_returns_empty_when_empty():
+    fc = _make_controller()
+    try:
+        assert fc._pop_bytes(100) == b""
+        assert fc._pop_bytes(0) == b""
+    finally:
+        await fc.close()
+
+
+# ---------------------------------------------------------------------------
+# drop_buffer
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_drop_buffer_resets_deque_size_and_utilization():
+    fc = _make_controller()
+    try:
+        fc(b"a" * 100)
+        fc(b"b" * 100)
+        fc._remote_buffer_utilization = 1234.5
+        fc.drop_buffer()
+        assert fc._local_buffer_size == 0
+        assert len(fc._local_buffer) == 0
+        assert fc._remote_buffer_utilization == 0.0
+    finally:
+        await fc.close()
+
+
+# ---------------------------------------------------------------------------
+# close()
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_close_non_graceful_cancels_immediately():
+    fc = _make_controller()
+    task = fc._flow_control
+    fc(b"x" * 1024)
+    await fc.close(gracefully=False)
+    with pytest.raises(asyncio.CancelledError):
+        await task
+    # Non-graceful close leaves the unsent buffer in place by design.
+    assert fc._local_buffer_size == 1024
+
+
+@pytest.mark.asyncio
+async def test_close_graceful_yields_loop_and_drains():
+    """The bug fix: close(gracefully=True) must `await asyncio.sleep`, not
+    `time.sleep`, so the flow_control task can actually run and drain the
+    buffer. We prove this by checking that close returns once
+    _local_buffer_size hits zero - which can only happen if flow_control
+    was given CPU during the wait.
+    """
+    fc = _make_controller()
+    try:
+        # Buffer one chunk worth of audio. Low water mark logic in
+        # flow_control will dispatch it on the next tick.
+        fc(b"\x00" * PSIZE_BYTES)
+        assert fc._local_buffer_size == PSIZE_BYTES
+
+        # If close used time.sleep, this would hang the event loop until
+        # the cancellation, and flow_control would never get a chance to
+        # send. Cap the wait so a regression times out instead of hanging.
+        await asyncio.wait_for(fc.close(gracefully=True), timeout=1.0)
+
+        assert fc._local_buffer_size == 0
+        fc._websocket_client.send.assert_awaited()
+    finally:
+        # Already cancelled above; second close() is a no-op safety net.
+        await fc.close()