fix(rtc): finish AudioMixer iteration when closed

livekit-rtc/livekit/rtc/audio_mixer.py

@@ -67,6 +67,7 @@ def __init__(
         # _ending signals that no new streams will be added,
         # but we continue processing until all streams are exhausted.
         self._ending: bool = False
+        self._closed: bool = False
         self._mixer_task: asyncio.Task = asyncio.create_task(self._mixer())
 
     def add_stream(self, stream: AsyncIterator[AudioFrame]) -> None:
@@ -106,29 +107,40 @@ def __aiter__(self) -> "AudioMixer":
         return self
 
     async def __anext__(self) -> AudioFrame:
+        if self._closed:
+            raise StopAsyncIteration
         item = await self._queue.get()
         if item is None:
+            self._closed = True
+            self._queue.put_nowait(None)
             raise StopAsyncIteration
         return item
 
     async def aclose(self) -> None:
         """
         Immediately stop mixing and close the mixer.
 
-        This cancels the mixing task, and any unconsumed output in the queue may be dropped.
+        This cancels the mixing task and drops any unconsumed output in the queue.
+        An iterator waiting for output is released, and subsequent iteration stops.
         Pending stream reads are cancelled and awaited, so a stream must let
         asyncio.CancelledError propagate once its cleanup is done.
         """
         self._ending = True
-        self._mixer_task.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            await self._mixer_task
-        pending = list(self._pending_reads.values())
-        self._pending_reads.clear()
-        for read in pending:
-            _discard_read(read)
-        # let the streams finish their cleanup before the mixer counts as closed
-        await asyncio.gather(*pending, return_exceptions=True)
+        self._closed = True
+        try:
+            self._mixer_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._mixer_task
+            pending = list(self._pending_reads.values())
+            self._pending_reads.clear()
+            for read in pending:
+                _discard_read(read)
+            # let the streams finish their cleanup before the mixer counts as closed
+            await asyncio.gather(*pending, return_exceptions=True)
+        finally:
+            while not self._queue.empty():
+                self._queue.get_nowait()
+            self._queue.put_nowait(None)
 
     def end_input(self) -> None:
         """

tests/rtc/test_mixer.py

@@ -191,3 +191,71 @@ async def slow_cleanup():
 
     await mixer.aclose()
     assert cleaned_up
+
+
+@pytest.mark.asyncio
+async def test_mixer_close_unblocks_waiting_consumer():
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1)
+    waiting = asyncio.create_task(mixer.__anext__())
+    await asyncio.sleep(0)
+    try:
+        await mixer.aclose()
+        with pytest.raises(StopAsyncIteration):
+            await asyncio.wait_for(waiting, timeout=1)
+    finally:
+        waiting.cancel()
+        await asyncio.gather(waiting, return_exceptions=True)
+        await mixer.aclose()
+
+
+@pytest.mark.asyncio
+async def test_mixer_stays_exhausted_after_end_input():
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1)
+    mixer.end_input()
+    try:
+        for _ in range(2):
+            with pytest.raises(StopAsyncIteration):
+                await asyncio.wait_for(mixer.__anext__(), timeout=1)
+    finally:
+        await mixer.aclose()
+
+
+@pytest.mark.asyncio
+async def test_mixer_close_does_not_block_on_full_output_queue():
+    second_frame = asyncio.Event()
+
+    async def frames():
+        yield AudioFrame(bytes(2), SAMPLE_RATE, 1, 1)
+        second_frame.set()
+        yield AudioFrame(bytes(2), SAMPLE_RATE, 1, 1)
+
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1, blocksize=1, capacity=1)
+    mixer.add_stream(frames())
+    try:
+        await asyncio.wait_for(second_frame.wait(), timeout=1)
+        await asyncio.wait_for(mixer.aclose(), timeout=1)
+        with pytest.raises(StopAsyncIteration):
+            await asyncio.wait_for(mixer.__anext__(), timeout=1)
+    finally:
+        await mixer.aclose()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("method", ["aclose", "end_input"])
+async def test_mixer_termination_releases_concurrent_readers(method):
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1)
+    readers = [asyncio.create_task(mixer.__anext__()) for _ in range(2)]
+    await asyncio.sleep(0)
+    try:
+        if method == "aclose":
+            await mixer.aclose()
+        else:
+            mixer.end_input()
+        for reader in readers:
+            with pytest.raises(StopAsyncIteration):
+                await asyncio.wait_for(reader, timeout=1)
+    finally:
+        for reader in readers:
+            reader.cancel()
+        await asyncio.gather(*readers, return_exceptions=True)
+        await mixer.aclose()