fix(rtc): keep an AudioMixer input alive after a read timeout

livekit-rtc/livekit/rtc/audio_mixer.py

@@ -8,6 +8,9 @@
 
 _Stream = AsyncIterator[AudioFrame]
 
+# consecutive timeouts after which a stream is dropped once end_input() was called
+_MAX_STALLED_ROUNDS = 10
+
 
 @dataclass
 class _Contribution:
@@ -52,6 +55,10 @@ def __init__(
         """
         self._streams: set[_Stream] = set()
         self._buffers: dict[_Stream, np.ndarray] = {}
+        # A read that timed out is kept for the next round: cancelling it would close an
+        # async generator stream for good.
+        self._pending_reads: dict[_Stream, asyncio.Future[AudioFrame]] = {}
+        self._stalled_rounds: dict[_Stream, int] = {}
         self._sample_rate: int = sample_rate
         self._num_channels: int = num_channels
         self._chunk_size: int = blocksize if blocksize > 0 else int(sample_rate // 10)
@@ -90,6 +97,10 @@ def remove_stream(self, stream: AsyncIterator[AudioFrame]) -> None:
         """
         self._streams.discard(stream)
         self._buffers.pop(stream, None)
+        self._stalled_rounds.pop(stream, None)
+        pending = self._pending_reads.pop(stream, None)
+        if pending is not None:
+            _discard_read(pending)
 
     def __aiter__(self) -> "AudioMixer":
         return self
@@ -105,11 +116,19 @@ 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.
+        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)
 
     def end_input(self) -> None:
         """
@@ -173,17 +192,34 @@ async def _get_contribution(
     ) -> _Contribution:
         had_data = buf.shape[0] > 0
         exhausted = False
-        while buf.shape[0] < self._chunk_size and not exhausted:
-            try:
-                frame = await asyncio.wait_for(
-                    stream.__anext__(), timeout=self._stream_timeout_ms / 1000
-                )
-            except asyncio.TimeoutError:
+        # a round scheduled before remove_stream() must not start a new read on that stream
+        while buf.shape[0] < self._chunk_size and not exhausted and stream in self._streams:
+            # The read stays in _pending_reads until it completes, so remove_stream()
+            # and aclose() can always cancel it.
+            read = self._pending_reads.get(stream)
+            if read is None:
+                read = asyncio.ensure_future(stream.__anext__())
+                self._pending_reads[stream] = read
+            done, _ = await asyncio.wait({read}, timeout=self._stream_timeout_ms / 1000)
+            if not done:
                 logger.warning(f"AudioMixer: stream {stream} timeout, ignoring")
+                stalled = self._stalled_rounds.get(stream, 0) + 1
+                self._stalled_rounds[stream] = stalled
+                # after end_input() a stream that keeps stalling is dropped, so the drain ends
+                exhausted = self._ending and stalled >= _MAX_STALLED_ROUNDS
                 break
+            self._stalled_rounds.pop(stream, None)
+            if self._pending_reads.get(stream) is read:
+                del self._pending_reads[stream]
+            try:
+                frame = read.result()
             except StopAsyncIteration:
                 exhausted = True
                 break
+            except asyncio.TimeoutError:
+                # raised by the stream itself: handled like the mixer's own timeout, as before
+                logger.warning(f"AudioMixer: stream {stream} timeout, ignoring")
+                break
             new_data = np.frombuffer(frame.data.tobytes(), dtype=np.int16).reshape(
                 -1, self._num_channels
             )
@@ -198,3 +234,14 @@ async def _get_contribution(
                 np.empty((0, self._num_channels), dtype=np.int16),
             )
         return _Contribution(stream, contrib, buf, had_data, exhausted)
+
+
+def _discard_read(read: "asyncio.Future[AudioFrame]") -> None:
+    """Cancel a read that is no longer wanted and consume its outcome whenever it ends."""
+    read.cancel()
+    read.add_done_callback(_consume_outcome)
+
+
+def _consume_outcome(read: "asyncio.Future[AudioFrame]") -> None:
+    if not read.cancelled():
+        read.exception()  # mark a failed read as handled

tests/rtc/test_mixer.py

@@ -1,9 +1,11 @@
 # type: ignore
 
+import asyncio
+
 import numpy as np
 import pytest
 
-from livekit.rtc import AudioMixer
+from livekit.rtc import AudioFrame, AudioMixer
 from livekit.rtc.utils import sine_wave_generator
 
 SAMPLE_RATE = 48000
@@ -56,3 +58,136 @@ async def test_mixer_two_sine_waves():
     # Assert that the peaks include 440Hz and 880Hz (with a tolerance of ±5 Hz)
     assert any(np.isclose(peak_freqs, 440, atol=5)), f"Expected 440 Hz in peaks, got: {peak_freqs}"
     assert any(np.isclose(peak_freqs, 880, atol=5)), f"Expected 880 Hz in peaks, got: {peak_freqs}"
+
+
+@pytest.mark.asyncio
+async def test_mixer_keeps_async_generator_after_a_timeout():
+    """
+    A stream that misses the timeout once must not be closed: cancelling the
+    pending read of an async generator ends it, which dropped the rest of its audio.
+    """
+    frames, block = 10, 480
+
+    async def late_once():
+        for i in range(frames):
+            if i == 3:
+                await asyncio.sleep(0.25)
+            yield AudioFrame(np.full(block, 1000, dtype=np.int16).tobytes(), SAMPLE_RATE, 1, block)
+
+    mixer = AudioMixer(
+        sample_rate=SAMPLE_RATE, num_channels=1, blocksize=block, stream_timeout_ms=100
+    )
+    mixer.add_stream(late_once())
+    mixer.end_input()
+
+    with_audio = 0
+    async for frame in mixer:
+        if np.frombuffer(frame.data.tobytes(), dtype=np.int16).any():
+            with_audio += 1
+    await mixer.aclose()
+
+    assert with_audio == frames
+
+
+@pytest.mark.asyncio
+async def test_mixer_drops_a_stalled_stream_after_end_input():
+    async def one_frame_then_stall():
+        yield AudioFrame(np.full(480, 1000, dtype=np.int16).tobytes(), SAMPLE_RATE, 1, 480)
+        await asyncio.Event().wait()
+
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1, blocksize=480, stream_timeout_ms=50)
+    mixer.add_stream(one_frame_then_stall())
+    mixer.end_input()
+
+    async def drain():
+        return [frame async for frame in mixer]
+
+    draining = asyncio.ensure_future(drain())
+    done, _ = await asyncio.wait({draining}, timeout=1)
+    await mixer.aclose()
+    assert done
+    assert len(draining.result()) == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("how", ["remove_stream", "aclose"])
+async def test_mixer_cancels_a_pending_read_when_the_stream_goes_away(how):
+    closed = asyncio.Event()
+
+    async def never_yields():
+        try:
+            await asyncio.Event().wait()
+            yield  # pragma: no cover
+        finally:
+            closed.set()
+
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1, blocksize=480, stream_timeout_ms=50)
+    stream = never_yields()
+    mixer.add_stream(stream)
+    await asyncio.sleep(0.12)  # at least one round has timed out on the stream
+
+    if how == "remove_stream":
+        mixer.remove_stream(stream)
+    else:
+        await mixer.aclose()
+
+    await asyncio.wait_for(closed.wait(), timeout=1)
+    assert not mixer._pending_reads
+    await mixer.aclose()
+
+
+@pytest.mark.asyncio
+async def test_mixer_starts_no_read_on_a_removed_stream():
+    # a round can be scheduled before remove_stream() and run after it
+    async def never_yields():
+        await asyncio.Event().wait()
+        yield  # pragma: no cover
+
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1, blocksize=480, stream_timeout_ms=10)
+    stream = never_yields()
+    mixer.add_stream(stream)
+    mixer.remove_stream(stream)
+
+    await mixer._get_contribution(stream, np.empty((0, 1), dtype=np.int16))
+    assert not mixer._pending_reads
+    await mixer.aclose()
+
+
+@pytest.mark.asyncio
+async def test_mixer_keeps_buffered_audio_when_a_stream_raises_timeout_error():
+    async def partial_then_timeout():
+        yield AudioFrame(np.full(240, 1000, dtype=np.int16).tobytes(), SAMPLE_RATE, 1, 240)
+        raise asyncio.TimeoutError
+
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1, blocksize=480)
+    mixer.add_stream(partial_then_timeout())
+    mixer.end_input()
+
+    with_audio = 0
+    async for frame in mixer:
+        if np.frombuffer(frame.data.tobytes(), dtype=np.int16).any():
+            with_audio += 1
+    await mixer.aclose()
+
+    assert with_audio == 1
+
+
+@pytest.mark.asyncio
+async def test_mixer_aclose_waits_for_a_cancelled_read_to_clean_up():
+    cleaned_up = False
+
+    async def slow_cleanup():
+        nonlocal cleaned_up
+        try:
+            await asyncio.Event().wait()
+            yield  # pragma: no cover
+        finally:
+            await asyncio.sleep(0.1)
+            cleaned_up = True
+
+    mixer = AudioMixer(sample_rate=SAMPLE_RATE, num_channels=1, blocksize=480, stream_timeout_ms=50)
+    mixer.add_stream(slow_cleanup())
+    await asyncio.sleep(0.12)
+
+    await mixer.aclose()
+    assert cleaned_up