diff --git a/livekit-rtc/livekit/rtc/room.py b/livekit-rtc/livekit/rtc/room.py index 8f89e8a5..38f0e3c2 100644 --- a/livekit-rtc/livekit/rtc/room.py +++ b/livekit-rtc/livekit/rtc/room.py @@ -21,9 +21,13 @@ from typing import Callable, Dict, Literal, Optional, cast, Mapping import warnings +from google.protobuf.descriptor import FieldDescriptor +from google.protobuf.message import Message + from .event_emitter import EventEmitter from ._ffi_client import FfiClient, FfiHandle from ._proto import ffi_pb2 as proto_ffi +from ._proto.handle_pb2 import FfiOwnedHandle from ._proto import participant_pb2 as proto_participant from ._proto import room_pb2 as proto_room from ._proto import stats_pb2 as proto_stats @@ -31,7 +35,7 @@ from ._proto.room_pb2 import ConnectionState, SimulateScenarioKind from ._proto.track_pb2 import TrackKind from ._proto.rpc_pb2 import RpcMethodInvocationEvent -from ._utils import BroadcastQueue +from ._utils import BroadcastQueue, Queue, task_done_logger from .e2ee import E2EEManager, E2EEOptions from .log import logger from .participant import ( @@ -167,6 +171,55 @@ def __init__(self, message: str): self.message = message +# how long a cancelled connect's cleanup waits for its room's end of stream after +# the disconnect, before it stops collecting the handles room events carry +_ABORTED_CONNECT_EOS_TIMEOUT = 5.0 + + +def _is_repeated(field: FieldDescriptor) -> bool: + # protobuf 7 drops `label` for `is_repeated`; protobuf 5 only has `label` + if hasattr(field, "is_repeated"): + return bool(field.is_repeated) + return bool(getattr(field, "label") == FieldDescriptor.LABEL_REPEATED) + + +def _dispose_owned_handles(msg: Message, *, keep: int = 0) -> None: + """Drop every FfiOwnedHandle in `msg`, at any depth, except `keep`.""" + for fd, value in msg.ListFields(): + message_type = fd.message_type + if fd.type != FieldDescriptor.TYPE_MESSAGE or message_type is None: + continue + if message_type.GetOptions().map_entry: + value_field = message_type.fields_by_name["value"] + if value_field.type != FieldDescriptor.TYPE_MESSAGE: + continue + items = list(value.values()) + elif _is_repeated(fd): + items = list(value) + else: + items = [value] + for item in items: + if item.DESCRIPTOR.full_name == FfiOwnedHandle.DESCRIPTOR.full_name: + if item.id != keep: + FfiHandle(item.id).dispose() + else: + _dispose_owned_handles(item, keep=keep) + + +async def _drain_room_events(queue: Queue[proto_ffi.FfiEvent], room_handle: int) -> None: + """Drop the handles in a room's events until its end of stream.""" + while True: + event = await queue.get() + queue.task_done() + if event.WhichOneof("message") != "room_event": + continue + if event.room_event.room_handle != room_handle: + continue + if event.room_event.HasField("eos"): + return + _dispose_owned_handles(event.room_event) + + class Room(EventEmitter[EventTypes]): def __init__( self, @@ -184,6 +237,7 @@ def __init__( self._room_queue = BroadcastQueue[proto_ffi.FfiEvent]() self._info = proto_room.RoomInfo() self._rpc_invocation_tasks: set[asyncio.Task] = set() + self._aborted_connect_tasks: set[asyncio.Task] = set() self._remote_participants: Dict[str, RemoteParticipant] = {} self._connection_state = ConnectionState.CONN_DISCONNECTED @@ -550,17 +604,35 @@ def on_participant_connected(participant): ) req.connect.options.rtc_config.ice_servers.extend(options.rtc_config.ice_servers) + # a cancelled earlier connect on this Room may still be joined server-side + # with the same identity. joining before it has left gets this attempt + # evicted as a duplicate, so wait for its cleanup first. + await self._wait_aborted_connects() + # subscribe before connecting so we don't miss any events self._ffi_queue = FfiClient.instance.queue.subscribe(self._loop) queue = FfiClient.instance.queue.subscribe() try: resp = FfiClient.instance.request(req) + except BaseException: + FfiClient.instance.queue.unsubscribe(queue) + raise + + try: cb: proto_ffi.FfiEvent = await queue.wait_for( lambda e: e.connect.async_id == resp.connect.async_id ) - finally: - FfiClient.instance.queue.unsubscribe(queue) + except BaseException: + # any exit here (cancellation, GeneratorExit when the coroutine is closed, + # KeyboardInterrupt) leaves the FFI server connecting and expecting a + # ReadyForRoomEvent once it answers. leaving that unanswered panics it, and + # the panic handler terminates the process, so close the room from a task + # that outlives this coroutine. + FfiClient.instance.queue.unsubscribe(self._ffi_queue) + self._close_aborted_connect(resp.connect.async_id, queue) + raise + FfiClient.instance.queue.unsubscribe(queue) if cb.connect.error: FfiClient.instance.queue.unsubscribe(self._ffi_queue) @@ -602,6 +674,89 @@ def on_participant_connected(participant): ready_req.ready_for_room_event.room_handle = self._ffi_handle.handle FfiClient.instance.request(ready_req) + def _close_aborted_connect(self, async_id: int, queue: Queue[proto_ffi.FfiEvent]) -> None: + """Close a room that connect() exited before it could own. + + Takes ownership of `queue`. The FFI server has no cancel path for an in-flight + connect, so the room has to be created and then disconnected. Without this the + room also stays joined server-side and reconnecting with the same identity + evicts the new session as a duplicate. + """ + + async def _close() -> None: + try: + cb: proto_ffi.FfiEvent = await queue.wait_for( + lambda e: e.connect.async_id == async_id + ) + if cb.connect.error: + return + + room_handle = FfiHandle(cb.connect.result.room.handle.id) + # nothing will own the participants and publications the connect + # returned, so drop their handles here instead of leaking them + _dispose_owned_handles(cb.connect.result, keep=room_handle.handle) + + ready_req = proto_ffi.FfiRequest() + ready_req.ready_for_room_event.room_handle = room_handle.handle + FfiClient.instance.request(ready_req) + + await self._request_disconnect( + room_handle.handle, DisconnectReason.CLIENT_INITIATED + ) + + # room events sent between the ready request and the disconnect + # (participant_connected, track_published, ...) carry owned handles + # too. drop them until the room's end of stream. + try: + await asyncio.wait_for( + _drain_room_events(queue, room_handle.handle), + _ABORTED_CONNECT_EOS_TIMEOUT, + ) + except asyncio.TimeoutError: + logger.warning( + "no end of stream from an aborted connect's room after %ss", + _ABORTED_CONNECT_EOS_TIMEOUT, + ) + room_handle.dispose() + finally: + FfiClient.instance.queue.unsubscribe(queue) + + try: + task = self._loop.create_task(_close()) + except RuntimeError: + # the loop is already closed: nothing can answer the FFI any more + FfiClient.instance.queue.unsubscribe(queue) + logger.warning("event loop closed before an aborted connect could be cleaned up") + return + self._aborted_connect_tasks.add(task) + task.add_done_callback(self._aborted_connect_tasks.discard) + # a failure here still ends in an FFI panic, so it must not be swallowed + task.add_done_callback(task_done_logger) + + async def _wait_aborted_connects(self) -> None: + """Wait for any room a cancelled connect() is still closing.""" + if not self._aborted_connect_tasks: + return + # shielded, because gather() cancels its children when it is cancelled. + # a caller who gives up on connect() or disconnect() would otherwise cancel + # the very cleanup that answers the FFI's wait, leaving it to time out and + # panic, which is the failure this path exists to prevent. + await asyncio.shield( + asyncio.gather(*tuple(self._aborted_connect_tasks), return_exceptions=True) + ) + + @staticmethod + async def _request_disconnect(room_handle: int, reason: DisconnectReason.ValueType) -> None: + req = proto_ffi.FfiRequest() + req.disconnect.room_handle = room_handle + req.disconnect.reason = reason + queue = FfiClient.instance.queue.subscribe() + try: + resp = FfiClient.instance.request(req) + await queue.wait_for(lambda e: e.disconnect.async_id == resp.disconnect.async_id) + finally: + FfiClient.instance.queue.unsubscribe(queue) + async def get_rtc_stats(self) -> RtcStats: if not self.isconnected(): raise RuntimeError("the room isn't connected") @@ -681,6 +836,10 @@ async def disconnect( self, *, reason: DisconnectReason.ValueType = DisconnectReason.CLIENT_INITIATED ) -> None: """Disconnects from the room.""" + # a cancelled connect may still be closing a room the FFI server opened. + # wait for it so disconnect() leaves nothing behind. + await self._wait_aborted_connects() + if not self.isconnected(): return @@ -688,15 +847,7 @@ async def disconnect( self._error_stream_readers() self._dispose_open_stream_writers() - req = proto_ffi.FfiRequest() - req.disconnect.room_handle = self._ffi_handle.handle # type: ignore - req.disconnect.reason = reason - queue = FfiClient.instance.queue.subscribe() - try: - resp = FfiClient.instance.request(req) - await queue.wait_for(lambda e: e.disconnect.async_id == resp.disconnect.async_id) - finally: - FfiClient.instance.queue.unsubscribe(queue) + await self._request_disconnect(self._ffi_handle.handle, reason) # type: ignore await self._task FfiClient.instance.queue.unsubscribe(self._ffi_queue) diff --git a/livekit-rtc/tests/test_connect_cancellation.py b/livekit-rtc/tests/test_connect_cancellation.py new file mode 100644 index 00000000..5564d702 --- /dev/null +++ b/livekit-rtc/tests/test_connect_cancellation.py @@ -0,0 +1,278 @@ +# 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. + +import asyncio + +import pytest + +from livekit import rtc +from livekit.rtc import room as room_mod +from livekit.rtc._ffi_client import FfiClient +from livekit.rtc._proto import ffi_pb2 as proto_ffi +from utils import wait_until # type: ignore[import-not-found] + +CONNECT_ASYNC_ID = 101 +DISCONNECT_ASYNC_ID = 202 +ROOM_HANDLE = 7 + + +class _FakeFfi: + """Records every FfiRequest and dropped handle, and answers the requests the + cancel path waits on the way the FFI server does.""" + + def __init__(self) -> None: + self.requests: list[proto_ffi.FfiRequest] = [] + self.disposed: list[int] = [] + # room events the server emits once the SDK says it is ready, by room handle + self.events_after_ready: dict[int, list[proto_ffi.FfiEvent]] = {} + self._next_connect_id = CONNECT_ASYNC_ID + self._next_disconnect_id = DISCONNECT_ASYNC_ID + + def kinds(self) -> list[str]: + return [req.WhichOneof("message") for req in self.requests] + + def request(self, req: proto_ffi.FfiRequest) -> proto_ffi.FfiResponse: + self.requests.append(req) + resp = proto_ffi.FfiResponse() + which = req.WhichOneof("message") + if which == "connect": + # the connect callback is delivered by the test, not here + resp.connect.async_id = self._next_connect_id + self._next_connect_id += 1 + elif which == "ready_for_room_event": + for event in self.events_after_ready.pop(req.ready_for_room_event.room_handle, []): + FfiClient.instance.queue.put(event) + elif which == "disconnect": + resp.disconnect.async_id = self._next_disconnect_id + self._next_disconnect_id += 1 + event = proto_ffi.FfiEvent() + event.disconnect.async_id = resp.disconnect.async_id + FfiClient.instance.queue.put(event) + eos = proto_ffi.FfiEvent() + eos.room_event.room_handle = req.disconnect.room_handle + eos.room_event.eos.SetInParent() + FfiClient.instance.queue.put(eos) + return resp + + +def _install_fake_ffi(monkeypatch: pytest.MonkeyPatch) -> _FakeFfi: + fake = _FakeFfi() + + class _FakeHandle: + """Stand-in for FfiHandle so a made-up handle id is never dropped natively.""" + + def __init__(self, handle: int) -> None: + self.handle = handle + + def dispose(self) -> None: + fake.disposed.append(self.handle) + + monkeypatch.setattr(FfiClient.instance, "request", fake.request) + monkeypatch.setattr(room_mod, "FfiHandle", _FakeHandle) + return fake + + +def _deliver_connect_callback( + async_id: int = CONNECT_ASYNC_ID, room_handle: int = ROOM_HANDLE +) -> proto_ffi.FfiEvent: + event = proto_ffi.FfiEvent() + event.connect.async_id = async_id + event.connect.result.room.handle.id = room_handle + FfiClient.instance.queue.put(event) + return event + + +async def _cancelled_connect(room: rtc.Room, fake: _FakeFfi) -> None: + sent = len(fake.requests) + task = asyncio.create_task(room.connect("ws://localhost:7880", "token")) + await wait_until(lambda: len(fake.requests) > sent, message="connect request never issued") + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + +async def test_cancelled_connect_answers_ready_and_closes_the_room( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake = _install_fake_ffi(monkeypatch) + requests = fake.requests + subscribers_before = len(FfiClient.instance.queue._subscribers) + + room = rtc.Room() + await _cancelled_connect(room, fake) + + # the FFI server does not cancel an in-flight connect: it answers, then waits for + # ReadyForRoomEvent. an unanswered wait panics it and the panic kills the process. + _deliver_connect_callback() + await room.disconnect() + + assert [req.WhichOneof("message") for req in requests] == [ + "connect", + "ready_for_room_event", + "disconnect", + ] + assert requests[1].ready_for_room_event.room_handle == ROOM_HANDLE + assert requests[2].disconnect.room_handle == ROOM_HANDLE + assert fake.disposed == [ROOM_HANDLE] + assert len(FfiClient.instance.queue._subscribers) == subscribers_before + + +async def test_cancelled_connect_leaves_no_pending_work_when_the_server_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + fake = _install_fake_ffi(monkeypatch) + subscribers_before = len(FfiClient.instance.queue._subscribers) + + room = rtc.Room() + await _cancelled_connect(room, fake) + + event = proto_ffi.FfiEvent() + event.connect.async_id = CONNECT_ASYNC_ID + event.connect.error = "could not connect" + FfiClient.instance.queue.put(event) + await room.disconnect() + + # there is no room to close, so nothing follows the connect + assert fake.kinds() == ["connect"] + assert fake.disposed == [] + assert len(FfiClient.instance.queue._subscribers) == subscribers_before + + +async def test_a_cancelled_disconnect_still_lets_the_cleanup_finish( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Giving up on disconnect() must not cancel the cleanup it is waiting on. + + gather() cancels its children when it is cancelled, so a caller who bounds + disconnect() with a timeout, or abandons it on shutdown, would cancel the + task that answers the FFI's wait for ReadyForRoomEvent. The wait then times + out, the FFI panics, and the panic handler kills the process: the exact + failure the rest of this file is about, reintroduced one layer up. + """ + fake = _install_fake_ffi(monkeypatch) + subscribers_before = len(FfiClient.instance.queue._subscribers) + + room = rtc.Room() + await _cancelled_connect(room, fake) + + # the cleanup is now parked on the connect callback, which has not arrived + closing = asyncio.create_task(room.disconnect()) + await asyncio.sleep(0) + closing.cancel() + with pytest.raises(asyncio.CancelledError): + await closing + + _deliver_connect_callback() + await asyncio.gather(*room._aborted_connect_tasks) + + assert fake.kinds() == ["connect", "ready_for_room_event", "disconnect"] + assert fake.requests[1].ready_for_room_event.room_handle == ROOM_HANDLE + assert fake.requests[2].disconnect.room_handle == ROOM_HANDLE + assert len(FfiClient.instance.queue._subscribers) == subscribers_before + + +async def test_a_retry_on_the_same_room_waits_for_the_cancelled_attempt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A retry loop like `wait_for(room.connect(...), 0.2)` with no disconnect(). + + If the retry joined while the cancelled attempt was still in flight, the + server would evict whichever session came second as DUPLICATE_IDENTITY and + the cleanup would then close the other, leaving the caller a dead room. + """ + fake = _install_fake_ffi(monkeypatch) + subscribers_before = len(FfiClient.instance.queue._subscribers) + + room = rtc.Room() + await _cancelled_connect(room, fake) + + retry = asyncio.create_task(room.connect("ws://localhost:7880", "token")) + for _ in range(5): + await asyncio.sleep(0) + assert fake.kinds() == ["connect"], "the retry joined before the first attempt left" + + # the first attempt's join completes, and its cleanup closes it + _deliver_connect_callback(CONNECT_ASYNC_ID, ROOM_HANDLE) + await wait_until(lambda: len(fake.requests) == 4, message="the retry never connected") + _deliver_connect_callback(CONNECT_ASYNC_ID + 1, ROOM_HANDLE + 1) + await retry + + assert fake.kinds() == [ + "connect", + "ready_for_room_event", + "disconnect", + "connect", + "ready_for_room_event", + ] + assert fake.requests[2].disconnect.room_handle == ROOM_HANDLE + assert fake.requests[4].ready_for_room_event.room_handle == ROOM_HANDLE + 1 + assert room.isconnected() + + await room.disconnect() + assert fake.requests[-1].disconnect.room_handle == ROOM_HANDLE + 1 + assert len(FfiClient.instance.queue._subscribers) == subscribers_before + + +async def test_a_cancelled_connect_into_a_populated_room_drops_every_handle( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Nothing owns the participants and publications a cancelled connect gets back, + nor the ones in room events sent before the disconnect, so the cleanup has to + drop their handles or a worker that keeps timing out connects leaks them.""" + fake = _install_fake_ffi(monkeypatch) + + room = rtc.Room() + await _cancelled_connect(room, fake) + + joined = proto_ffi.FfiEvent() + joined.room_event.room_handle = ROOM_HANDLE + joined.room_event.participant_connected.info.handle.id = 14 + fake.events_after_ready[ROOM_HANDLE] = [joined] + + event = proto_ffi.FfiEvent() + event.connect.async_id = CONNECT_ASYNC_ID + result = event.connect.result + result.room.handle.id = ROOM_HANDLE + result.local_participant.handle.id = 11 + remote = result.participants.add() + remote.participant.handle.id = 12 + remote.publications.add().handle.id = 13 + FfiClient.instance.queue.put(event) + await room.disconnect() + + assert fake.kinds() == ["connect", "ready_for_room_event", "disconnect"] + assert sorted(fake.disposed) == [ROOM_HANDLE, 11, 12, 13, 14] + # the room goes last, after everything it owned + assert fake.disposed[-1] == ROOM_HANDLE + + +async def test_an_exit_other_than_cancellation_still_closes_the_room( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """GeneratorExit (a coroutine closed during shutdown) or KeyboardInterrupt at the + connect await must hand the room off too, or the FFI is left waiting.""" + fake = _install_fake_ffi(monkeypatch) + subscribers_before = len(FfiClient.instance.queue._subscribers) + + room = rtc.Room() + coro = room.connect("ws://localhost:7880", "token") + coro.send(None) # runs up to the await on the connect callback + assert fake.kinds() == ["connect"] + coro.close() # raises GeneratorExit at that await + + _deliver_connect_callback() + await room.disconnect() + + assert fake.kinds() == ["connect", "ready_for_room_event", "disconnect"] + assert len(FfiClient.instance.queue._subscribers) == subscribers_before