fix: manage WebRTC pipeline shutdown
This commit is contained in:
@@ -41,7 +41,10 @@ async def lifespan(_app: FastAPI):
|
||||
await sync_interface_definitions()
|
||||
await sync_default_tools()
|
||||
await recover_interrupted_documents()
|
||||
yield
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
await voice_webrtc.shutdown_active_sessions()
|
||||
|
||||
|
||||
app = FastAPI(title="AI Video Assistant 平台 - 后端", lifespan=lifespan)
|
||||
|
||||
@@ -11,6 +11,8 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import Coroutine
|
||||
from typing import Any
|
||||
|
||||
from db.session import SessionLocal
|
||||
from fastapi import APIRouter, Body, Depends, Request, WebSocket
|
||||
@@ -27,6 +29,116 @@ from services.webrtc_ice import aiortc_ice_servers, client_ice_servers
|
||||
|
||||
router = APIRouter(tags=["voice"])
|
||||
_http_peers: dict[str, object] = {}
|
||||
_active_connections: set[object] = set()
|
||||
_pipeline_tasks: set[asyncio.Task[None]] = set()
|
||||
_connection_tasks: dict[object, asyncio.Task[None]] = {}
|
||||
PIPELINE_CLOSE_GRACE_SECONDS = 10.0
|
||||
|
||||
|
||||
def _consume_pipeline_result(task: asyncio.Task[None], connection: object) -> None:
|
||||
"""Retrieve background exceptions and release the task's strong reference."""
|
||||
_pipeline_tasks.discard(task)
|
||||
if _connection_tasks.get(connection) is task:
|
||||
_connection_tasks.pop(connection, None)
|
||||
try:
|
||||
error = task.exception()
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"WebRTC pipeline 已取消: task={task.get_name()}")
|
||||
return
|
||||
if error is not None:
|
||||
logger.opt(exception=error).error(
|
||||
f"WebRTC pipeline 异常结束: task={task.get_name()}"
|
||||
)
|
||||
|
||||
|
||||
def _start_pipeline_task(
|
||||
connection: object,
|
||||
coroutine: Coroutine[Any, Any, None],
|
||||
) -> asyncio.Task[None]:
|
||||
"""Start one strongly referenced pipeline task for a WebRTC connection."""
|
||||
connection_id = str(getattr(connection, "pc_id", "unknown"))
|
||||
task = asyncio.create_task(
|
||||
coroutine,
|
||||
name=f"webrtc-pipeline:{connection_id}",
|
||||
)
|
||||
_active_connections.add(connection)
|
||||
_pipeline_tasks.add(task)
|
||||
_connection_tasks[connection] = task
|
||||
task.add_done_callback(
|
||||
lambda completed, connection=connection: _consume_pipeline_result(
|
||||
completed,
|
||||
connection,
|
||||
)
|
||||
)
|
||||
return task
|
||||
|
||||
|
||||
async def _wait_for_pipeline_close(
|
||||
task: asyncio.Task[None] | None,
|
||||
*,
|
||||
connection_id: str,
|
||||
) -> None:
|
||||
"""Let transport disconnect finish normally, then cancel a stuck pipeline."""
|
||||
if task is None:
|
||||
return
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
asyncio.shield(task),
|
||||
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
||||
)
|
||||
return
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
f"WebRTC pipeline 关闭超过 {PIPELINE_CLOSE_GRACE_SECONDS:g} 秒,"
|
||||
f"执行取消: pc_id={connection_id}"
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception:
|
||||
# The done callback owns exception reporting and retrieval.
|
||||
return
|
||||
|
||||
task.cancel()
|
||||
done, _pending = await asyncio.wait(
|
||||
{task},
|
||||
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
||||
)
|
||||
if not done:
|
||||
logger.error(f"WebRTC pipeline 取消后仍未退出: pc_id={connection_id}")
|
||||
|
||||
|
||||
async def shutdown_active_sessions() -> None:
|
||||
"""Close active peers and drain every managed pipeline during app shutdown."""
|
||||
connections = list(_active_connections)
|
||||
if connections:
|
||||
await asyncio.gather(
|
||||
*(connection.disconnect() for connection in connections),
|
||||
return_exceptions=True,
|
||||
)
|
||||
_active_connections.difference_update(connections)
|
||||
_http_peers.clear()
|
||||
|
||||
tasks = list(_pipeline_tasks)
|
||||
if not tasks:
|
||||
return
|
||||
done, pending = await asyncio.wait(
|
||||
tasks,
|
||||
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
if pending:
|
||||
cancelled, stuck = await asyncio.wait(
|
||||
pending,
|
||||
timeout=PIPELINE_CLOSE_GRACE_SECONDS,
|
||||
)
|
||||
if stuck:
|
||||
logger.error(f"应用关闭时仍有 {len(stuck)} 个 WebRTC pipeline 未退出")
|
||||
else:
|
||||
cancelled = set()
|
||||
logger.info(
|
||||
f"WebRTC 会话清理完成: normal={len(done)} cancelled={len(cancelled)}"
|
||||
)
|
||||
|
||||
|
||||
@router.get("/api/webrtc/ice-servers", dependencies=[Depends(require_admin)])
|
||||
@@ -155,7 +267,12 @@ async def _handle_offer_payload(payload, peers):
|
||||
restart_pc = bool(payload.get("restart_pc"))
|
||||
if restart_pc and pc_id and pc_id in peers:
|
||||
old_pc = peers.pop(pc_id)
|
||||
old_task = _connection_tasks.get(old_pc)
|
||||
await old_pc.disconnect()
|
||||
await _wait_for_pipeline_close(
|
||||
old_task,
|
||||
connection_id=str(getattr(old_pc, "pc_id", pc_id)),
|
||||
)
|
||||
|
||||
if pc_id and pc_id in peers:
|
||||
pc = peers[pc_id]
|
||||
@@ -187,25 +304,31 @@ async def _handle_offer_payload(payload, peers):
|
||||
await pc.initialize(sdp=offer.sdp, type=offer.type)
|
||||
peers[pc.pc_id] = pc
|
||||
|
||||
@pc.event_handler("closed")
|
||||
async def _on_closed(conn: SmallWebRTCConnection):
|
||||
peers.pop(conn.pc_id, None)
|
||||
|
||||
# 后台跑管线:WebRTC transport + 解析出的运行时配置
|
||||
transport = build_webrtc_transport(
|
||||
pc,
|
||||
video_in_enabled=vision_enabled,
|
||||
)
|
||||
asyncio.create_task(
|
||||
pipeline_task = _start_pipeline_task(
|
||||
pc,
|
||||
run_pipeline(
|
||||
transport,
|
||||
cfg,
|
||||
vision_enabled=vision_enabled,
|
||||
assistant_id=offer.assistant_id,
|
||||
channel="webrtc",
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
@pc.event_handler("closed")
|
||||
async def _on_closed(conn: SmallWebRTCConnection):
|
||||
peers.pop(conn.pc_id, None)
|
||||
_active_connections.discard(conn)
|
||||
await _wait_for_pipeline_close(
|
||||
pipeline_task,
|
||||
connection_id=conn.pc_id,
|
||||
)
|
||||
|
||||
answer = pc.get_answer()
|
||||
return {
|
||||
"pc_id": answer["pc_id"],
|
||||
|
||||
@@ -169,6 +169,21 @@ class ConversationRecorder:
|
||||
logger.error(f"保存对话文本失败,不影响本次通话: {exc}")
|
||||
|
||||
async def finish(self, *, status: str = "completed") -> None:
|
||||
"""Finish the session even if the owning pipeline is being cancelled."""
|
||||
finish_task = asyncio.create_task(
|
||||
self._finish(status=status),
|
||||
name=f"conversation-recorder-finish:{self.session_id}",
|
||||
)
|
||||
try:
|
||||
await asyncio.shield(finish_task)
|
||||
except asyncio.CancelledError:
|
||||
# Shield prevents an interrupted pipeline cleanup from abandoning a
|
||||
# checked-out asyncpg connection. Preserve cancellation only after
|
||||
# the short database cleanup has returned the connection.
|
||||
await asyncio.shield(finish_task)
|
||||
raise
|
||||
|
||||
async def _finish(self, *, status: str) -> None:
|
||||
async with self._lock:
|
||||
try:
|
||||
async with SessionLocal() as db:
|
||||
|
||||
@@ -810,14 +810,22 @@ async def run_pipeline(
|
||||
)
|
||||
runner = WorkerRunner(handle_sigint=False)
|
||||
run_status = "completed"
|
||||
can_finalize_recorder = True
|
||||
try:
|
||||
await runner.add_workers(worker)
|
||||
await runner.run()
|
||||
except GeneratorExit:
|
||||
# A coroutine being synchronously closed cannot perform async cleanup.
|
||||
# The WebRTC route now prevents this by retaining every pipeline task;
|
||||
# keep this guard so a future unmanaged caller cannot corrupt asyncpg.
|
||||
run_status = "failed"
|
||||
can_finalize_recorder = False
|
||||
raise
|
||||
except Exception:
|
||||
run_status = "failed"
|
||||
raise
|
||||
finally:
|
||||
if recorder:
|
||||
if recorder and can_finalize_recorder:
|
||||
await recorder.finish(status=run_status)
|
||||
logger.info("管线已结束")
|
||||
|
||||
@@ -971,13 +979,18 @@ async def run_realtime_pipeline(
|
||||
)
|
||||
runner = WorkerRunner(handle_sigint=False)
|
||||
run_status = "completed"
|
||||
can_finalize_recorder = True
|
||||
try:
|
||||
await runner.add_workers(worker)
|
||||
await runner.run()
|
||||
except GeneratorExit:
|
||||
run_status = "failed"
|
||||
can_finalize_recorder = False
|
||||
raise
|
||||
except Exception:
|
||||
run_status = "failed"
|
||||
raise
|
||||
finally:
|
||||
if recorder:
|
||||
if recorder and can_finalize_recorder:
|
||||
await recorder.finish(status=run_status)
|
||||
logger.info("Realtime 管线已结束")
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, patch
|
||||
@@ -8,6 +9,30 @@ from services.conversation_history import ConversationRecorder
|
||||
|
||||
|
||||
class ConversationRecorderTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_finish_waits_for_database_cleanup_when_cancelled(self):
|
||||
recorder = ConversationRecorder("conv_test")
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
completed = asyncio.Event()
|
||||
|
||||
async def finish_database_write(*, status: str):
|
||||
self.assertEqual(status, "failed")
|
||||
started.set()
|
||||
await release.wait()
|
||||
completed.set()
|
||||
|
||||
recorder._finish = finish_database_write
|
||||
task = asyncio.create_task(recorder.finish(status="failed"))
|
||||
await started.wait()
|
||||
task.cancel()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.assertFalse(task.done())
|
||||
release.set()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await task
|
||||
self.assertTrue(completed.is_set())
|
||||
|
||||
async def test_fixed_reply_transcript_keeps_workflow_metadata(self):
|
||||
recorder = ConversationRecorder("conv_test")
|
||||
recorder._append = AsyncMock()
|
||||
|
||||
65
backend/tests/test_voice_webrtc_lifecycle.py
Normal file
65
backend/tests/test_voice_webrtc_lifecycle.py
Normal file
@@ -0,0 +1,65 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from routes import voice_webrtc
|
||||
|
||||
|
||||
class FakeConnection:
|
||||
pc_id = "pc_test"
|
||||
|
||||
|
||||
class WebRTCPipelineLifecycleTest(unittest.IsolatedAsyncioTestCase):
|
||||
async def asyncTearDown(self):
|
||||
tasks = list(voice_webrtc._pipeline_tasks)
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
if tasks:
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
voice_webrtc._pipeline_tasks.clear()
|
||||
voice_webrtc._connection_tasks.clear()
|
||||
voice_webrtc._active_connections.clear()
|
||||
|
||||
async def test_pipeline_task_is_retained_until_completion(self):
|
||||
connection = FakeConnection()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def pipeline():
|
||||
await release.wait()
|
||||
|
||||
task = voice_webrtc._start_pipeline_task(connection, pipeline())
|
||||
self.assertIn(task, voice_webrtc._pipeline_tasks)
|
||||
self.assertIs(voice_webrtc._connection_tasks[connection], task)
|
||||
|
||||
release.set()
|
||||
await task
|
||||
await asyncio.sleep(0)
|
||||
|
||||
self.assertNotIn(task, voice_webrtc._pipeline_tasks)
|
||||
self.assertNotIn(connection, voice_webrtc._connection_tasks)
|
||||
|
||||
async def test_stuck_pipeline_is_cancelled_after_close_grace_period(self):
|
||||
connection = FakeConnection()
|
||||
cancelled = asyncio.Event()
|
||||
|
||||
async def pipeline():
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
task = voice_webrtc._start_pipeline_task(connection, pipeline())
|
||||
with patch.object(voice_webrtc, "PIPELINE_CLOSE_GRACE_SECONDS", 0.01):
|
||||
await voice_webrtc._wait_for_pipeline_close(
|
||||
task,
|
||||
connection_id=connection.pc_id,
|
||||
)
|
||||
|
||||
self.assertTrue(cancelled.is_set())
|
||||
self.assertTrue(task.cancelled())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user