zwf
154 lines
4.9 KiB
Python
154 lines
4.9 KiB
Python
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")
|