"""进程内后台任务登记与生命周期收口。""" from __future__ import annotations import asyncio import concurrent.futures from collections.abc import Coroutine from dataclasses import dataclass from functools import partial from typing import Any, Callable @dataclass(frozen=True, slots=True) class TaskRecord: """记录一个后台任务的所有者,便于关停阶段按责任域收口。""" owner: str task: asyncio.Task[Any] cancel_on_shutdown: bool class TaskRegistry: """管理由宿主创建的进程内后台任务,并提供统一取消与等待入口。""" def __init__(self) -> None: """初始化空任务登记表。""" self._records: dict[asyncio.Task[Any], TaskRecord] = {} self._shutdown_cancel_requested: set[asyncio.Task[Any]] = set() self._shutdown_timeout_reported: set[asyncio.Task[Any]] = set() self._accepting = True @property def records(self) -> tuple[TaskRecord, ...]: """返回当前仍未完成的任务快照。""" return tuple( record for record in self._records.values() if not record.task.done() ) def create( self, coroutine: Coroutine[Any, Any, Any], *, owner: str, cancel_on_shutdown: bool = True, ) -> asyncio.Task[Any]: """创建并登记后台任务,任务完成后自动从登记表移除。""" if not self._accepting: coroutine.close() raise RuntimeError("后台任务登记器正在关闭,不能再创建新任务") task = asyncio.create_task(coroutine, name=owner) self.register( task, owner=owner, cancel_on_shutdown=cancel_on_shutdown, ) return task def create_sync( self, function: Callable[..., Any], *args: Any, owner: str, **kwargs: Any, ) -> asyncio.Task[Any]: """在线程池执行同步后台函数并登记其异步生命周期。""" return self.create( asyncio.to_thread(partial(function, *args, **kwargs)), owner=owner, cancel_on_shutdown=False, ) def submit_threadsafe( self, coroutine: Coroutine[Any, Any, Any], *, loop: asyncio.AbstractEventLoop, owner: str, cancel_on_shutdown: bool = True, ) -> concurrent.futures.Future[Any]: """从宿主线程提交协程,并在目标循环内原子登记 owner 后执行。""" completion: concurrent.futures.Future[Any] = concurrent.futures.Future() task_holder: dict[str, asyncio.Task[Any]] = {} def mirror_completion(task: asyncio.Task[Any]) -> None: """把登记任务的真实终态镜像给跨线程调用方。""" if completion.done(): return if task.cancelled(): completion.cancel() return exception = task.exception() if exception is not None: completion.set_exception(exception) else: completion.set_result(task.result()) def submit_on_loop() -> None: """在目标循环内完成 accepting 检查、任务创建和 owner 登记。""" if completion.cancelled(): coroutine.close() return try: task = self.create( coroutine, owner=owner, cancel_on_shutdown=cancel_on_shutdown, ) except Exception as error: if not completion.done(): completion.set_exception(error) loop.call_exception_handler( { "message": "MoviePilot 跨线程后台任务提交失败", "exception": error, "owner": owner, } ) return task_holder["task"] = task task.add_done_callback(mirror_completion) if completion.cancelled() and not task.done(): task.cancel() def cancel_registered_task( submitted: concurrent.futures.Future[Any], ) -> None: """调用方取消 completion 时,把取消请求转交目标循环中的真实任务。""" if not submitted.cancelled(): return task = task_holder.get("task") if task is not None and not task.done(): try: loop.call_soon_threadsafe(task.cancel) except RuntimeError: pass completion.add_done_callback(cancel_registered_task) try: loop.call_soon_threadsafe(submit_on_loop) except RuntimeError: coroutine.close() raise return completion def register( self, task: asyncio.Task[Any], *, owner: str, cancel_on_shutdown: bool = True, ) -> asyncio.Task[Any]: """登记已有任务并绑定责任域。""" if not self._accepting: task.cancel() raise RuntimeError("后台任务登记器正在关闭,不能再登记新任务") task.set_name(owner) self._records[task] = TaskRecord( owner=owner, task=task, cancel_on_shutdown=cancel_on_shutdown, ) task.add_done_callback(self._discard) return task def _discard(self, task: asyncio.Task[Any]) -> None: """移除已结束任务,并把未处理异常交给事件循环统一报告。""" record = self._records.pop(task, None) self._shutdown_cancel_requested.discard(task) self._shutdown_timeout_reported.discard(task) if task.cancelled(): return exception = task.exception() if exception is not None: task.get_loop().call_exception_handler( { "message": "MoviePilot 后台任务执行失败", "exception": exception, "task": task, "owner": record.owner if record else task.get_name(), } ) async def shutdown(self, *, timeout_seconds: float = 10.0) -> bool: """停止接收并有限等待存量任务,返回全部 owner 是否真实收敛。""" self._accepting = False records = self.records tasks = [record.task for record in records] for record in records: if ( record.cancel_on_shutdown and record.task not in self._shutdown_cancel_requested ): self._shutdown_cancel_requested.add(record.task) record.task.cancel() if tasks: await asyncio.wait(tasks, timeout=timeout_seconds) unfinished = tuple( record for record in records if not record.task.done() and record.task not in self._shutdown_timeout_reported ) if unfinished: self._shutdown_timeout_reported.update( record.task for record in unfinished ) asyncio.get_running_loop().call_exception_handler( { "message": "MoviePilot 后台任务未在关停预算内结束", "owners": tuple(record.owner for record in unfinished), "tasks": tuple(record.task for record in unfinished), "timeout_seconds": timeout_seconds, } ) return all(record.task.done() for record in records) _default_registry = TaskRegistry() _runtime_registry: TaskRegistry | None = None def configure_task_registry(registry: TaskRegistry | None) -> None: """由启动组合根发布当前 lifespan 的任务登记器。""" global _runtime_registry _runtime_registry = registry def get_task_registry() -> TaskRegistry: """返回当前宿主任务登记器,未启动完整 lifespan 时保留测试兼容回退。""" return _runtime_registry or _default_registry