Merge pull request #2424 from pipecat-ai/aleix/system-frame-queues-fix

FrameProcessor: fix race condition on FrameProcessorQueue

src/pipecat/processors/frame_processor.py

@@ -38,7 +38,10 @@ from pipecat.observers.base_observer import BaseObserver, FramePushed
 from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics
 from pipecat.utils.asyncio.task_manager import BaseTaskManager
 from pipecat.utils.asyncio.watchdog_event import WatchdogEvent
-from pipecat.utils.asyncio.watchdog_priority_queue import WatchdogPriorityQueue
+from pipecat.utils.asyncio.watchdog_priority_queue import (
+    WatchdogPriorityCancelSentinel,
+    WatchdogPriorityQueue,
+)
 from pipecat.utils.asyncio.watchdog_queue import WatchdogQueue
 from pipecat.utils.base_object import BaseObject
 
@@ -79,16 +82,8 @@ class FrameProcessorQueue(WatchdogPriorityQueue):
     """A priority queue for systems frames and other frames.
 
     This is a specialized queue for frame processors that separates and
-    prioritizes system frames over other frames.
-
-    This queue uses two internal `WatchdogQueue` instances:
-    - One for system-level frames (`SystemFrame`)
-    - One for regular frames
-
-    It ensures that `SystemFrame` objects are processed before any other
-    frames. Additionally, it uses an `asyncio.Event` to signal when new items
-    have been added to either queue, allowing consumers to wait efficiently when
-    the queue is empty.
+    prioritizes system frames over other frames. It ensures that `SystemFrame`
+    objects are processed before any other frames by using a priority queue.
 
     """
 
@@ -103,26 +98,33 @@ class FrameProcessorQueue(WatchdogPriorityQueue):
 
         """
         super().__init__(manager)
-        self.__counter = 0
-        self.__system_counter = 0
+        self.__high_counter = 0
+        self.__low_counter = 0
 
-    async def put(self, item: Tuple[Frame, FrameDirection, FrameCallback]):
-        """Put an item into the appropriate queue.
+    async def put(
+        self, item: WatchdogPriorityCancelSentinel | Tuple[Frame, FrameDirection, FrameCallback]
+    ):
+        """Put an item into the priority queue.
 
-        System frames (`SystemFrame`) are placed into the system queue and all others
-        into the regular queue. Signals the event to wake up any waiting consumers.
+        System frames (`SystemFrame`) have higher priority than any other
+        frames. If a non-frame item (e.g. a watchdog cancellation sentinel) is
+        provided it will have the highest priority.
 
         Args:
             item (Any): The item to enqueue.
 
         """
+        if isinstance(item, WatchdogPriorityCancelSentinel):
+            await super().put(item)
+            return
+
         frame, _, _ = item
         if isinstance(frame, SystemFrame):
-            await super().put((self.HIGH_PRIORITY, self.__system_counter, item))
-            self.__system_counter += 1
+            self.__high_counter += 1
+            await super().put((self.HIGH_PRIORITY, self.__high_counter, item))
         else:
-            await super().put((self.LOW_PRIORITY, self.__counter, item))
-            self.__counter += 1
+            self.__low_counter += 1
+            await super().put((self.LOW_PRIORITY, self.__low_counter, item))
 
     async def get(self) -> Any:
         """Retrieve the next item from the queue.

src/pipecat/utils/asyncio/watchdog_priority_queue.py

@@ -20,7 +20,17 @@ from pipecat.utils.asyncio.task_manager import BaseTaskManager
 
 
 @dataclass
-class _WatchdogPriorityCancelSentinel:
+class WatchdogPriorityCancelSentinel:
+    """Sentinel object used in priority queues to force cancellation.
+
+    An instance of this class is typically inserted into a
+    `WatchdogPriorityQueue` to act as a high-priority marker asyncio task
+    cancellation. The `__lt__` method always returns `True`, ensuring that the
+    sentinel is considered "less than" any other item in the queue, and
+    therefore processed before anything else.
+
+    """
+
     def __lt__(self, other):
         return True
 
@@ -62,7 +72,7 @@ class WatchdogPriorityQueue(asyncio.PriorityQueue):
         else:
             get_result = await super().get()
 
-        if isinstance(get_result, _WatchdogPriorityCancelSentinel):
+        if isinstance(get_result, WatchdogPriorityCancelSentinel):
             logger.trace(
                 "Received WatchdogPriorityCancelSentinel, throwing CancelledError to force cancelling"
             )
@@ -91,7 +101,7 @@ class WatchdogPriorityQueue(asyncio.PriorityQueue):
         forces the task to raise CancelledError when consumed, ensuring proper
         task termination.
         """
-        super().put_nowait(_WatchdogPriorityCancelSentinel())
+        super().put_nowait(WatchdogPriorityCancelSentinel())
 
     async def _watchdog_get(self):
         """Get item from queue while periodically resetting watchdog timer."""

src/pipecat/utils/asyncio/watchdog_queue.py

@@ -20,7 +20,14 @@ from pipecat.utils.asyncio.task_manager import BaseTaskManager
 
 
 @dataclass
-class _WatchdogQueueCancelSentinel:
+class WatchdogQueueCancelSentinel:
+    """Sentinel object used in queues to force cancellation.
+
+    An instance of this class is typically inserted into a `WatchdogQueue` to
+    act as a marker for asyncio task cancellation.
+
+    """
+
     pass
 
 
@@ -61,7 +68,7 @@ class WatchdogQueue(asyncio.Queue):
         else:
             get_result = await super().get()
 
-        if isinstance(get_result, _WatchdogQueueCancelSentinel):
+        if isinstance(get_result, WatchdogQueueCancelSentinel):
             logger.trace(
                 "Received WatchdogQueueCancelFrame, throwing CancelledError to force cancelling"
             )
@@ -90,7 +97,7 @@ class WatchdogQueue(asyncio.Queue):
         forces the task to raise CancelledError when consumed, ensuring proper
         task termination.
         """
-        super().put_nowait(_WatchdogQueueCancelSentinel())
+        super().put_nowait(WatchdogQueueCancelSentinel())
 
     async def _watchdog_get(self):
         """Get item from queue while periodically resetting watchdog timer."""

tests/test_watchdog_queue.py

@@ -0,0 +1,65 @@
+#
+# Copyright (c) 2024-2025 Daily
+#
+# SPDX-License-Identifier: BSD 2-Clause License
+#
+
+import asyncio
+import unittest
+
+from pipecat.utils.asyncio.task_manager import TaskManager
+from pipecat.utils.asyncio.watchdog_priority_queue import WatchdogPriorityQueue
+from pipecat.utils.asyncio.watchdog_queue import WatchdogQueue
+
+
+class TestWatchdogQueue(unittest.IsolatedAsyncioTestCase):
+    async def test_simple_item(self):
+        queue = WatchdogQueue(TaskManager())
+        await queue.put(1)
+        await queue.put(2)
+        await queue.put(3)
+        self.assertEqual(await queue.get(), 1)
+        queue.task_done()
+        self.assertEqual(await queue.get(), 2)
+        queue.task_done()
+        self.assertEqual(await queue.get(), 3)
+        queue.task_done()
+
+    async def test_watchdog_sentinel(self):
+        queue = WatchdogQueue(TaskManager())
+        await queue.put(1)
+        self.assertEqual(await queue.get(), 1)
+        queue.task_done()
+        # The get should throw an exception.
+        queue.cancel()
+        try:
+            await queue.get()
+            assert False
+        except asyncio.CancelledError:
+            assert True
+
+
+class TestWatchdogPriorityQueue(unittest.IsolatedAsyncioTestCase):
+    async def test_simple_item(self):
+        queue = WatchdogPriorityQueue(TaskManager())
+        await queue.put((3, 1))
+        await queue.put((2, 1))
+        await queue.put((1, 1))
+        self.assertEqual(await queue.get(), (1, 1))
+        queue.task_done()
+        self.assertEqual(await queue.get(), (2, 1))
+        queue.task_done()
+        self.assertEqual(await queue.get(), (3, 1))
+        queue.task_done()
+
+    async def test_watchdog_sentinel(self):
+        queue = WatchdogPriorityQueue(TaskManager())
+        await queue.put((0, 1))
+        # The get should throw an exception because the watchdog sentinel has
+        # higher priority.
+        queue.cancel()
+        try:
+            await queue.get()
+            assert False
+        except asyncio.CancelledError:
+            assert True