Add zero-copy video frames and surface dropped frames

livekit-rtc/livekit/rtc/_utils.py

@@ -17,7 +17,7 @@
 from collections import deque
 import ctypes
 import random
-from typing import Any, Callable, Generator, Generic, List, TypeVar, Union
+from typing import Any, Callable, Generator, Generic, List, Optional, TypeVar, Union
 
 logger = logging.getLogger("livekit")
 
@@ -76,14 +76,24 @@ def get_address(data: Union[bytes, bytearray, memoryview]) -> int:
 
 
 class RingQueue(Generic[T]):
-    def __init__(self, capacity: int = 0) -> None:
+    def __init__(self, capacity: int = 0, on_discard: Optional[Callable[[T], None]] = None) -> None:
+        """
+        Args:
+            capacity: Maximum queued items, or 0 for unbounded.
+            on_discard: Called with the oldest item when a put evicts it.
+                Lets an owner of out-of-band resources (an SDK buffer, say)
+                release them instead of waiting for the collector.
+        """
         self._capacity = capacity
         self._queue: deque[T] = deque()
         self._event = asyncio.Event()
+        self._on_discard = on_discard
 
     def put(self, item: T) -> None:
         if self._capacity > 0 and len(self._queue) == self._capacity:
-            self._queue.popleft()
+            discarded = self._queue.popleft()
+            if self._on_discard is not None:
+                self._on_discard(discarded)
         self._queue.append(item)
         self._event.set()
 
@@ -93,6 +103,10 @@ async def get(self) -> T:
         self._event.clear()
         return self._queue.popleft()
 
+    def get_nowait(self) -> Optional[T]:
+        """Returns the next item, or None when the queue is empty."""
+        return self._queue.popleft() if self._queue else None
+
 
 class Queue(asyncio.Queue[T]):
     """asyncio.Queue with utility functions."""

livekit-rtc/livekit/rtc/video_frame.py

@@ -50,6 +50,11 @@ def __init__(
         """
         data = _ensure_compatible_buffer(data)
 
+        # Set for a borrowed frame only; see `_borrow_owned_info`. An owned
+        # frame holds its own copy of the pixels and has nothing to release.
+        self._owner: Optional[FfiHandle] = None
+        self._closed = False
+
         self._width = width
         self._height = height
         self._type = type
@@ -92,11 +97,79 @@ def data(self) -> memoryview:
 
         Returns:
             memoryview: The raw pixel data of the video frame as a memoryview object.
+
+        Raises:
+            RuntimeError: If this is a borrowed frame that has already been
+                closed. See :meth:`close`.
         """
+        if self._closed:
+            raise RuntimeError(
+                "the video frame has been closed; its pixel data belonged to the SDK "
+                "and has been released. Copy what you need before the frame closes."
+            )
         return memoryview(self._data)
 
+    @property
+    def is_borrowed(self) -> bool:
+        """
+        Whether this frame reads the SDK's own pixel buffer rather than a copy.
+
+        A borrowed frame's :attr:`data` is only valid until the frame is closed;
+        see :meth:`close`. Frames are owned unless the stream was created with
+        ``zero_copy=True``.
+        """
+        return self._owner is not None
+
+    def close(self) -> None:
+        """
+        Releases a borrowed frame's pixel buffer back to the SDK.
+
+        No-op for an owned frame, and safe to call more than once.
+
+        After this returns, :attr:`data` raises. **Anything else still pointing
+        at the buffer -- a ``memoryview`` taken from :attr:`data`, a ``numpy``
+        array built with ``np.frombuffer`` -- points at freed memory and must
+        not be read.**
+
+        There is deliberately no check for that. CPython reports an outstanding
+        export only against the exact buffer object it was taken from, and a
+        view derived from :attr:`data` re-exports the underlying array instead,
+        so a check here would catch almost nothing while reading as a
+        guarantee. Use ``np.array(frame.data)`` rather than
+        ``np.frombuffer(frame.data)`` when the result has to outlive the frame.
+
+        :class:`VideoStream` closes the previous borrowed frame before yielding
+        the next one, so a loop that does not hold onto frames needs no explicit
+        call. Use the frame as a context manager to release it sooner.
+        """
+        if self._closed or self._owner is None:
+            return
+
+        self._closed = True
+        if isinstance(self._data, memoryview):
+            # Invalidates the frame's own view, so a later read of a stale
+            # `frame.data` fails loudly instead of reading freed memory.
+            self._data.release()
+        self._owner.dispose()
+        self._owner = None
+
+    def __enter__(self) -> "VideoFrame":
+        return self
+
+    def __exit__(self, *_: Any) -> None:
+        self.close()
+
+    def __del__(self) -> None:
+        # Backstop for a borrowed frame the application never closed and the
+        # stream never superseded. Best effort: __del__ must not raise.
+        try:
+            self.close()
+        except Exception:
+            pass
+
     @staticmethod
     def _from_owned_info(owned_info: proto_video.OwnedVideoBuffer) -> "VideoFrame":
+        """Copies the pixels out of the SDK buffer and releases the handle."""
         info = owned_info.info
         data_len = _get_plane_length(info.type, info.width, info.height)
         cdata = (ctypes.c_uint8 * data_len).from_address(info.data_ptr)
@@ -110,6 +183,30 @@ def _from_owned_info(owned_info: proto_video.OwnedVideoBuffer) -> "VideoFrame":
         FfiHandle(owned_info.handle.id).dispose()
         return frame
 
+    @staticmethod
+    def _borrow_owned_info(owned_info: proto_video.OwnedVideoBuffer) -> "VideoFrame":
+        """
+        Wraps the SDK's buffer in place, without copying it.
+
+        The returned frame owns the FFI handle and keeps the buffer alive until
+        the frame is closed. This is what makes the frame's lifetime matter:
+        an owned frame is independent of the SDK, a borrowed one is not.
+        """
+        info = owned_info.info
+        data_len = _get_plane_length(info.type, info.width, info.height)
+        cdata = (ctypes.c_uint8 * data_len).from_address(info.data_ptr)
+        frame = VideoFrame(
+            width=info.width,
+            height=info.height,
+            type=info.type,
+            # `cast` is not optional: a view of a ctypes array carries format
+            # "<B", which memoryview refuses to index. Casting to "B" gives the
+            # same plain bytes view the copying path produces from a bytearray.
+            data=memoryview(cdata).cast("B"),
+        )
+        frame._owner = FfiHandle(owned_info.handle.id)
+        return frame
+
     def _proto_info(self) -> proto_video.VideoBufferInfo:
         info = proto_video.VideoBufferInfo()
         addr = get_address(self.data)

livekit-rtc/livekit/rtc/video_stream.py

@@ -35,26 +35,77 @@ class VideoFrameEvent:
     rotation: proto_video_frame.VideoRotation
     metadata: proto_video_frame.FrameMetadata | None = None
 
+    dropped_frames: int = 0
+    """
+    Frames this stream decoded but never delivered, counted since the stream
+    was created.
+
+    A frame is dropped when the SDK's internal queue overflows because this
+    application did not read frames fast enough. No WebRTC stat reports the
+    loss -- ``framesDecoded`` counts these frames and ``framesDropped`` does
+    not -- so this counter is the only way to observe it.
+
+    The value is cumulative and monotonic: compare it against the previous
+    event's to see how many frames were lost in between. It stays at zero when
+    the SDK in use predates the counter.
+    """
+
+
+def _close_discarded(event: VideoFrameEvent | None) -> None:
+    if event is not None:
+        event.frame.close()
+
+
+def _dropped_frames(frame_received: proto_video_frame.VideoFrameReceived) -> int:
+    """
+    Reads the dropped-frame counter, tolerating an SDK that predates it.
+
+    The field arrived with the native counter; an older `livekit_ffi` does not
+    define it, and reporting zero is the honest answer there -- we have no
+    information either way.
+
+    `getattr` rather than attribute access so this also builds against the
+    generated stubs from before the field existed. Once the vendored
+    `rust-sdks` is bumped past it, this can become a plain read.
+    """
+    return int(getattr(frame_received, "dropped_frames", 0))
+
 
 class VideoStream:
-    """VideoStream is a stream of video frames received from a RemoteTrack."""
+    """VideoStream is a stream of video frames received from a RemoteTrack.
+
+    With ``zero_copy=True`` each frame reads the SDK's own pixel buffer instead
+    of a copy of it, which removes a full frame-sized copy per frame. The
+    buffer belongs to the SDK and is only valid while the frame is open; see
+    :meth:`VideoFrame.close` and the ``zero_copy`` argument.
+    """
 
     def __init__(
         self,
         track: Track,
         loop: Optional[asyncio.AbstractEventLoop] = None,
         capacity: int = 0,
         format: Optional[proto_video_frame.VideoBufferType.ValueType] = None,
+        zero_copy: bool = False,
         **kwargs: Any,
     ) -> None:
         self._loop = loop or asyncio.get_event_loop()
+        self._zero_copy = zero_copy
+        # The most recent borrowed frame handed to the application. Closed when
+        # the next one is requested, so a plain `async for` never leaks a
+        # buffer even if the application does not close frames itself.
+        self._borrowed: VideoFrame | None = None
         # 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,
             filter_fn=lambda e: e.WhichOneof("message") == "video_stream_event",
         )
-        self._queue: RingQueue[VideoFrameEvent | None] = RingQueue(capacity)
+        # A borrowed frame evicted by a bounded queue has to be closed, or its
+        # SDK buffer stays out until the collector gets to it.
+        self._queue: RingQueue[VideoFrameEvent | None] = RingQueue(
+            capacity, on_discard=_close_discarded if zero_copy else None
+        )
         self._track: Track | None = track
         self._format = format
         self._capacity = capacity
@@ -82,13 +133,15 @@ def from_participant(
         loop: Optional[asyncio.AbstractEventLoop] = None,
         format: Optional[proto_video_frame.VideoBufferType.ValueType] = None,
         capacity: int = 0,
+        zero_copy: bool = False,
     ) -> VideoStream:
         return VideoStream(
             participant=participant,
             track_source=track_source,
             loop=loop,
             capacity=capacity,
             format=format,
+            zero_copy=zero_copy,
             track=None,  # type: ignore
         )
 
@@ -100,16 +153,24 @@ def from_track(
         loop: Optional[asyncio.AbstractEventLoop] = None,
         format: Optional[proto_video_frame.VideoBufferType.ValueType] = None,
         capacity: int = 0,
+        zero_copy: bool = False,
     ) -> VideoStream:
         return VideoStream(
             track=track,
             loop=loop,
             capacity=capacity,
             format=format,
+            zero_copy=zero_copy,
         )
 
     def __del__(self) -> None:
-        FfiClient.instance.queue.unsubscribe(self._ffi_queue)
+        # Tolerate a half-built stream: if __init__ raised before subscribing
+        # (a failed FFI request, say), there is nothing to unsubscribe, and
+        # raising from __del__ only produces an unraisable-exception warning
+        # that hides the original error.
+        ffi_queue = getattr(self, "_ffi_queue", None)
+        if ffi_queue is not None:
+            FfiClient.instance.queue.unsubscribe(ffi_queue)
 
     def _create_owned_stream(self) -> Any:
         assert self._track is not None
@@ -145,7 +206,10 @@ async def _run(self) -> None:
             if video_event.HasField("frame_received"):
                 frame_received = video_event.frame_received
                 owned_buffer_info = frame_received.buffer
-                frame = VideoFrame._from_owned_info(owned_buffer_info)
+                if self._zero_copy:
+                    frame = VideoFrame._borrow_owned_info(owned_buffer_info)
+                else:
+                    frame = VideoFrame._from_owned_info(owned_buffer_info)
                 metadata: proto_video_frame.FrameMetadata | None = None
                 if frame_received.HasField("metadata"):
                     metadata = proto_video_frame.FrameMetadata()
@@ -156,6 +220,7 @@ async def _run(self) -> None:
                     timestamp_us=frame_received.timestamp_us,
                     rotation=frame_received.rotation,
                     metadata=metadata,
+                    dropped_frames=_dropped_frames(frame_received),
                 )
 
                 self._queue.put(event)
@@ -167,6 +232,16 @@ async def _run(self) -> None:
     async def aclose(self) -> None:
         self._ffi_handle.dispose()
         await self._task
+        self._release_borrowed()
+        # Frames still queued hold SDK buffers; release them with the stream.
+        while (queued := self._queue.get_nowait()) is not None:
+            queued.frame.close()
+
+    def _release_borrowed(self) -> None:
+        """Closes the borrowed frame handed out last, if there is one."""
+        if self._borrowed is not None:
+            self._borrowed.close()
+            self._borrowed = None
 
     def _is_event(self, e: proto_ffi.FfiEvent) -> bool:
         return e.video_stream_event.stream_handle == self._ffi_handle.handle
@@ -175,11 +250,19 @@ def __aiter__(self) -> AsyncIterator[VideoFrameEvent]:
         return self
 
     async def __anext__(self) -> VideoFrameEvent:
+        # Asking for the next frame ends the previous one's borrow. This is
+        # what keeps a plain `async for` from accumulating SDK buffers when
+        # the application never closes frames itself.
+        self._release_borrowed()
+
         if self._task.done():
             raise StopAsyncIteration
 
         item = await self._queue.get()
         if item is None:
             raise StopAsyncIteration
 
+        if self._zero_copy:
+            self._borrowed = item.frame
+
         return item