首次提交

paste/core/aio_pool.py

@@ -0,0 +1,153 @@
+import asyncio
+import datetime
+from typing import Optional, Callable, Set, Any
+
+import psutil
+from dateutil.relativedelta import relativedelta
+
+from paste.core.logging import echo_log
+
+MAX_BACKGROUND_TASKS = 3000
+"""
+最大任务数量，根据服务器内存调整。
+"""
+
+aio_loop: Optional[asyncio.AbstractEventLoop] = None
+"""
+异步循环对象。
+"""
+
+aio_runner: Optional[Callable] = None
+"""
+异步方法运行器，对应:: 
+
+ asyncio.events.AbstractEventLoop.run_until_complete()
+
+方法的返回值 。
+"""
+
+global_background_tasks: Set[asyncio.Task] = set()
+"""
+全局后台任务池。
+"""
+
+
+async def run_background_task(coro, max_tasks: int = MAX_BACKGROUND_TASKS):
+    """
+    运行后台任务。
+
+    :param coro: 要在后台执行的协程。
+    :param max_tasks: 背压总量
+    """
+    global MAX_BACKGROUND_TASKS
+    if max_tasks != MAX_BACKGROUND_TASKS:
+        MAX_BACKGROUND_TASKS = max_tasks
+
+    if len(global_background_tasks) >= MAX_BACKGROUND_TASKS:
+        # 增加背压控制
+        await asyncio.wait(global_background_tasks, return_when=asyncio.FIRST_COMPLETED)
+
+    task = asyncio.create_task(coro)
+    global_background_tasks.add(task)
+    # 任务完成后自动移除
+    task.add_done_callback(global_background_tasks.discard)
+
+
+def get_aio_loop():
+    """
+    这里必须采用方法，在适当的时间点创建事件循环对象，否则会导致服务无法启动。
+    主要是测试到 EventLoop 之间的冲突，或异步事件已经在运行，导致无法顺利执行。
+
+    :return: 事件循环对象
+    """
+    global aio_loop
+    if aio_loop:
+        return aio_loop
+    else:
+        try:
+            # 尝试获取当前运行中的事件循环
+            aio_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            # 如果没有运行中的循环，才创建新的
+            aio_loop = asyncio.new_event_loop()
+        return aio_loop
+
+
+def get_aio_runner():
+    """
+    这里必须采用方法，在适当的时间点创建事件循环对象，否则会导致服务无法启动。
+    主要是测试到 EventLoop 之间的冲突，或异步事件已经在运行，导致无法顺利执行。
+
+    :return: 运行器对象
+    """
+    global aio_loop, aio_runner
+    if aio_runner:
+        return aio_runner
+    else:
+        aio_runner = get_aio_loop().run_until_complete
+        return aio_runner
+
+
+def process_info(pid: int):
+    """
+    若传入的 PID 存在，则返回进程信息，否则返回 None。
+
+    :param pid: 进程 ID
+    :return: 进程信息
+    """
+    try:
+        process = psutil.Process(pid)
+        _delta = relativedelta(datetime.datetime.now(), datetime.datetime.fromtimestamp(process.create_time()))
+        _d, _h, _m, _s = _delta.days, _delta.hours, _delta.minutes, _delta.seconds
+        process.cpu_percent()
+        return {
+            'name': process.name(),
+            'cpu_usage': f"{process.cpu_percent()}%",
+            'memory_usage': f"{process.memory_info().rss / (1024 * 1024):.3f}MB",
+            'running_time': f"{_d}天{_h}时{_m}分{_s}秒",
+        }
+    except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+        return None
+
+
+async def gather_with_retry(*coro_constructors: Callable[[], Any],
+                            max_retries: int = 3,
+                            eba: int = 0.5,
+                            return_exceptions: bool = False) -> tuple[Any, ...]:
+    """
+    封装 asyncio.gather，支持对一组并发任务进行 N 次重试。
+
+    与 asyncio.gather 对齐：
+    - 使用 *args 传参，而非列表
+    - 支持 return_exceptions 参数
+    - 返回 list 类型，顺序一致
+
+    但要求：每个参数必须是一个**无参函数**，调用后返回一个 awaitable。
+    这样才能在每次重试时重新创建任务。
+
+    Args:
+        *coro_constructors: Callable 对象，async 方法要用 lambda 封装
+        max_retries (int): 最大重试次数（总尝试次数 = max_retries + 1）
+        eba (int): 指数退避的起始等待时间
+        return_exceptions (bool): 若为 True，异常作为结果返回，不抛出
+
+    Returns:
+        list: 所有任务结果列表。若 return_exceptions=True，异常也会作为列表元素。
+
+    Raises:
+        Exception: 当 return_exceptions=False 且所有重试都失败时，抛出最后一次尝试中的第一个异常。
+    """
+    for attempt in range(max_retries + 1):
+        try:
+            tasks = [ctor() for ctor in coro_constructors]
+            results = await asyncio.gather(*tasks, return_exceptions=return_exceptions)
+            return tuple(results)
+        except Exception as e:
+            if attempt == max_retries:
+                echo_log(f"共执行 {max_retries} 次后全部失败: {str(e)}")
+            else:
+                echo_log(f"执行第 {attempt + 1} 次重试失败.")
+            # 指数退避
+            await asyncio.sleep(eba * (attempt + 1))
+
+    raise RuntimeError("Unreachable")