periodic_tasks.py

"""后台任务的所属启动、停止信号与取消安全等待（保留周期任务导出名）。"""
import asyncio
import logging

logger = logging.getLogger("ComfyUI-Ranking.PeriodicTasks")


def start_periodic_task(state, name, factory):
    current = getattr(state, name + "_worker", None)
    if current is not None and not current.done():
        return current
    stop = asyncio.Event()
    coroutine = factory(stop)
    try:
        task = asyncio.create_task(coroutine)
    except BaseException:
        coroutine.close()
        raise
    setattr(state, name + "_stop", stop)
    setattr(state, name + "_worker", task)

    def consume(completed):
        if not completed.cancelled():
            error = completed.exception()
            if error is not None:
                logger.warning("后台任务结束异常: task=%s | error_type=%s", name, type(error).__name__)
    task.add_done_callback(consume)
    return task


def request_periodic_stop(state, names):
    tasks = []
    for name in names:
        stop = getattr(state, name + "_stop", None)
        if stop is not None:
            stop.set()
        task = getattr(state, name + "_worker", None)
        if task is not None:
            tasks.append(task)
    return tasks


async def wait_periodic_tasks(tasks):
    cancelled = None
    for task in tasks:
        while not task.done():
            try:
                await asyncio.shield(task)
            except asyncio.CancelledError as error:
                # worker自身取消与关闭等待者取消分别处理。
                if asyncio.current_task().cancelling() or not task.cancelled():
                    cancelled = error
            except Exception:
                break
        if not task.cancelled():
            # 业务异常已有所属回调诊断，不阻止其它任务收尾和最终同步。
            task.exception()
    if cancelled is not None:
        raise cancelled