long_articles_agentic_supply/supply/infra/shared/async_tasks.py

import asyncio
from typing import Callable, Coroutine, List, Any, Dict


async def run_tasks_with_async_worker_group(
    task_list: List[Any],
    handler: Callable[[Any], Coroutine[Any, Any, None]],
    *,
    description: str = None,
    unit: str = "item",
    max_concurrency: int = 20,
    fail_fast: bool = False,
    show_progress: bool = True,
) -> Dict[str, Any]:
    """使用异步 worker pool 处理 I/O 密集型任务

    Args:
        task_list: 任务对象列表
        handler: 异步处理函数，接受单个任务对象
        description: 进度条描述
        unit: 进度条单位
        max_concurrency: 最大并发 worker 数
        fail_fast: True 时任一 worker 异常立即中止全部
        show_progress: 是否显示 tqdm 进度条（需要安装 tqdm）

    Returns:
        {"total_task": int, "processed_task": int, "errors": [(idx, obj, exc), ...]}
    """
    if not task_list:
        return {"total_task": 0, "processed_task": 0, "errors": []}

    total_task = len(task_list)
    processed_task = 0
    errors: List[tuple] = []

    queue: asyncio.Queue = asyncio.Queue()
    cancel_event = asyncio.Event()
    counter_lock = asyncio.Lock()

    # 进度条（可选）
    if show_progress:
        try:
            from tqdm.asyncio import tqdm as async_tqdm

            pbar = async_tqdm(total=total_task, unit=unit, desc=description)
        except ImportError:
            pbar = None
    else:
        pbar = None

    async def worker(worker_id: int):
        nonlocal processed_task
        while True:
            item = await queue.get()
            if item is None:
                queue.task_done()
                break

            idx, task_obj = item
            try:
                if cancel_event.is_set():
                    return

                await handler(task_obj)
                async with counter_lock:
                    processed_task += 1

            except Exception as e:
                if fail_fast:
                    cancel_event.set()
                    raise
                errors.append((idx, task_obj, e))

            finally:
                if pbar:
                    pbar.update()
                queue.task_done()

    workers = [asyncio.create_task(worker(i)) for i in range(max_concurrency)]

    try:
        for index, task in enumerate(task_list, start=1):
            await queue.put((index, task))

        await queue.join()

    except Exception:
        for w in workers:
            w.cancel()
        raise

    finally:
        for _ in range(max_concurrency):
            await queue.put(None)

        await asyncio.gather(*workers, return_exceptions=True)
        if pbar:
            pbar.close()

    return {
        "total_task": total_task,
        "processed_task": processed_task,
        "errors": errors,
    }


async def run_tasks_with_asyncio_task_group(
    task_list: List[Any],
    handler: Callable[[Any], Coroutine[Any, Any, None]],
    *,
    description: str = None,
    unit: str = "item",
    max_concurrency: int = 20,
    fail_fast: bool = False,
    show_progress: bool = True,
) -> Dict[str, Any]:
    """使用 asyncio.TaskGroup 处理 I/O 密集型任务

    Args:
        task_list: 任务对象列表
        handler: 异步处理函数，接受单个任务对象
        description: 进度条描述
        unit: 进度条单位
        max_concurrency: 最大并发数（通过 Semaphore 控制）
        fail_fast: True 时任一任务异常立即中止全部
        show_progress: 是否显示 tqdm 进度条

    Returns:
        {"total_task": int, "processed_task": int, "errors": [(idx, obj, exc), ...]}
    """
    if not task_list:
        return {"total_task": 0, "processed_task": 0, "errors": []}

    processed_task = 0
    total_task = len(task_list)
    errors: List[tuple] = []
    semaphore = asyncio.Semaphore(max_concurrency)

    # 进度条（可选）
    if show_progress:
        try:
            from tqdm.asyncio import tqdm as async_tqdm

            pbar = async_tqdm(total=total_task, unit=unit, desc=description)
        except ImportError:
            pbar = None
    else:
        pbar = None

    lock = asyncio.Lock()

    async def _run_single_task(task_obj: Any, idx: int):
        nonlocal processed_task
        async with semaphore:
            try:
                await handler(task_obj)
                async with lock:
                    processed_task += 1
            except Exception as e:
                if fail_fast:
                    raise e
                errors.append((idx, task_obj, e))
            finally:
                if pbar:
                    pbar.update()

    async with asyncio.TaskGroup() as task_group:
        for index, task in enumerate(task_list, start=1):
            task_group.create_task(
                _run_single_task(task, index),
                name=f"processing {description}-{index}",
            )

    if pbar:
        pbar.close()

    return {
        "total_task": total_task,
        "processed_task": processed_task,
        "errors": errors,
    }