fix: ensure that data streams readers are never orphaned when there's an error reading via ffi

Author: 1egoman

livekit-rtc/livekit/rtc/data_stream.py

@@ -89,6 +89,25 @@ def _byte_stream_info_from_proto(info: proto_data_stream.ByteStreamInfo) -> Byte
 
 
 class TextStreamReader:
+    """An incoming text stream.
+
+    Use as an async iterator to receive chunks, or :meth:`read_all` to collect
+    the whole payload::
+
+        async for chunk in reader:
+            process(chunk)
+
+    A reader subscribes to the FFI event queue as soon as it is handed to your
+    handler, so one that is never iterated to completion — abandoned after a
+    ``break``, or received by a handler that never reads it — must be closed
+    with :meth:`close`, or its subscription lives for the rest of the process.
+    Iterating to the end closes the reader for you.
+
+    If the stream terminates abnormally (aborted by the sender, oversized, or
+    the room disconnects mid-stream), :class:`StreamError` is raised instead of
+    a normal ``StopAsyncIteration``.
+    """
+
     def __init__(
         self,
         owned_info: proto_data_stream.OwnedTextStreamReader,
@@ -151,25 +170,80 @@ async def read_all(self) -> str:
             final_string += chunk
         return final_string
 
+    def _unsubscribe(self) -> None:
+        """Drops the FFI queue subscription without ending the stream.
+
+        Events already delivered stay in the queue and remain readable; only
+        the global subscription, and the filter it runs against every FFI
+        event, goes away. Safe to call repeatedly — unsubscribing a queue that
+        is no longer registered is a no-op.
+
+        Only the subscription is released: the reader handle was consumed by
+        the read_incremental request in __init__, so there is nothing left to
+        dispose (and disposing it would drop a handle the FFI server no longer
+        owns).
+        """
+        FfiClient.instance.queue.unsubscribe(self._queue)
+
     def _close(self) -> None:
         if not self._closed:
             self._closed = True
-            FfiClient.instance.queue.unsubscribe(self._queue)
+            self._unsubscribe()
             if self._on_close is not None:
                 self._on_close()
 
+    def close(self) -> None:
+        """Explicitly close the reader and unsubscribe.
+
+        Call this on a reader you stop consuming before it ends. Closing a
+        reader that already reached end-of-stream is a no-op, and iterating a
+        closed reader raises ``StopAsyncIteration`` (or :class:`StreamError`
+        if the stream had already failed).
+        """
+        self._close()
+
+    async def aclose(self) -> None:
+        self.close()
+
     def _signal_disconnect(self) -> None:
         """Injects a synthetic EOS-with-error event so pending reads wake up
-        and raise StreamError when the room disconnects mid-stream."""
+        and raise StreamError when the room disconnects mid-stream.
+
+        Also drops the queue subscription, which would otherwise outlive the
+        room: no further events can arrive once the room is gone, and the
+        injected one is already queued. The reader is deliberately left open
+        so chunks that arrived before the disconnect are still delivered
+        before the StreamError, as they were before this was unsubscribed
+        here.
+        """
         if self._closed:
             return
         event = proto_ffi.FfiEvent()
         event.text_stream_reader_event.reader_handle = self._reader_handle
         event.text_stream_reader_event.eos.error.description = _DISCONNECT_ERROR
         self._queue.put_nowait(event)
+        self._unsubscribe()
 
 
 class ByteStreamReader:
+    """An incoming byte stream.
+
+    Use as an async iterator to receive chunks::
+
+        async for chunk in reader:
+            process(chunk)
+
+    A reader subscribes to the FFI event queue as soon as it is handed to your
+    handler, so one that is never iterated to completion — abandoned after a
+    ``break``, or received by a handler that never reads it — must be closed
+    with :meth:`close`, or its subscription lives for the rest of the process.
+    Iterating to the end closes the reader for you.
+
+    If the stream terminates abnormally (aborted by the sender, oversized, or
+    the room disconnects mid-stream), :class:`StreamError` is raised instead of
+    a normal ``StopAsyncIteration``.
+    """
+
     def __init__(
         self,
         owned_info: proto_data_stream.OwnedByteStreamReader,
@@ -225,22 +299,59 @@ async def __anext__(self) -> bytes:
     def info(self) -> ByteStreamInfo:
         return self._info
 
+    def _unsubscribe(self) -> None:
+        """Drops the FFI queue subscription without ending the stream.
+
+        Events already delivered stay in the queue and remain readable; only
+        the global subscription, and the filter it runs against every FFI
+        event, goes away. Safe to call repeatedly — unsubscribing a queue that
+        is no longer registered is a no-op.
+
+        Only the subscription is released: the reader handle was consumed by
+        the read_incremental request in __init__, so there is nothing left to
+        dispose (and disposing it would drop a handle the FFI server no longer
+        owns).
+        """
+        FfiClient.instance.queue.unsubscribe(self._queue)
+
     def _close(self) -> None:
         if not self._closed:
             self._closed = True
-            FfiClient.instance.queue.unsubscribe(self._queue)
+            self._unsubscribe()
             if self._on_close is not None:
                 self._on_close()
 
+    def close(self) -> None:
+        """Explicitly close the reader and unsubscribe.
+
+        Call this on a reader you stop consuming before it ends. Closing a
+        reader that already reached end-of-stream is a no-op, and iterating a
+        closed reader raises ``StopAsyncIteration`` (or :class:`StreamError`
+        if the stream had already failed).
+        """
+        self._close()
+
+    async def aclose(self) -> None:
+        self.close()
+
     def _signal_disconnect(self) -> None:
         """Injects a synthetic EOS-with-error event so pending reads wake up
-        and raise StreamError when the room disconnects mid-stream."""
+        and raise StreamError when the room disconnects mid-stream.
+
+        Also drops the queue subscription, which would otherwise outlive the
+        room: no further events can arrive once the room is gone, and the
+        injected one is already queued. The reader is deliberately left open
+        so chunks that arrived before the disconnect are still delivered
+        before the StreamError, as they were before this was unsubscribed
+        here.
+        """
         if self._closed:
             return
         event = proto_ffi.FfiEvent()
         event.byte_stream_reader_event.reader_handle = self._reader_handle
         event.byte_stream_reader_event.eos.error.description = _DISCONNECT_ERROR
         self._queue.put_nowait(event)
+        self._unsubscribe()
 
 
 class BaseStreamWriter:

livekit-rtc/livekit/rtc/room.py

@@ -1196,10 +1196,18 @@ def on_close() -> None:
         byte_stream_handler(byte_reader, opened.participant_identity)
 
     def _error_stream_readers(self) -> None:
-        """Wakes up any pending stream reads with a StreamError on disconnect."""
-        for text_reader in self._text_stream_readers.values():
+        """Wakes up any pending stream reads with a StreamError on disconnect.
+
+        This also drops each reader's FFI queue subscription, which would
+        otherwise outlive the room: a reader unsubscribes itself only once a
+        read consumes its end-of-stream event, so one the application never
+        finished reading would stay subscribed for the life of the process.
+        The readers stay open, so a read can still drain whatever arrived
+        before the disconnect and then raise StreamError.
+        """
+        for text_reader in list(self._text_stream_readers.values()):
             text_reader._signal_disconnect()
-        for byte_reader in self._byte_stream_readers.values():
+        for byte_reader in list(self._byte_stream_readers.values()):
             byte_reader._signal_disconnect()
         self._text_stream_readers.clear()
         self._byte_stream_readers.clear()

tests/rtc/test_e2e_data_streams.py

@@ -26,6 +26,7 @@
 import pytest
 
 from livekit import api, rtc
+from livekit.rtc._ffi_client import FfiClient
 
 
 def skip_if_no_credentials() -> Any:
@@ -61,6 +62,26 @@ def pseudo_random_text(length: int, seed: int = 0x5EED) -> str:
     return "".join(rng.choices(string.ascii_lowercase, k=length))
 
 
+def reader_is_subscribed(reader: Any) -> bool:
+    """True while the reader's filtered subscriber is still on the FFI queue.
+
+    Checked by identity against the reader's own queue rather than by counting
+    subscribers, so unrelated churn (per-request callback subscriptions, the
+    room's own subscribers going away at disconnect) can't affect the result.
+    """
+    queue = reader._queue
+    return any(q is queue for q, _, _ in FfiClient.instance.queue._subscribers)
+
+
+async def wait_until_unsubscribed(reader: Any, timeout: float = 5.0) -> None:
+    deadline = asyncio.get_event_loop().time() + timeout
+    while asyncio.get_event_loop().time() < deadline:
+        if not reader_is_subscribed(reader):
+            return
+        await asyncio.sleep(0.02)
+    raise AssertionError("reader is still subscribed to the FFI event queue")
+
+
 async def connect_rooms(
     room_name: str,
     *,
@@ -441,6 +462,145 @@ async def read() -> None:
         await sender.disconnect()
 
 
+@pytest.mark.asyncio
+@skip_if_no_credentials()  # type: ignore[untyped-decorator]
+async def test_abandoned_reader_releases_subscription_on_close() -> None:
+    """A reader abandoned mid-iteration never consumes its terminal eos event,
+    so it stays subscribed to the FFI queue until close() is called."""
+    receiver, sender = await connect_rooms(unique_room_name("ds-abandon"))
+    try:
+        got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future()
+        abandoned = asyncio.Event()
+
+        def handler(reader: rtc.TextStreamReader, _identity: str) -> None:
+            got_reader.set_result(reader)
+
+            async def read() -> None:
+                async for _chunk in reader:
+                    break  # walk away without draining to end-of-stream
+                abandoned.set()
+
+            asyncio.create_task(read())
+
+        receiver.register_text_stream_handler("abandon-topic", handler)
+
+        await sender.local_participant.send_text(pseudo_random_text(50_000), topic="abandon-topic")
+        reader = await asyncio.wait_for(got_reader, timeout=10.0)
+        await asyncio.wait_for(abandoned.wait(), timeout=10.0)
+
+        assert reader_is_subscribed(reader)
+        reader.close()
+        await wait_until_unsubscribed(reader)
+
+        reader.close()  # idempotent
+    finally:
+        await asyncio.gather(receiver.disconnect(), sender.disconnect())
+
+
+@pytest.mark.asyncio
+@skip_if_no_credentials()  # type: ignore[untyped-decorator]
+async def test_never_read_reader_releases_subscription_on_close() -> None:
+    """A handler that never iterates its reader still leaves a subscription
+    behind, because readers subscribe eagerly on construction."""
+    receiver, sender = await connect_rooms(unique_room_name("ds-noread"))
+    writer = None
+    try:
+        got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future()
+
+        def handler(reader: rtc.TextStreamReader, _identity: str) -> None:
+            got_reader.set_result(reader)  # never read
+
+        receiver.register_text_stream_handler("noread-topic", handler)
+
+        writer = await sender.local_participant.stream_text(topic="noread-topic")
+        await writer.write("some data")
+
+        reader = await asyncio.wait_for(got_reader, timeout=10.0)
+        assert reader_is_subscribed(reader)
+
+        await reader.aclose()
+        await wait_until_unsubscribed(reader)
+    finally:
+        if writer is not None:
+            await writer.aclose()
+        await asyncio.gather(receiver.disconnect(), sender.disconnect())
+
+
+@pytest.mark.asyncio
+@skip_if_no_credentials()  # type: ignore[untyped-decorator]
+async def test_disconnect_unsubscribes_unread_reader() -> None:
+    """Disconnecting must drop the subscriptions of readers the application
+    never consumed, without the application closing them.
+
+    The reader is left open, so a read starting after the disconnect still
+    drains whatever arrived beforehand and then reports the failure — the
+    behaviour callers had before disconnect released the subscription.
+    """
+    receiver, sender = await connect_rooms(unique_room_name("ds-disc-unread"))
+    writer = None
+    try:
+        got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future()
+
+        def handler(reader: rtc.TextStreamReader, _identity: str) -> None:
+            got_reader.set_result(reader)  # never read
+
+        receiver.register_text_stream_handler("disc-unread-topic", handler)
+
+        writer = await sender.local_participant.stream_text(topic="disc-unread-topic")
+        await writer.write("buffered ")
+        await writer.write("data")
+
+        reader = await asyncio.wait_for(got_reader, timeout=10.0)
+        assert reader_is_subscribed(reader)
+        # let the chunks land in the reader's queue while nothing is reading
+        await asyncio.sleep(0.5)
+
+        await receiver.disconnect()
+        await wait_until_unsubscribed(reader)
+
+        # chunks received before the disconnect are still delivered, and the
+        # stream then terminates with StreamError rather than a clean EOF
+        chunks = []
+        with pytest.raises(rtc.StreamError):
+            async for chunk in reader:
+                chunks.append(chunk)
+        assert "".join(chunks) == "buffered data"
+    finally:
+        if writer is not None:
+            await writer.aclose()
+        await sender.disconnect()
+
+
+@pytest.mark.asyncio
+@skip_if_no_credentials()  # type: ignore[untyped-decorator]
+async def test_fully_read_reader_needs_no_close() -> None:
+    """Reading to end-of-stream unsubscribes on its own; close() afterwards is
+    a no-op."""
+    receiver, sender = await connect_rooms(unique_room_name("ds-drain"))
+    try:
+        got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future()
+        received: asyncio.Future[str] = asyncio.get_event_loop().create_future()
+
+        def handler(reader: rtc.TextStreamReader, _identity: str) -> None:
+            got_reader.set_result(reader)
+
+            async def read() -> None:
+                received.set_result(await reader.read_all())
+
+            asyncio.create_task(read())
+
+        receiver.register_text_stream_handler("drain-topic", handler)
+
+        await sender.local_participant.send_text("hello", topic="drain-topic")
+        assert await asyncio.wait_for(received, timeout=10.0) == "hello"
+
+        reader = await asyncio.wait_for(got_reader, timeout=5.0)
+        await wait_until_unsubscribed(reader)
+        reader.close()
+    finally:
+        await asyncio.gather(receiver.disconnect(), sender.disconnect())
+
+
 @pytest.mark.asyncio
 @skip_if_no_credentials()  # type: ignore[untyped-decorator]
 async def test_send_file_round_trip(tmp_path: Any) -> None: