fix: Use filter_fn instead of event_types to preserve generic abstraction

Addresses review feedback: FfiQueue is Generic[T], so we can't assume
item has WhichOneof method. Instead, use a filter_fn callback that the
caller provides - this keeps FfiQueue generic while allowing filtering.

- FfiQueue.subscribe() now takes optional filter_fn: Callable[[T], bool]
- AudioStream/VideoStream provide the filter that knows the concrete type
- Tests updated to use filter_fn approach

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: Hormold

livekit-rtc/livekit/rtc/_ffi_client.py

@@ -24,7 +24,7 @@
 import platform
 import atexit
 import threading
-from typing import Generic, List, Optional, Set, TypeVar
+from typing import Callable, Generic, List, Optional, TypeVar
 
 from ._proto import ffi_pb2 as proto_ffi
 from ._utils import Queue, classproperty
@@ -95,25 +95,21 @@ def __repr__(self) -> str:
 class FfiQueue(Generic[T]):
     def __init__(self) -> None:
         self._lock = threading.RLock()
-        # Format: (queue, loop, event_types or None)
+        # Format: (queue, loop, filter_fn or None)
         self._subscribers: List[
-            tuple[Queue[T], asyncio.AbstractEventLoop, Optional[Set[str]]]
+            tuple[Queue[T], asyncio.AbstractEventLoop, Optional[Callable[[T], bool]]]
         ] = []
 
     def put(self, item: T) -> None:
-        # Get event type for filtering (if item has WhichOneof method)
-        which = None
-        try:
-            which = item.WhichOneof("message")  # type: ignore
-        except Exception:
-            pass
-
         with self._lock:
-            for queue, loop, event_types in self._subscribers:
-                # Filter: if event_types specified and we know the type, skip non-matching
-                if event_types is not None and which is not None:
-                    if which not in event_types:
-                        continue
+            for queue, loop, filter_fn in self._subscribers:
+                # If filter provided, skip items that don't match
+                if filter_fn is not None:
+                    try:
+                        if not filter_fn(item):
+                            continue
+                    except Exception:
+                        pass  # On filter error, deliver the item
 
                 try:
                     loop.call_soon_threadsafe(queue.put_nowait, item)
@@ -125,22 +121,23 @@ def put(self, item: T) -> None:
     def subscribe(
         self,
         loop: Optional[asyncio.AbstractEventLoop] = None,
-        event_types: Optional[Set[str]] = None,
+        filter_fn: Optional[Callable[[T], bool]] = None,
     ) -> Queue[T]:
         """Subscribe to FFI events.
 
         Args:
             loop: Event loop to use (defaults to current).
-            event_types: Optional set of event type names to receive (e.g., {"audio_stream_event"}).
-                        If None, receives all events (original behavior).
+            filter_fn: Optional filter function. If provided, only items where
+                      filter_fn(item) returns True will be delivered.
+                      If None, receives all events (original behavior).
 
         Returns:
             Queue to receive events from.
         """
         with self._lock:
             queue = Queue[T]()
             loop = loop or asyncio.get_event_loop()
-            self._subscribers.append((queue, loop, event_types))
+            self._subscribers.append((queue, loop, filter_fn))
             return queue
 
     def unsubscribe(self, queue: Queue[T]) -> None:

livekit-rtc/livekit/rtc/audio_stream.py

@@ -105,7 +105,8 @@ def __init__(
         # Only subscribe to audio_stream_event to avoid unnecessary memory allocations
         # from other event types (room_event, track_event, etc.)
         self._ffi_queue = FfiClient.instance.queue.subscribe(
-            self._loop, event_types={"audio_stream_event"}
+            self._loop,
+            filter_fn=lambda e: e.WhichOneof("message") == "audio_stream_event",
         )
         self._queue: RingQueue[AudioFrameEvent | None] = RingQueue(capacity)
 

livekit-rtc/livekit/rtc/video_stream.py

@@ -50,7 +50,8 @@ def __init__(
         # Only subscribe to video_stream_event to avoid unnecessary memory allocations
         # from other event types (room_event, track_event, etc.)
         self._ffi_queue = FfiClient.instance.queue.subscribe(
-            self._loop, event_types={"video_stream_event"}
+            self._loop,
+            filter_fn=lambda e: e.WhichOneof("message") == "video_stream_event",
         )
         self._queue: RingQueue[VideoFrameEvent | None] = RingQueue(capacity)
         self._track: Track | None = track

tests/rtc/test_ffi_queue.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for FfiQueue event filtering functionality.
+"""Tests for FfiQueue filter_fn functionality.
 
-These tests verify the event_types filtering feature of FfiQueue without
+These tests verify the filter_fn feature of FfiQueue without
 requiring the native FFI library.
 """
 
 import asyncio
 import threading
 from dataclasses import dataclass
-from typing import Generic, List, Optional, Set, TypeVar
+from typing import Callable, Generic, List, Optional, TypeVar
 from unittest.mock import MagicMock
 
 import pytest
@@ -47,26 +47,23 @@ def empty(self) -> bool:
 
 
 class FfiQueue(Generic[T]):
-    """Copy of FfiQueue with event_types filtering for testing."""
+    """Copy of FfiQueue with filter_fn for testing."""
 
     def __init__(self) -> None:
         self._lock = threading.RLock()
         self._subscribers: List[
-            tuple[Queue[T], asyncio.AbstractEventLoop, Optional[Set[str]]]
+            tuple[Queue[T], asyncio.AbstractEventLoop, Optional[Callable[[T], bool]]]
         ] = []
 
     def put(self, item: T) -> None:
-        which = None
-        try:
-            which = item.WhichOneof("message")  # type: ignore
-        except Exception:
-            pass
-
         with self._lock:
-            for queue, loop, event_types in self._subscribers:
-                if event_types is not None and which is not None:
-                    if which not in event_types:
-                        continue
+            for queue, loop, filter_fn in self._subscribers:
+                if filter_fn is not None:
+                    try:
+                        if not filter_fn(item):
+                            continue
+                    except Exception:
+                        pass  # On filter error, deliver the item
 
                 try:
                     loop.call_soon_threadsafe(queue.put_nowait, item)
@@ -76,12 +73,12 @@ def put(self, item: T) -> None:
     def subscribe(
         self,
         loop: Optional[asyncio.AbstractEventLoop] = None,
-        event_types: Optional[Set[str]] = None,
+        filter_fn: Optional[Callable[[T], bool]] = None,
     ) -> Queue[T]:
         with self._lock:
             queue = Queue[T]()
             loop = loop or asyncio.get_event_loop()
-            self._subscribers.append((queue, loop, event_types))
+            self._subscribers.append((queue, loop, filter_fn))
             return queue
 
     def unsubscribe(self, queue: Queue[T]) -> None:
@@ -102,8 +99,8 @@ def WhichOneof(self, field: str) -> str:
         return self._message_type
 
 
-class TestFfiQueueEventFiltering:
-    """Test suite for FfiQueue event_types filtering."""
+class TestFfiQueueFilterFn:
+    """Test suite for FfiQueue filter_fn functionality."""
 
     @pytest.fixture
     def event_loop(self):
@@ -113,11 +110,10 @@ def event_loop(self):
         loop.close()
 
     def test_subscribe_without_filter_receives_all_events(self, event_loop):
-        """Subscriber without event_types filter receives all events."""
+        """Subscriber without filter_fn receives all events."""
         queue = FfiQueue()
-        sub = queue.subscribe(event_loop, event_types=None)
+        sub = queue.subscribe(event_loop, filter_fn=None)
 
-        # Send various event types
         events = [
             MockFfiEvent("room_event"),
             MockFfiEvent("audio_stream_event"),
@@ -128,22 +124,22 @@ def test_subscribe_without_filter_receives_all_events(self, event_loop):
         for event in events:
             queue.put(event)
 
-        # Run event loop to process callbacks
         event_loop.run_until_complete(asyncio.sleep(0.01))
 
-        # Should receive all 4 events
         received = []
         while not sub.empty():
             received.append(sub.get_nowait())
 
         assert len(received) == 4
 
     def test_subscribe_with_filter_receives_only_matching_events(self, event_loop):
-        """Subscriber with event_types filter only receives matching events."""
+        """Subscriber with filter_fn only receives matching events."""
         queue = FfiQueue()
-        sub = queue.subscribe(event_loop, event_types={"audio_stream_event"})
+        sub = queue.subscribe(
+            event_loop,
+            filter_fn=lambda e: e.WhichOneof("message") == "audio_stream_event",
+        )
 
-        # Send various event types
         events = [
             MockFfiEvent("room_event"),
             MockFfiEvent("audio_stream_event"),
@@ -155,10 +151,8 @@ def test_subscribe_with_filter_receives_only_matching_events(self, event_loop):
         for event in events:
             queue.put(event)
 
-        # Run event loop to process callbacks
         event_loop.run_until_complete(asyncio.sleep(0.01))
 
-        # Should receive only 2 audio_stream_events
         received = []
         while not sub.empty():
             received.append(sub.get_nowait())
@@ -170,16 +164,16 @@ def test_multiple_subscribers_different_filters(self, event_loop):
         """Multiple subscribers can have different filters."""
         queue = FfiQueue()
 
-        # Subscriber 1: only audio events
-        sub_audio = queue.subscribe(event_loop, event_types={"audio_stream_event"})
-
-        # Subscriber 2: only video events
-        sub_video = queue.subscribe(event_loop, event_types={"video_stream_event"})
-
-        # Subscriber 3: all events
-        sub_all = queue.subscribe(event_loop, event_types=None)
+        sub_audio = queue.subscribe(
+            event_loop,
+            filter_fn=lambda e: e.WhichOneof("message") == "audio_stream_event",
+        )
+        sub_video = queue.subscribe(
+            event_loop,
+            filter_fn=lambda e: e.WhichOneof("message") == "video_stream_event",
+        )
+        sub_all = queue.subscribe(event_loop, filter_fn=None)
 
-        # Send mixed events
         events = [
             MockFfiEvent("room_event"),
             MockFfiEvent("audio_stream_event"),
@@ -192,7 +186,6 @@ def test_multiple_subscribers_different_filters(self, event_loop):
 
         event_loop.run_until_complete(asyncio.sleep(0.01))
 
-        # Count received events
         audio_count = 0
         while not sub_audio.empty():
             sub_audio.get_nowait()
@@ -208,15 +201,17 @@ def test_multiple_subscribers_different_filters(self, event_loop):
             sub_all.get_nowait()
             all_count += 1
 
-        assert audio_count == 2  # 2 audio events
-        assert video_count == 1  # 1 video event
-        assert all_count == 4  # all 4 events
+        assert audio_count == 2
+        assert video_count == 1
+        assert all_count == 4
 
     def test_filter_with_multiple_event_types(self, event_loop):
-        """Filter can accept multiple event types."""
+        """Filter can match multiple event types."""
         queue = FfiQueue()
         sub = queue.subscribe(
-            event_loop, event_types={"audio_stream_event", "video_stream_event"}
+            event_loop,
+            filter_fn=lambda e: e.WhichOneof("message")
+            in {"audio_stream_event", "video_stream_event"},
         )
 
         events = [
@@ -235,48 +230,46 @@ def test_filter_with_multiple_event_types(self, event_loop):
         while not sub.empty():
             received.append(sub.get_nowait())
 
-        # Should receive audio and video events only
         assert len(received) == 2
         types = {e._message_type for e in received}
         assert types == {"audio_stream_event", "video_stream_event"}
 
     def test_unsubscribe_works_with_filtered_subscriber(self, event_loop):
         """Unsubscribe correctly removes filtered subscriber."""
         queue = FfiQueue()
-        sub = queue.subscribe(event_loop, event_types={"audio_stream_event"})
+        sub = queue.subscribe(
+            event_loop,
+            filter_fn=lambda e: e.WhichOneof("message") == "audio_stream_event",
+        )
 
         queue.put(MockFfiEvent("audio_stream_event"))
         event_loop.run_until_complete(asyncio.sleep(0.01))
 
-        # Should have received 1 event
         assert not sub.empty()
 
-        # Unsubscribe
         queue.unsubscribe(sub)
 
-        # Clear the queue
         while not sub.empty():
             sub.get_nowait()
 
-        # Send more events
         queue.put(MockFfiEvent("audio_stream_event"))
         event_loop.run_until_complete(asyncio.sleep(0.01))
 
-        # Should not receive after unsubscribe
         assert sub.empty()
 
-    def test_event_without_which_oneof_passes_through(self, event_loop):
-        """Events without WhichOneof method pass through to all subscribers."""
+    def test_filter_error_delivers_item(self, event_loop):
+        """If filter_fn raises, item is still delivered."""
         queue = FfiQueue()
-        sub = queue.subscribe(event_loop, event_types={"audio_stream_event"})
 
-        # Event without WhichOneof
-        plain_event = MagicMock(spec=[])  # No WhichOneof method
+        def bad_filter(e):
+            raise ValueError("oops")
 
-        queue.put(plain_event)
+        sub = queue.subscribe(event_loop, filter_fn=bad_filter)
+
+        queue.put(MockFfiEvent("audio_stream_event"))
         event_loop.run_until_complete(asyncio.sleep(0.01))
 
-        # Should still receive it (can't filter without type info)
+        # Item should be delivered despite filter error
         received = []
         while not sub.empty():
             received.append(sub.get_nowait())
@@ -300,18 +293,20 @@ def test_filtering_reduces_callback_calls(self, event_loop):
         # Create 10 subscribers, each only wants audio events
         subscribers = []
         for _ in range(10):
-            sub = queue.subscribe(event_loop, event_types={"audio_stream_event"})
+            sub = queue.subscribe(
+                event_loop,
+                filter_fn=lambda e: e.WhichOneof("message") == "audio_stream_event",
+            )
             subscribers.append(sub)
 
         # Generate 1000 events, only 5% are audio
         events = []
         for i in range(1000):
-            if i < 50:  # 5% audio events
+            if i < 50:
                 events.append(MockFfiEvent("audio_stream_event"))
             else:
                 events.append(MockFfiEvent("room_event"))
 
-        # Process all events
         for event in events:
             queue.put(event)
 
@@ -324,7 +319,3 @@ def test_filtering_reduces_callback_calls(self, event_loop):
                 sub.get_nowait()
                 count += 1
             assert count == 50
-
-        # Total callbacks made: 10 subscribers × 50 audio events = 500
-        # Without filtering: 10 subscribers × 1000 events = 10,000
-        # This is a 95% reduction in callback/object creation