diff --git a/livekit-rtc/livekit/rtc/audio_source.py b/livekit-rtc/livekit/rtc/audio_source.py index 63cc1a5d..64180a43 100644 --- a/livekit-rtc/livekit/rtc/audio_source.py +++ b/livekit-rtc/livekit/rtc/audio_source.py @@ -14,8 +14,9 @@ from __future__ import annotations -import time import asyncio +import threading +import time from ._ffi_client import FfiHandle, FfiClient from ._proto import audio_frame_pb2 as proto_audio_frame @@ -53,7 +54,12 @@ def __init__( """ self._sample_rate = sample_rate self._num_channels = num_channels - self._loop = loop or asyncio.get_event_loop() + self._loop: asyncio.AbstractEventLoop | None = loop + if self._loop is None: + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + self._loop = None req = proto_ffi.FfiRequest() req.new_audio_source.type = proto_audio_frame.AudioSourceType.AUDIO_SOURCE_NATIVE @@ -67,6 +73,8 @@ def __init__( self._last_capture = 0.0 self._q_size = 0.0 + self._capture_gen = 0 + self._lock = threading.Lock() self._join_handle: asyncio.TimerHandle | None = None self._join_fut: asyncio.Future[None] | None = None @@ -83,7 +91,8 @@ def num_channels(self) -> int: @property def queued_duration(self) -> float: """The current duration (in seconds) of audio data queued for playback.""" - return max(self._q_size - time.monotonic() + self._last_capture, 0.0) + with self._lock: + return max(self._q_size - time.monotonic() + self._last_capture, 0.0) def clear_queue(self) -> None: """ @@ -97,7 +106,16 @@ def clear_queue(self) -> None: req = proto_ffi.FfiRequest() req.clear_audio_buffer.source_handle = self._ffi_handle.handle _ = FfiClient.instance.request(req) - self._release_waiter() + with self._lock: + self._last_capture = 0.0 + self._q_size = 0.0 + target_gen = self._capture_gen + self._release_waiter_threadsafe(gen=target_gen) + + def _get_loop(self) -> asyncio.AbstractEventLoop: + if self._loop is None or self._loop.is_closed(): + self._loop = asyncio.get_running_loop() + return self._loop async def capture_frame(self, frame: AudioFrame) -> None: """ @@ -116,27 +134,38 @@ async def capture_frame(self, frame: AudioFrame) -> None: Exception: If there is an error during frame capture. """ - if frame.samples_per_channel == 0 or self._ffi_handle.disposed: + if frame.samples_per_channel == 0 or getattr(self._ffi_handle, "disposed", False): return - now = time.monotonic() - elapsed = 0.0 if self._last_capture == 0.0 else now - self._last_capture - self._q_size += frame.samples_per_channel / self.sample_rate - elapsed - self._last_capture = now + loop = self._get_loop() - if self._join_handle: - self._join_handle.cancel() + with self._lock: + self._capture_gen += 1 + current_gen = self._capture_gen - if self._join_fut is None: - self._join_fut = self._loop.create_future() + now = time.monotonic() + elapsed = 0.0 if self._last_capture == 0.0 else now - self._last_capture + self._q_size += frame.samples_per_channel / self.sample_rate - elapsed + self._last_capture = now + q_size = self._q_size + + if self._join_handle: + self._join_handle.cancel() - self._join_handle = self._loop.call_later(self._q_size, self._release_waiter) + if ( + self._join_fut is None + or self._join_fut.get_loop() is not loop + or self._join_fut.get_loop().is_closed() + ): + self._join_fut = loop.create_future() + + self._join_handle = loop.call_later(q_size, self._release_waiter, current_gen) req = proto_ffi.FfiRequest() req.capture_audio_frame.source_handle = self._ffi_handle.handle req.capture_audio_frame.buffer.CopyFrom(frame._proto_info()) - queue = FfiClient.instance.queue.subscribe(loop=self._loop) + queue = FfiClient.instance.queue.subscribe(loop=loop) try: resp = FfiClient.instance.request(req) cb: proto_ffi.FfiEvent = await queue.wait_for( @@ -162,20 +191,49 @@ async def wait_for_playout(self) -> None: await asyncio.shield(self._join_fut) - def _release_waiter(self) -> None: - if self._join_fut is None: - return # could be None when clear_queue is called + def _release_waiter(self, gen: int | None = None) -> None: + with self._lock: + if gen is not None and gen != self._capture_gen: + return - if not self._join_fut.done(): - self._join_fut.set_result(None) + if self._join_handle is not None: + self._join_handle.cancel() + self._join_handle = None - self._last_capture = 0.0 - self._q_size = 0.0 - self._join_fut = None + if self._join_fut is None: + return # could be None when clear_queue is called - async def aclose(self) -> None: - """Close the audio source + if not self._join_fut.done(): + self._join_fut.set_result(None) + + self._last_capture = 0.0 + self._q_size = 0.0 + self._join_fut = None + + def _release_waiter_threadsafe(self, gen: int | None = None) -> None: + try: + running_loop = asyncio.get_running_loop() + except RuntimeError: + running_loop = None + + if self._loop is not None and self._loop.is_running() and running_loop is not self._loop: + self._loop.call_soon_threadsafe(self._release_waiter, gen) + else: + self._release_waiter(gen) + + def close(self) -> None: + """Close the audio source. This method cleans up resources associated with the audio source. """ + self._release_waiter_threadsafe() self._ffi_handle.dispose() + + async def aclose(self) -> None: + """Close the audio source without blocking the event loop. + + This method cleans up resources associated with the audio source. + """ + self._release_waiter() + if not getattr(self._ffi_handle, "disposed", False): + await asyncio.to_thread(self._ffi_handle.dispose) diff --git a/livekit-rtc/livekit/rtc/video_source.py b/livekit-rtc/livekit/rtc/video_source.py index c8dd2a02..6b2c3947 100644 --- a/livekit-rtc/livekit/rtc/video_source.py +++ b/livekit-rtc/livekit/rtc/video_source.py @@ -14,6 +14,8 @@ from __future__ import annotations +import asyncio + from ._ffi_client import FfiHandle, FfiClient from ._proto import ffi_pb2 as proto_ffi from ._proto import video_frame_pb2 as proto_video @@ -61,6 +63,9 @@ def capture_frame( metadata: proto_video.FrameMetadata | None = None, ) -> None: """Capture a frame, optionally attaching frame metadata.""" + if getattr(self._ffi_handle, "disposed", False): + return + req = proto_ffi.FfiRequest() req.capture_video_frame.source_handle = self._ffi_handle.handle req.capture_video_frame.buffer.CopyFrom(frame._proto_info()) @@ -70,5 +75,17 @@ def capture_frame( req.capture_video_frame.metadata.CopyFrom(metadata) FfiClient.instance.request(req) - async def aclose(self) -> None: + def close(self) -> None: + """Close the video source. + + This method cleans up resources associated with the video source. + """ self._ffi_handle.dispose() + + async def aclose(self) -> None: + """Close the video source without blocking the event loop. + + This method cleans up resources associated with the video source. + """ + if not getattr(self._ffi_handle, "disposed", False): + await asyncio.to_thread(self._ffi_handle.dispose) diff --git a/livekit-rtc/tests/test_ffi_handle_lifecycle.py b/livekit-rtc/tests/test_ffi_handle_lifecycle.py index 08750bfa..442a7bae 100644 --- a/livekit-rtc/tests/test_ffi_handle_lifecycle.py +++ b/livekit-rtc/tests/test_ffi_handle_lifecycle.py @@ -14,7 +14,6 @@ from __future__ import annotations -import gc import subprocess import sys @@ -43,25 +42,46 @@ def test_active_runtime_invalid_handle_asserts() -> None: def test_live_handles_marked_consumed_and_no_assert_after_dispose() -> None: - handle_id = _allocate_real_handle() - handle = FfiHandle(handle_id) - assert not handle.disposed + code = """ +import gc +from livekit.rtc._ffi_client import FfiClient, FfiHandle +from livekit.rtc._proto import audio_frame_pb2, ffi_pb2 + +req = ffi_pb2.FfiRequest() +req.new_audio_source.type = audio_frame_pb2.AUDIO_SOURCE_NATIVE +req.new_audio_source.sample_rate = 48000 +req.new_audio_source.num_channels = 1 +req.new_audio_source.queue_size_ms = 100 +resp = FfiClient.instance.request(req) +handle_id = resp.new_audio_source.source.handle.id +assert handle_id != 0 + +handle = FfiHandle(int(handle_id)) +assert not handle.disposed - # Global dispose marks all live handles consumed and sets _ffi_disposed - FfiClient.instance._dispose() - assert handle.disposed +# Global dispose marks all live handles consumed and sets _ffi_disposed +FfiClient.instance._dispose() +assert handle.disposed - # Explicit dispose or __del__ after global dispose should be a no-op and not assert - handle.dispose() - handle.__del__() +# Explicit dispose or __del__ after global dispose should be a no-op and not assert +handle.dispose() +handle.__del__() - # Newly instantiated handle after global dispose also safely no-ops without assert - late_handle = FfiHandle(12345) - late_handle.dispose() +# Newly instantiated handle after global dispose also safely no-ops without assert +late_handle = FfiHandle(12345) +late_handle.dispose() - del handle - del late_handle - gc.collect() +del handle +del late_handle +gc.collect() +""" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "AssertionError" not in result.stderr def test_subprocess_exit_with_undisposed_handle_has_no_assertion_error() -> None: diff --git a/livekit-rtc/tests/test_source_stream_close.py b/livekit-rtc/tests/test_source_stream_close.py new file mode 100644 index 00000000..ab97b5db --- /dev/null +++ b/livekit-rtc/tests/test_source_stream_close.py @@ -0,0 +1,198 @@ +# Copyright 2026 LiveKit, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from livekit import rtc + + +@pytest.mark.asyncio +async def test_audio_source_aclose_releases_waiter_and_cancels_timer() -> None: + source = rtc.AudioSource(48000, 1) + frame = rtc.AudioFrame.create(48000, 1, 48000) + await asyncio.wait_for(source.capture_frame(frame), timeout=5.0) + + assert source._join_fut is not None + assert source._join_handle is not None + + await asyncio.wait_for(source.aclose(), timeout=5.0) + + assert source._ffi_handle.disposed + assert source._join_fut is None + assert source._join_handle is None + + # After aclose, wait_for_playout should return immediately without hanging + await asyncio.wait_for(source.wait_for_playout(), timeout=2.0) + + # Repeated calls are idempotent + await asyncio.wait_for(source.aclose(), timeout=2.0) + + +def test_audio_source_sync_close() -> None: + source = rtc.AudioSource(48000, 1) + assert not source._ffi_handle.disposed + + source.close() + assert source._ffi_handle.disposed + + # Repeated close is idempotent + source.close() + + +def test_audio_source_cross_thread_close() -> None: + source = rtc.AudioSource(48000, 1) + assert not source._ffi_handle.disposed + + t = threading.Thread(target=source.close) + t.start() + t.join() + + assert source._ffi_handle.disposed + + +@pytest.mark.asyncio +async def test_audio_source_capture_after_close_is_noop() -> None: + source = rtc.AudioSource(48000, 1) + source.close() + + frame = rtc.AudioFrame.create(48000, 1, 480) + # Should safely return without error + await asyncio.wait_for(source.capture_frame(frame), timeout=2.0) + + +def test_video_source_sync_close() -> None: + source = rtc.VideoSource(640, 480) + assert not source._ffi_handle.disposed + + source.close() + assert source._ffi_handle.disposed + + # Repeated close is idempotent + source.close() + + +@pytest.mark.asyncio +async def test_video_source_aclose_and_capture_after_close() -> None: + source = rtc.VideoSource(640, 480) + assert not source._ffi_handle.disposed + + await asyncio.wait_for(source.aclose(), timeout=5.0) + assert source._ffi_handle.disposed + + # Capture after close is safely ignored + frame = rtc.VideoFrame( + 640, + 480, + rtc.VideoBufferType.RGBA, + b"\x00" * (640 * 480 * 4), + ) + source.capture_frame(frame) + + # Repeated aclose is idempotent + await asyncio.wait_for(source.aclose(), timeout=2.0) + + +def test_audio_source_init_outside_loop_and_capture_in_asyncio_run() -> None: + source = rtc.AudioSource(48000, 1) + assert source._loop is None + frame = rtc.AudioFrame.create(48000, 1, 480) + + async def _capture() -> None: + await asyncio.wait_for(source.capture_frame(frame), timeout=5.0) + assert source._loop is not None + assert source._loop.is_running() + await asyncio.wait_for(source.aclose(), timeout=5.0) + + asyncio.run(_capture()) + + +def test_audio_source_cross_thread_clear_queue() -> None: + source = rtc.AudioSource(48000, 1) + frame = rtc.AudioFrame.create(48000, 1, 48000) + + async def _run() -> None: + await source.capture_frame(frame) + assert source._join_fut is not None + + t = threading.Thread(target=source.clear_queue) + t.start() + t.join() + + # Playout waiter should be released by cross-thread clear_queue + await asyncio.wait_for(source.wait_for_playout(), timeout=2.0) + await source.aclose() + + asyncio.run(_run()) + + +def test_audio_source_clear_queue_does_not_cancel_subsequent_capture() -> None: + source = rtc.AudioSource(48000, 1, queue_size_ms=10000) + frame1 = rtc.AudioFrame.create(48000, 1, 48000) + frame2 = rtc.AudioFrame.create(48000, 1, 48000) + + async def _run() -> None: + await source.capture_frame(frame1) + gen1 = source._capture_gen + + # Simulate clear_queue scheduled from another thread + loop = asyncio.get_running_loop() + loop.call_soon(source._release_waiter, gen1) + + # Capture a subsequent frame before yielding + await source.capture_frame(frame2) + assert source._capture_gen > gen1 + assert source._join_fut is not None + + # Allow queued clear callback to execute + await asyncio.sleep(0.01) + + # The subsequent frame's playout future and timer must not be cleared + assert source._join_fut is not None + assert not source._join_fut.done() + + await source.aclose() + + asyncio.run(_run()) + + +def test_audio_source_cross_thread_clear_queue_resets_queued_duration() -> None: + source = rtc.AudioSource(48000, 1, queue_size_ms=10000) + frame1 = rtc.AudioFrame.create(48000, 1, 48000) + frame2 = rtc.AudioFrame.create(48000, 1, 48000) + + async def _run() -> None: + await source.capture_frame(frame1) + assert source.queued_duration > 0.8 + + t = threading.Thread(target=source.clear_queue) + t.start() + t.join() + + # Queued duration should be reset immediately + assert source.queued_duration == 0.0 + + # Capture a new frame immediately after cross-thread clear + await source.capture_frame(frame2) + + # Queued duration should reflect only the new frame (~1.0s), not ~2.0s + assert 0.8 < source.queued_duration <= 1.05 + + await source.aclose() + + asyncio.run(_run())