fix(livekit): prevent memory leak when video_in_enabled is False
This commit is contained in:
@@ -539,11 +539,14 @@ class LiveKitTransportClient:
|
||||
elif track.kind == rtc.TrackKind.KIND_VIDEO:
|
||||
logger.info(f"Video track subscribed: {track.sid} from participant {participant.sid}")
|
||||
self._video_tracks[participant.sid] = track
|
||||
video_stream = rtc.VideoStream(track)
|
||||
self._task_manager.create_task(
|
||||
self._process_video_stream(video_stream, participant.sid),
|
||||
f"{self}::_process_video_stream",
|
||||
)
|
||||
# Only process video stream if video input is enabled to prevent
|
||||
# unbounded queue growth when there is no consumer for video frames.
|
||||
if self._params.video_in_enabled:
|
||||
video_stream = rtc.VideoStream(track)
|
||||
self._task_manager.create_task(
|
||||
self._process_video_stream(video_stream, participant.sid),
|
||||
f"{self}::_process_video_stream",
|
||||
)
|
||||
await self._callbacks.on_video_track_subscribed(participant.sid)
|
||||
|
||||
async def _async_on_track_unsubscribed(
|
||||
|
||||
133
tests/test_livekit_transport.py
Normal file
133
tests/test_livekit_transport.py
Normal file
@@ -0,0 +1,133 @@
|
||||
#
|
||||
# Copyright (c) 2024-2026, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""Tests for LiveKit transport video stream handling.
|
||||
|
||||
Regression tests for issue #3116: Memory leak when video_in_enabled=False
|
||||
but video tracks are subscribed.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from livekit import rtc
|
||||
|
||||
from pipecat.transports.livekit.transport import (
|
||||
LiveKitCallbacks,
|
||||
LiveKitParams,
|
||||
LiveKitTransportClient,
|
||||
)
|
||||
|
||||
|
||||
class TestLiveKitTransportClient(unittest.IsolatedAsyncioTestCase):
|
||||
"""Tests for LiveKitTransportClient video stream handling."""
|
||||
|
||||
def _create_client(self, video_in_enabled: bool) -> LiveKitTransportClient:
|
||||
"""Create a LiveKitTransportClient with the specified video_in_enabled setting."""
|
||||
params = LiveKitParams(video_in_enabled=video_in_enabled)
|
||||
|
||||
callbacks = LiveKitCallbacks(
|
||||
on_connected=AsyncMock(),
|
||||
on_disconnected=AsyncMock(),
|
||||
on_before_disconnect=AsyncMock(),
|
||||
on_participant_connected=AsyncMock(),
|
||||
on_participant_disconnected=AsyncMock(),
|
||||
on_audio_track_subscribed=AsyncMock(),
|
||||
on_audio_track_unsubscribed=AsyncMock(),
|
||||
on_video_track_subscribed=AsyncMock(),
|
||||
on_video_track_unsubscribed=AsyncMock(),
|
||||
on_data_received=AsyncMock(),
|
||||
on_first_participant_joined=AsyncMock(),
|
||||
)
|
||||
|
||||
client = LiveKitTransportClient(
|
||||
url="wss://test.livekit.cloud",
|
||||
token="test-token",
|
||||
room_name="test-room",
|
||||
params=params,
|
||||
callbacks=callbacks,
|
||||
transport_name="test-transport",
|
||||
)
|
||||
|
||||
# Mock the task manager
|
||||
client._task_manager = MagicMock()
|
||||
client._task_manager.create_task = MagicMock()
|
||||
|
||||
return client
|
||||
|
||||
def _create_mock_video_track(self) -> tuple:
|
||||
"""Create mock video track, publication, and participant."""
|
||||
mock_track = MagicMock()
|
||||
mock_track.kind = rtc.TrackKind.KIND_VIDEO
|
||||
mock_track.sid = "test-track-sid"
|
||||
|
||||
mock_publication = MagicMock()
|
||||
|
||||
mock_participant = MagicMock()
|
||||
mock_participant.sid = "test-participant-sid"
|
||||
|
||||
return mock_track, mock_publication, mock_participant
|
||||
|
||||
def _was_video_stream_task_created(self, client: LiveKitTransportClient) -> bool:
|
||||
"""Check if _process_video_stream task was created."""
|
||||
for call in client._task_manager.create_task.call_args_list:
|
||||
task_name = call[0][1] if len(call[0]) > 1 else call[1].get("name", "")
|
||||
if "_process_video_stream" in task_name:
|
||||
return True
|
||||
return False
|
||||
|
||||
async def test_video_stream_not_started_when_video_in_disabled(self):
|
||||
"""Test that _process_video_stream is NOT started when video_in_enabled=False.
|
||||
|
||||
This prevents unbounded queue growth when there is no consumer for video frames.
|
||||
Regression test for issue #3116.
|
||||
"""
|
||||
client = self._create_client(video_in_enabled=False)
|
||||
mock_track, mock_publication, mock_participant = self._create_mock_video_track()
|
||||
|
||||
# Call the track subscribed handler
|
||||
await client._async_on_track_subscribed(mock_track, mock_publication, mock_participant)
|
||||
|
||||
# Verify that create_task was NOT called for video stream processing
|
||||
self.assertFalse(
|
||||
self._was_video_stream_task_created(client),
|
||||
"Video stream processing should NOT be started when video_in_enabled=False",
|
||||
)
|
||||
|
||||
# Verify that the callback was still called
|
||||
client._callbacks.on_video_track_subscribed.assert_called_once_with(mock_participant.sid)
|
||||
|
||||
# Verify that the track was still added to _video_tracks
|
||||
self.assertIn(mock_participant.sid, client._video_tracks)
|
||||
|
||||
async def test_video_stream_started_when_video_in_enabled(self):
|
||||
"""Test that _process_video_stream IS started when video_in_enabled=True."""
|
||||
from unittest.mock import patch
|
||||
|
||||
client = self._create_client(video_in_enabled=True)
|
||||
mock_track, mock_publication, mock_participant = self._create_mock_video_track()
|
||||
|
||||
# Mock rtc.VideoStream to avoid needing a real LiveKit connection
|
||||
with patch("pipecat.transports.livekit.transport.rtc.VideoStream"):
|
||||
# Call the track subscribed handler
|
||||
await client._async_on_track_subscribed(mock_track, mock_publication, mock_participant)
|
||||
|
||||
# Verify that create_task WAS called for video stream processing
|
||||
self.assertTrue(
|
||||
self._was_video_stream_task_created(client),
|
||||
"Video stream processing SHOULD be started when video_in_enabled=True",
|
||||
)
|
||||
|
||||
# Verify that the callback was called
|
||||
client._callbacks.on_video_track_subscribed.assert_called_once_with(mock_participant.sid)
|
||||
|
||||
# Verify that the track was added to _video_tracks
|
||||
self.assertIn(mock_participant.sid, client._video_tracks)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user