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:
|
elif track.kind == rtc.TrackKind.KIND_VIDEO:
|
||||||
logger.info(f"Video track subscribed: {track.sid} from participant {participant.sid}")
|
logger.info(f"Video track subscribed: {track.sid} from participant {participant.sid}")
|
||||||
self._video_tracks[participant.sid] = track
|
self._video_tracks[participant.sid] = track
|
||||||
video_stream = rtc.VideoStream(track)
|
# Only process video stream if video input is enabled to prevent
|
||||||
self._task_manager.create_task(
|
# unbounded queue growth when there is no consumer for video frames.
|
||||||
self._process_video_stream(video_stream, participant.sid),
|
if self._params.video_in_enabled:
|
||||||
f"{self}::_process_video_stream",
|
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)
|
await self._callbacks.on_video_track_subscribed(participant.sid)
|
||||||
|
|
||||||
async def _async_on_track_unsubscribed(
|
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