fix: manage WebRTC pipeline shutdown

This commit is contained in:
Xin Wang
2026-08-05 13:48:29 +08:00
parent 10bea48d72
commit 16156028c9
6 changed files with 253 additions and 9 deletions
+129 -6
View File
@@ -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"],