fix: backport abandoned data stream writer drop from node-sdks a6f8a35

Author: 1egoman

livekit-rtc/livekit/rtc/data_stream.py

@@ -378,6 +378,21 @@ def __init__(
     def _provisional_timestamp(self) -> int:
         return int(datetime.datetime.now().timestamp() * 1000)
 
+    def _register_open(self) -> None:
+        """Records the freshly opened writer so that, if it is never closed,
+        the room can drop its native handle on disconnect."""
+        assert self._writer_handle is not None
+        self._local_participant._open_stream_writers.add(self._writer_handle)
+
+    def _unregister_open(self) -> None:
+        """Forgets the writer because a close request has been issued for it.
+
+        The close consumes the handle on the native side, so it must no longer
+        be dropped by the disconnect cleanup.
+        """
+        if self._writer_handle is not None:
+            self._local_participant._open_stream_writers.discard(self._writer_handle)
+
     async def _wait_for_callback(
         self, req: proto_ffi.FfiRequest, callback_field: str, response_field: str
     ) -> proto_ffi.FfiEvent:
@@ -403,6 +418,10 @@ async def aclose(
         if self._writer_handle is None:
             raise RuntimeError("Stream is not open")
         self._closed = True
+        # unregister before sending: the request consumes the handle natively,
+        # so it must not be dropped again by the disconnect cleanup even if the
+        # close itself reports an error
+        self._unregister_open()
         await self._send_close(reason=reason, attributes=attributes)
 
     async def _send_close(self, *, reason: str, attributes: Optional[Dict[str, str]]) -> None:
@@ -461,6 +480,7 @@ async def _send_header(self) -> None:
         cb = await self._wait_for_callback(req, "text_stream_open", "text_stream_open")
         owned = cb.text_stream_open.writer
         self._writer_handle = owned.handle.id
+        self._register_open()
         self._info = _text_stream_info_from_proto(owned.info)
         if self._info.size is None and self._total_size is not None:
             # total_size is not transmitted for text streams; keep it local
@@ -546,6 +566,7 @@ async def _send_header(self) -> None:
         cb = await self._wait_for_callback(req, "byte_stream_open", "byte_stream_open")
         owned = cb.byte_stream_open.writer
         self._writer_handle = owned.handle.id
+        self._register_open()
         self._info = _byte_stream_info_from_proto(owned.info)
 
     async def write(self, data: bytes) -> None:

livekit-rtc/livekit/rtc/participant.py

@@ -246,6 +246,14 @@ def __init__(
         self._room_queue = room_queue
         self._track_publications: dict[str, LocalTrackPublication] = {}
         self._rpc_handlers: Dict[str, RpcHandler] = {}
+        # Handle ids of data stream writers that have been opened but not yet
+        # closed. The FFI close request consumes the handle (take_handle), so
+        # an entry is removed as soon as a close is sent for it; whatever is
+        # left when the room disconnects is dropped by
+        # _dispose_open_stream_writers(). Only the ids are kept — wrapping them
+        # in FfiHandle would make GC drop a handle the close request already
+        # consumed.
+        self._open_stream_writers: set[int] = set()
 
     @property
     def track_publications(self) -> Mapping[str, LocalTrackPublication]:
@@ -254,6 +262,22 @@ def track_publications(self) -> Mapping[str, LocalTrackPublication]:
         """
         return self._track_publications
 
+    def _dispose_open_stream_writers(self) -> None:
+        """Drops the native writers of any data streams that were opened but
+        never closed — a writer the caller abandoned, or one left behind by a
+        failed write. Their handles would otherwise be held by the FFI server
+        for the lifetime of the process.
+
+        The room is already gone at this point, so the handles are dropped
+        directly rather than closed: there is no end-of-stream left to deliver.
+        """
+        for writer_handle in self._open_stream_writers:
+            try:
+                FfiHandle(writer_handle).dispose()
+            except Exception:
+                logger.exception("failed to dispose data stream writer handle")
+        self._open_stream_writers.clear()
+
     async def publish_data(
         self,
         payload: Union[bytes, str],

livekit-rtc/livekit/rtc/room.py

@@ -686,6 +686,7 @@ async def disconnect(
 
         await self._drain_rpc_invocation_tasks()
         self._error_stream_readers()
+        self._dispose_open_stream_writers()
 
         req = proto_ffi.FfiRequest()
         req.disconnect.room_handle = self._ffi_handle.handle  # type: ignore
@@ -735,6 +736,7 @@ async def _listen_task(self) -> None:
         # Clean up any pending RPC invocation tasks
         await self._drain_rpc_invocation_tasks()
         self._error_stream_readers()
+        self._dispose_open_stream_writers()
 
     def _on_rpc_method_invocation(self, rpc_invocation: RpcMethodInvocationEvent) -> None:
         if self._local_participant is None:
@@ -1195,6 +1197,16 @@ def on_close() -> None:
         self._byte_stream_readers[handle_id] = byte_reader
         byte_stream_handler(byte_reader, opened.participant_identity)
 
+    def _dispose_open_stream_writers(self) -> None:
+        """Drops the native writers of data streams left open at disconnect.
+
+        A writer releases its handle only when it is closed, so one the
+        application abandoned — or left behind after a failed write — would
+        otherwise be held by the FFI server for the life of the process.
+        """
+        if self._local_participant is not None:
+            self._local_participant._dispose_open_stream_writers()
+
     def _error_stream_readers(self) -> None:
         """Wakes up any pending stream reads with a StreamError on disconnect.
 

tests/rtc/test_e2e_data_streams.py

@@ -26,7 +26,7 @@
 import pytest
 
 from livekit import api, rtc
-from livekit.rtc._ffi_client import FfiClient
+from livekit.rtc._ffi_client import FfiClient, FfiHandle
 
 
 def skip_if_no_credentials() -> Any:
@@ -73,6 +73,21 @@ def reader_is_subscribed(reader: Any) -> bool:
     return any(q is queue for q, _, _ in FfiClient.instance.queue._subscribers)
 
 
+def ffi_handle_is_dropped(handle_id: int) -> bool:
+    """True when the FFI server no longer holds ``handle_id``.
+
+    Probes by trying to drop it: ``FfiHandle.dispose()`` asserts that the
+    native drop succeeded, so a handle that is already gone raises. Destructive
+    when the handle *is* still held — it gets dropped — so only call this once,
+    at the end of a test.
+    """
+    try:
+        FfiHandle(handle_id).dispose()
+    except AssertionError:
+        return True
+    return False
+
+
 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:
@@ -601,6 +616,49 @@ async def read() -> None:
         await asyncio.gather(receiver.disconnect(), sender.disconnect())
 
 
+@pytest.mark.asyncio
+@skip_if_no_credentials()  # type: ignore[untyped-decorator]
+async def test_writer_handle_untracked_after_close() -> None:
+    """Closing a writer consumes its handle natively, so it must drop out of
+    the open-writer set and not be dropped again at disconnect."""
+    receiver, sender = await connect_rooms(unique_room_name("ds-writer-close"))
+    try:
+        local = sender.local_participant
+        writer = await local.stream_text(topic="writer-close-topic")
+        handle = writer._writer_handle
+        assert handle is not None
+        assert handle in local._open_stream_writers
+
+        await writer.write("data")
+        await writer.aclose()
+
+        assert handle not in local._open_stream_writers
+    finally:
+        await asyncio.gather(receiver.disconnect(), sender.disconnect())
+
+
+@pytest.mark.asyncio
+@skip_if_no_credentials()  # type: ignore[untyped-decorator]
+async def test_abandoned_writer_disposed_on_disconnect() -> None:
+    """A writer opened and never closed holds its native writer in the FFI
+    handle table; disconnecting must drop it."""
+    receiver, sender = await connect_rooms(unique_room_name("ds-writer-abandon"))
+    try:
+        writer = await sender.local_participant.stream_text(topic="writer-abandon-topic")
+        handle = writer._writer_handle
+        assert handle is not None
+
+        await writer.write("data")  # abandoned: aclose() is never called
+
+        await sender.disconnect()
+
+        # asserted against the FFI handle table rather than the SDK's
+        # bookkeeping, so this fails on an actually-leaked native writer
+        assert ffi_handle_is_dropped(handle)
+    finally:
+        await receiver.disconnect()
+
+
 @pytest.mark.asyncio
 @skip_if_no_credentials()  # type: ignore[untyped-decorator]
 async def test_send_file_round_trip(tmp_path: Any) -> None: