From ee98e5f1ddabc6697b03075097d3699c5af7ecc8 Mon Sep 17 00:00:00 2001 From: Swayam Gupta Date: Wed, 28 Oct 2026 20:20:50 +0530 Subject: [PATCH] fix(rtc): deliver media stream events only to their own stream --- livekit-rtc/livekit/rtc/audio_source.py | 12 +- livekit-rtc/livekit/rtc/audio_stream.py | 12 +- livekit-rtc/livekit/rtc/video_stream.py | 12 +- tests/rtc/test_media_ffi_filter.py | 172 ++++++++++++++++++++++++ 4 files changed, 201 insertions(+), 7 deletions(-) create mode 100644 tests/rtc/test_media_ffi_filter.py diff --git a/livekit-rtc/livekit/rtc/audio_source.py b/livekit-rtc/livekit/rtc/audio_source.py index 63cc1a5d..680cc803 100644 --- a/livekit-rtc/livekit/rtc/audio_source.py +++ b/livekit-rtc/livekit/rtc/audio_source.py @@ -136,9 +136,19 @@ async def capture_frame(self, frame: AudioFrame) -> None: 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) + # The async_id is only known once the request returns; until then every capture + # callback is let through, and wait_for drops the ones of other captures. + async_id: int | None = None + queue = FfiClient.instance.queue.subscribe( + loop=self._loop, + filter_fn=lambda e: ( + e.WhichOneof("message") == "capture_audio_frame" + and (async_id is None or e.capture_audio_frame.async_id == async_id) + ), + ) try: resp = FfiClient.instance.request(req) + async_id = resp.capture_audio_frame.async_id cb: proto_ffi.FfiEvent = await queue.wait_for( lambda e: e.capture_audio_frame.async_id == resp.capture_audio_frame.async_id ) diff --git a/livekit-rtc/livekit/rtc/audio_stream.py b/livekit-rtc/livekit/rtc/audio_stream.py index 97772f4f..727bf17b 100644 --- a/livekit-rtc/livekit/rtc/audio_stream.py +++ b/livekit-rtc/livekit/rtc/audio_stream.py @@ -106,11 +106,16 @@ def __init__( self._num_channels = num_channels self._frame_size_ms = frame_size_ms self._loop = loop or asyncio.get_event_loop() - # Only subscribe to audio_stream_event to avoid unnecessary memory allocations - # from other event types (room_event, track_event, etc.) + # Subscribe before creating the stream so none of its events are missed. Its + # handle is only known once the request returns; until then every audio stream + # event is let through, and _run drops the ones of other streams. + handle_id: int | None = None self._ffi_queue = FfiClient.instance.queue.subscribe( self._loop, - filter_fn=lambda e: e.WhichOneof("message") == "audio_stream_event", + filter_fn=lambda e: ( + e.WhichOneof("message") == "audio_stream_event" + and (handle_id is None or e.audio_stream_event.stream_handle == handle_id) + ), ) self._queue: RingQueue[AudioFrameEvent | None] = RingQueue(capacity) @@ -133,6 +138,7 @@ def __init__( else: stream = self._create_owned_stream() self._ffi_handle = FfiHandle(stream.handle.id) + handle_id = stream.handle.id self._info = stream.info # Start _run only after _ffi_handle is set, so a failure above doesn't orphan diff --git a/livekit-rtc/livekit/rtc/video_stream.py b/livekit-rtc/livekit/rtc/video_stream.py index eeced550..07938fe5 100644 --- a/livekit-rtc/livekit/rtc/video_stream.py +++ b/livekit-rtc/livekit/rtc/video_stream.py @@ -48,11 +48,16 @@ def __init__( **kwargs: Any, ) -> None: self._loop = loop or asyncio.get_event_loop() - # Only subscribe to video_stream_event to avoid unnecessary memory allocations - # from other event types (room_event, track_event, etc.) + # Subscribe before creating the stream so none of its events are missed. Its + # handle is only known once the request returns; until then every video stream + # event is let through, and _run drops the ones of other streams. + handle_id: int | None = None self._ffi_queue = FfiClient.instance.queue.subscribe( self._loop, - filter_fn=lambda e: e.WhichOneof("message") == "video_stream_event", + filter_fn=lambda e: ( + e.WhichOneof("message") == "video_stream_event" + and (handle_id is None or e.video_stream_event.stream_handle == handle_id) + ), ) self._queue: RingQueue[VideoFrameEvent | None] = RingQueue(capacity) self._track: Track | None = track @@ -68,6 +73,7 @@ def __init__( stream = self._create_owned_stream() self._ffi_handle = FfiHandle(stream.handle.id) + handle_id = stream.handle.id self._info = stream.info self._task = self._loop.create_task(self._run()) diff --git a/tests/rtc/test_media_ffi_filter.py b/tests/rtc/test_media_ffi_filter.py new file mode 100644 index 00000000..50742dee --- /dev/null +++ b/tests/rtc/test_media_ffi_filter.py @@ -0,0 +1,172 @@ +# 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. + +"""Exercise the media streams' and AudioSource's FFI subscriptions without a native FFI.""" + +import asyncio +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Union +from unittest.mock import patch + +import pytest + +from livekit import rtc +from livekit.rtc._ffi_client import FfiClient, FfiQueue +from livekit.rtc._proto import ffi_pb2 as proto_ffi +from livekit.rtc._proto import track_pb2 as proto_track + + +@pytest.fixture +def queue(monkeypatch: pytest.MonkeyPatch) -> Iterator[FfiQueue[proto_ffi.FfiEvent]]: + queue = FfiQueue[proto_ffi.FfiEvent]() + handles = iter(range(100, 200)) + + def request(req: proto_ffi.FfiRequest) -> proto_ffi.FfiResponse: + response = proto_ffi.FfiResponse() + which = req.WhichOneof("message") + if which == "new_audio_stream": + response.new_audio_stream.stream.handle.id = next(handles) + elif which == "new_video_stream": + response.new_video_stream.stream.handle.id = next(handles) + elif which == "new_audio_source": + response.new_audio_source.source.handle.id = next(handles) + elif which == "capture_audio_frame": + response.capture_audio_frame.async_id = 1 + else: + raise AssertionError(f"unexpected FFI request: {which}") + return response + + # A different PID prevents synthetic handles from being dropped natively. + monkeypatch.setattr( + FfiClient, "_instance", SimpleNamespace(queue=queue, request=request, _pid=-1) + ) + yield queue + assert not queue._subscribers + + +def _create_stream(event_type: str) -> Union[rtc.AudioStream, rtc.VideoStream]: + track = rtc.Track(proto_track.OwnedTrack()) + if event_type == "audio_stream_event": + return rtc.AudioStream(track) + return rtc.VideoStream(track) + + +def _stream_event(event_type: str, stream_handle: int, *, eos: bool = False) -> proto_ffi.FfiEvent: + event = proto_ffi.FfiEvent() + stream_event = getattr(event, event_type) + stream_event.stream_handle = stream_handle + if eos: + stream_event.eos.SetInParent() + return event + + +@pytest.mark.parametrize("event_type", ["audio_stream_event", "video_stream_event"]) +async def test_stream_does_not_receive_other_streams_events( + queue: FfiQueue[proto_ffi.FfiEvent], event_type: str +) -> None: + first = _create_stream(event_type) + second = _create_stream(event_type) + loop = asyncio.get_running_loop() + event = _stream_event(event_type, second._ffi_handle.handle) + + try: + with patch.object( + loop, "call_soon_threadsafe", wraps=loop.call_soon_threadsafe + ) as schedule: + for _ in range(100): + queue.put(event) + + # Only the stream the events belong to should incur an event-loop wakeup. + assert schedule.call_count == 100 + assert all(call.args[0] == second._ffi_queue.put_nowait for call in schedule.call_args_list) + finally: + for stream in (first, second): + queue.put(_stream_event(event_type, stream._ffi_handle.handle, eos=True)) + await asyncio.wait_for(stream.aclose(), 1) + + +@pytest.mark.parametrize("event_type", ["audio_stream_event", "video_stream_event"]) +async def test_stream_receives_events_sent_while_it_is_created( + queue: FfiQueue[proto_ffi.FfiEvent], + monkeypatch: pytest.MonkeyPatch, + event_type: str, +) -> None: + instance = FfiClient._instance + assert instance is not None + create = instance.request + + def request(req: proto_ffi.FfiRequest) -> proto_ffi.FfiResponse: + response = create(req) + which = req.WhichOneof("message") + assert which is not None + # the native stream can end before the request that created it returns + handle = getattr(response, which).stream.handle.id + queue.put(_stream_event(event_type, handle, eos=True)) + return response + + monkeypatch.setattr(instance, "request", request) + stream = _create_stream(event_type) + await asyncio.wait_for(stream.aclose(), 1) + + +async def test_capture_frame_only_receives_its_own_callback( + queue: FfiQueue[proto_ffi.FfiEvent], +) -> None: + source = rtc.AudioSource(48000, 1) + capture = asyncio.create_task(source.capture_frame(rtc.AudioFrame.create(48000, 1, 480))) + await asyncio.sleep(0) + + loop = asyncio.get_running_loop() + room_event = proto_ffi.FfiEvent() + room_event.room_event.SetInParent() + other_capture = proto_ffi.FfiEvent() + other_capture.capture_audio_frame.async_id = 2 + try: + with patch.object( + loop, "call_soon_threadsafe", wraps=loop.call_soon_threadsafe + ) as schedule: + for event in (room_event, _stream_event("audio_stream_event", 1), other_capture): + queue.put(event) + + # Neither other events nor another capture's callback should wake the capture up. + assert schedule.call_count == 0 + finally: + callback = proto_ffi.FfiEvent() + callback.capture_audio_frame.async_id = 1 + queue.put(callback) + await asyncio.wait_for(capture, 1) + await source.aclose() + + +async def test_capture_frame_receives_callback_sent_before_request_returns( + queue: FfiQueue[proto_ffi.FfiEvent], monkeypatch: pytest.MonkeyPatch +) -> None: + instance = FfiClient._instance + assert instance is not None + create = instance.request + + def request(req: proto_ffi.FfiRequest) -> proto_ffi.FfiResponse: + response = create(req) + if req.WhichOneof("message") == "capture_audio_frame": + # the native capture can complete before the request that started it returns + callback = proto_ffi.FfiEvent() + callback.capture_audio_frame.async_id = response.capture_audio_frame.async_id + queue.put(callback) + return response + + monkeypatch.setattr(instance, "request", request) + source = rtc.AudioSource(48000, 1) + await asyncio.wait_for(source.capture_frame(rtc.AudioFrame.create(48000, 1, 480)), 1) + await source.aclose()