133 lines
4.2 KiB
Python
133 lines
4.2 KiB
Python
"""Shared ownership of peer connections and their pipeline tasks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Coroutine
|
|
from typing import Any
|
|
|
|
from loguru import logger
|
|
|
|
|
|
active_connections: set[object] = set()
|
|
pipeline_tasks: set[asyncio.Task[None]] = set()
|
|
connection_tasks: dict[object, asyncio.Task[None]] = {}
|
|
connection_cleanup_tasks: set[asyncio.Task[None]] = set()
|
|
DEFAULT_CLOSE_GRACE_SECONDS = 10.0
|
|
|
|
|
|
async def _disconnect_connection(connection: object) -> None:
|
|
"""Release media even when construction failed before transport startup."""
|
|
try:
|
|
async with asyncio.timeout(DEFAULT_CLOSE_GRACE_SECONDS):
|
|
await connection.disconnect()
|
|
except Exception as exc:
|
|
logger.warning(f"Realtime 连接清理失败: {type(exc).__name__}")
|
|
|
|
|
|
def _consume_pipeline_result(task: asyncio.Task[None], connection: object) -> None:
|
|
pipeline_tasks.discard(task)
|
|
if connection_tasks.get(connection) is task:
|
|
connection_tasks.pop(connection, None)
|
|
if connection in active_connections:
|
|
active_connections.discard(connection)
|
|
# Run after the pipeline task is done: closed handlers can safely wait
|
|
# for that task without making pipeline shutdown wait on itself.
|
|
cleanup = asyncio.create_task(
|
|
_disconnect_connection(connection),
|
|
name=f"realtime-disconnect:{getattr(connection, 'pc_id', id(connection))}",
|
|
)
|
|
connection_cleanup_tasks.add(cleanup)
|
|
cleanup.add_done_callback(connection_cleanup_tasks.discard)
|
|
try:
|
|
error = task.exception()
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Realtime pipeline 已取消: task={task.get_name()}")
|
|
return
|
|
if error is not None:
|
|
logger.opt(exception=error).error(
|
|
f"Realtime pipeline 异常结束: task={task.get_name()}"
|
|
)
|
|
|
|
|
|
def start_pipeline_task(
|
|
connection: object,
|
|
coroutine: Coroutine[Any, Any, None],
|
|
*,
|
|
protocol: str,
|
|
) -> asyncio.Task[None]:
|
|
connection_id = str(getattr(connection, "pc_id", id(connection)))
|
|
task = asyncio.create_task(
|
|
coroutine,
|
|
name=f"{protocol}-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,
|
|
timeout: float = DEFAULT_CLOSE_GRACE_SECONDS,
|
|
) -> None:
|
|
if task is None:
|
|
return
|
|
try:
|
|
await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
|
|
return
|
|
except TimeoutError:
|
|
logger.warning(
|
|
f"Realtime pipeline 关闭超过 {timeout:g} 秒，执行取消: "
|
|
f"connection_id={connection_id}"
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
return
|
|
|
|
task.cancel()
|
|
done, _pending = await asyncio.wait({task}, timeout=timeout)
|
|
if not done:
|
|
logger.error(f"Realtime pipeline 取消后仍未退出: connection_id={connection_id}")
|
|
|
|
|
|
async def shutdown_active_sessions(
|
|
*,
|
|
timeout: float = DEFAULT_CLOSE_GRACE_SECONDS,
|
|
) -> None:
|
|
connections = list(active_connections)
|
|
if connections:
|
|
active_connections.difference_update(connections)
|
|
await asyncio.gather(
|
|
*(_disconnect_connection(connection) for connection in connections),
|
|
return_exceptions=True,
|
|
)
|
|
|
|
cleanup_tasks = list(connection_cleanup_tasks)
|
|
if cleanup_tasks:
|
|
await asyncio.gather(*cleanup_tasks, return_exceptions=True)
|
|
|
|
tasks = list(pipeline_tasks)
|
|
if not tasks:
|
|
return
|
|
done, pending = await asyncio.wait(tasks, timeout=timeout)
|
|
for task in pending:
|
|
task.cancel()
|
|
cancelled: set[asyncio.Task[None]] = set()
|
|
if pending:
|
|
cancelled, stuck = await asyncio.wait(pending, timeout=timeout)
|
|
if stuck:
|
|
logger.error(f"应用关闭时仍有 {len(stuck)} 个 Realtime pipeline 未退出")
|
|
logger.info(
|
|
f"Realtime 会话清理完成: normal={len(done)} cancelled={len(cancelled)}"
|
|
)
|