diff --git a/livekit-rtc/livekit/rtc/_ffi_client.py b/livekit-rtc/livekit/rtc/_ffi_client.py index 096b7ea5..0938b28d 100644 --- a/livekit-rtc/livekit/rtc/_ffi_client.py +++ b/livekit-rtc/livekit/rtc/_ffi_client.py @@ -115,13 +115,16 @@ def __repr__(self) -> str: class FfiQueue(Generic[T]): def __init__(self) -> None: self._lock = threading.RLock() + self._put_lock = threading.RLock() # Format: (queue, loop, filter_fn or None) self._subscribers: List[ tuple[Queue[T], asyncio.AbstractEventLoop, Optional[Callable[[T], bool]]] ] = [] def put(self, item: T) -> None: - with self._lock: + # _put_lock keeps concurrent puts in order at every subscriber. subscribe and + # unsubscribe replace the list instead of mutating it, so they don't wait for a put + with self._put_lock: for queue, loop, filter_fn in self._subscribers: # If filter provided, skip items that don't match if filter_fn is not None: @@ -157,7 +160,7 @@ def subscribe( with self._lock: queue = Queue[T]() loop = loop or asyncio.get_event_loop() - self._subscribers.append((queue, loop, filter_fn)) + self._subscribers = [*self._subscribers, (queue, loop, filter_fn)] return queue def unsubscribe(self, queue: Queue[T]) -> None: @@ -165,7 +168,7 @@ def unsubscribe(self, queue: Queue[T]) -> None: # looping here is ok, since we don't expect a lot of subscribers for i, (q, _, _) in enumerate(self._subscribers): if q == queue: - self._subscribers.pop(i) + self._subscribers = self._subscribers[:i] + self._subscribers[i + 1 :] break diff --git a/tests/rtc/test_ffi_queue.py b/tests/rtc/test_ffi_queue.py index 5f76aa6e..ef0c65ac 100644 --- a/tests/rtc/test_ffi_queue.py +++ b/tests/rtc/test_ffi_queue.py @@ -19,6 +19,7 @@ """ import asyncio +import threading from collections.abc import Generator from dataclasses import dataclass @@ -266,3 +267,93 @@ def test_filtering_reduces_callback_calls(self, event_loop: asyncio.AbstractEven sub.get_nowait() count += 1 assert count == 50 + + +class TestFfiQueueConcurrency: + """Test put running concurrently with other puts and with subscribe/unsubscribe.""" + + @pytest.fixture + def event_loop(self) -> Generator[asyncio.AbstractEventLoop, None, None]: + loop = asyncio.new_event_loop() + yield loop + loop.close() + + def test_subscribe_and_unsubscribe_do_not_wait_for_put( + self, event_loop: asyncio.AbstractEventLoop + ) -> None: + """subscribe/unsubscribe return while put is still delivering on another thread.""" + queue: FfiQueue[MockFfiEvent] = FfiQueue() + delivering = threading.Event() + release = threading.Event() + released: list[bool] = [] + + def blocking_filter(e: MockFfiEvent) -> bool: + delivering.set() + released.append(release.wait(timeout=5)) + return True + + sub = queue.subscribe(event_loop, filter_fn=blocking_filter) + put_thread = threading.Thread(target=queue.put, args=(MockFfiEvent("room_event"),)) + put_thread.start() + assert delivering.wait(timeout=5) + + queue.unsubscribe(queue.subscribe(event_loop)) + release.set() + put_thread.join() + event_loop.run_until_complete(asyncio.sleep(0.01)) + + assert released == [True] + assert sub.qsize() == 1 + + def test_unsubscribe_during_put_does_not_skip_other_subscribers( + self, event_loop: asyncio.AbstractEventLoop + ) -> None: + """A subscriber removed during put does not make the next one miss the item.""" + queue: FfiQueue[MockFfiEvent] = FfiQueue() + + def unsubscribe_self(e: MockFfiEvent) -> bool: + queue.unsubscribe(sub) + return True + + sub = queue.subscribe(event_loop, filter_fn=unsubscribe_self) + other = queue.subscribe(event_loop) + + queue.put(MockFfiEvent("room_event")) + event_loop.run_until_complete(asyncio.sleep(0.01)) + + assert other.qsize() == 1 + + def test_concurrent_puts_keep_their_order(self, event_loop: asyncio.AbstractEventLoop) -> None: + """A put from another thread does not overtake one that is still delivering.""" + queue: FfiQueue[MockFfiEvent] = FfiQueue() + delivering = threading.Event() + second_done = threading.Event() + + def slow_filter(e: MockFfiEvent) -> bool: + if e.WhichOneof("message") == "room_event": + delivering.set() + # only set in time if the second put overtook this one + second_done.wait(timeout=0.2) + return True + + def put_second() -> None: + queue.put(MockFfiEvent("unpublish_track")) + second_done.set() + + queue.subscribe(event_loop, filter_fn=slow_filter) + sub = queue.subscribe(event_loop) + + first = threading.Thread(target=queue.put, args=(MockFfiEvent("room_event"),)) + first.start() + assert delivering.wait(timeout=5) + second = threading.Thread(target=put_second) + second.start() + first.join() + second.join() + event_loop.run_until_complete(asyncio.sleep(0.01)) + + received = [] + while not sub.empty(): + received.append(sub.get_nowait()._message_type) + + assert received == ["room_event", "unpublish_track"]