110 lines
3.1 KiB
Python
110 lines
3.1 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]] = {}
|
||
|
|
DEFAULT_CLOSE_GRACE_SECONDS = 10.0
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
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:
|
||
|
|
await asyncio.gather(
|
||
|
|
*(connection.disconnect() for connection in connections),
|
||
|
|
return_exceptions=True,
|
||
|
|
)
|
||
|
|
active_connections.difference_update(connections)
|
||
|
|
|
||
|
|
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)}"
|
||
|
|
)
|
||
|
|
|