diff --git a/livekit-rtc/livekit/rtc/audio_mixer.py b/livekit-rtc/livekit/rtc/audio_mixer.py index f293184a..0eb36aac 100644 --- a/livekit-rtc/livekit/rtc/audio_mixer.py +++ b/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,8 +107,12 @@ 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 @@ -115,20 +120,27 @@ 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: """ diff --git a/tests/rtc/test_mixer.py b/tests/rtc/test_mixer.py index ffb450cc..0f97a3e1 100644 --- a/tests/rtc/test_mixer.py +++ b/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()