225 lines
7.9 KiB
Python
225 lines
7.9 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import os
|
||
import threading
|
||
from concurrent.futures import Future, TimeoutError as FutureTimeoutError
|
||
from typing import Awaitable, TypeVar
|
||
|
||
from app.config import settings
|
||
|
||
logger = logging.getLogger("video_gen")
|
||
|
||
T = TypeVar("T")
|
||
|
||
_thread_local = threading.local()
|
||
_single_loop_lock = threading.RLock()
|
||
_single_loop: asyncio.AbstractEventLoop | None = None
|
||
_single_loop_thread: threading.Thread | None = None
|
||
_single_loop_pid: int | None = None
|
||
_single_loop_ready: threading.Event | None = None
|
||
|
||
|
||
async def _dispose_async_resources() -> None:
|
||
"""释放当前 async loop 内缓存的异步资源。
|
||
|
||
Celery soft time limit 会打断同步等待 future.result() 的线程；如果不主动
|
||
cancel coroutine 并释放 engine/redis，后台 loop 里残留的协程可能继续占用
|
||
SQLAlchemy QueuePool 连接，后续任务就会出现 QueuePool timeout。
|
||
"""
|
||
try:
|
||
from app.services.redis_registry_service import close_registry_redis
|
||
|
||
await close_registry_redis()
|
||
except Exception:
|
||
logger.debug("关闭 Celery Redis registry 连接失败", exc_info=True)
|
||
|
||
try:
|
||
from app.models.base import engine
|
||
|
||
await engine.dispose()
|
||
except Exception:
|
||
logger.debug("dispose Celery SQLAlchemy engine 失败", exc_info=True)
|
||
|
||
|
||
def _runner_mode() -> str:
|
||
mode = str(getattr(settings, "CELERY_ASYNC_RUNNER_MODE", "single_loop") or "single_loop").strip().lower()
|
||
if mode not in {"single_loop", "direct"}:
|
||
return "single_loop"
|
||
return mode
|
||
|
||
|
||
def _get_or_create_thread_local_loop() -> asyncio.AbstractEventLoop:
|
||
"""兼容旧方案：当前线程持有一个长期 event loop。
|
||
|
||
仅作为降级模式使用。长期推荐 single_loop，避免 Windows threads 下
|
||
多线程 event loop 复用 asyncpg / redis.asyncio 连接对象。
|
||
"""
|
||
pid = os.getpid()
|
||
loop = getattr(_thread_local, "loop", None)
|
||
loop_pid = getattr(_thread_local, "pid", None)
|
||
|
||
if loop is None or loop.is_closed() or loop_pid != pid:
|
||
loop = asyncio.new_event_loop()
|
||
_thread_local.loop = loop
|
||
_thread_local.pid = pid
|
||
|
||
return loop
|
||
|
||
|
||
def _single_loop_worker(loop: asyncio.AbstractEventLoop, ready: threading.Event) -> None:
|
||
asyncio.set_event_loop(loop)
|
||
ready.set()
|
||
try:
|
||
loop.run_forever()
|
||
finally:
|
||
pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
|
||
if pending:
|
||
for task in pending:
|
||
task.cancel()
|
||
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
||
|
||
loop.run_until_complete(_dispose_async_resources())
|
||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||
loop.close()
|
||
|
||
|
||
def _get_or_create_single_loop() -> asyncio.AbstractEventLoop:
|
||
"""获取当前 Celery 进程内唯一 async event loop。
|
||
|
||
Linux prefork：每个 Celery 子进程各自一个 loop。
|
||
Windows threads：同一 worker 进程内所有任务线程共享同一个 loop。
|
||
"""
|
||
global _single_loop, _single_loop_thread, _single_loop_pid, _single_loop_ready
|
||
|
||
pid = os.getpid()
|
||
with _single_loop_lock:
|
||
if (
|
||
_single_loop is not None
|
||
and not _single_loop.is_closed()
|
||
and _single_loop_thread is not None
|
||
and _single_loop_thread.is_alive()
|
||
and _single_loop_pid == pid
|
||
):
|
||
return _single_loop
|
||
|
||
# fork 后 pid 变化，必须丢弃父进程状态，重新创建子进程自己的 loop。
|
||
_single_loop = asyncio.new_event_loop()
|
||
_single_loop_pid = pid
|
||
_single_loop_ready = threading.Event()
|
||
_single_loop_thread = threading.Thread(
|
||
target=_single_loop_worker,
|
||
args=(_single_loop, _single_loop_ready),
|
||
name=f"celery-async-runner-{pid}",
|
||
daemon=True,
|
||
)
|
||
_single_loop_thread.start()
|
||
_single_loop_ready.wait(timeout=5)
|
||
return _single_loop
|
||
|
||
|
||
def _cancel_future_and_reset_loop(future: Future[T] | None, *, reason: str) -> None:
|
||
"""取消当前协程并重置当前进程内 event loop。"""
|
||
if future is not None and not future.done():
|
||
future.cancel()
|
||
try:
|
||
future.result(timeout=2)
|
||
except Exception:
|
||
pass
|
||
|
||
logger.warning("Celery async_runner 正在重置 event loop。reason=%s", reason)
|
||
close_loop()
|
||
|
||
|
||
def run_async(coro: Awaitable[T]) -> T:
|
||
"""Celery 同步 task 调用异步协程的统一入口。
|
||
|
||
默认 single_loop 模式：
|
||
- 一个 Celery 子进程只有一个专用 event loop；
|
||
- 所有 asyncpg / redis.asyncio 操作都在这个 loop 内创建和使用；
|
||
- 避免 got Future attached to a different loop；
|
||
- 当 Celery soft time limit 打断 future.result() 时，主动 cancel 后台协程并
|
||
释放连接池，避免 QueuePool 被残留任务长期占用。
|
||
|
||
降级 direct 模式：
|
||
- 兼容旧的线程本地 loop 方案；
|
||
- 如果使用 direct，建议同时开启 CELERY_DB_USE_NULLPOOL=true。
|
||
"""
|
||
if _runner_mode() == "direct":
|
||
loop = _get_or_create_thread_local_loop()
|
||
try:
|
||
return loop.run_until_complete(coro)
|
||
except BaseException:
|
||
try:
|
||
if not loop.is_closed():
|
||
loop.run_until_complete(_dispose_async_resources())
|
||
finally:
|
||
close_loop()
|
||
raise
|
||
|
||
loop = _get_or_create_single_loop()
|
||
try:
|
||
running_loop = asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
running_loop = None
|
||
|
||
if running_loop is loop:
|
||
raise RuntimeError("run_async() 不能在 Celery async_runner 的事件循环内部被同步调用")
|
||
|
||
future: Future[T] = asyncio.run_coroutine_threadsafe(coro, loop)
|
||
try:
|
||
return future.result()
|
||
except FutureTimeoutError:
|
||
_cancel_future_and_reset_loop(future, reason="future_result_timeout")
|
||
raise
|
||
except BaseException:
|
||
# Celery SoftTimeLimitExceeded/worker shutdown 等异常会从这里抛出。
|
||
# 必须重置 loop，否则后台协程继续运行会拖住 DB 连接池。
|
||
_cancel_future_and_reset_loop(future, reason="base_exception")
|
||
raise
|
||
|
||
|
||
def close_loop() -> None:
|
||
"""关闭当前进程内 async runner loop。"""
|
||
global _single_loop, _single_loop_thread, _single_loop_pid, _single_loop_ready
|
||
|
||
# 关闭 single_loop。
|
||
with _single_loop_lock:
|
||
loop = _single_loop
|
||
thread = _single_loop_thread
|
||
if loop is not None and not loop.is_closed() and thread is not None and thread.is_alive():
|
||
try:
|
||
cleanup_future = asyncio.run_coroutine_threadsafe(_dispose_async_resources(), loop)
|
||
cleanup_future.result(timeout=5)
|
||
except Exception:
|
||
logger.debug("关闭 loop 前清理 async 资源失败", exc_info=True)
|
||
try:
|
||
loop.call_soon_threadsafe(loop.stop)
|
||
thread.join(timeout=5)
|
||
except Exception:
|
||
logger.debug("关闭 Celery async_runner loop 失败", exc_info=True)
|
||
|
||
_single_loop = None
|
||
_single_loop_thread = None
|
||
_single_loop_pid = None
|
||
_single_loop_ready = None
|
||
|
||
# 关闭 direct 降级模式的线程本地 loop。
|
||
loop = getattr(_thread_local, "loop", None)
|
||
if loop is not None and not loop.is_closed():
|
||
try:
|
||
pending = [task for task in asyncio.all_tasks(loop) if not task.done()]
|
||
for task in pending:
|
||
task.cancel()
|
||
if pending:
|
||
loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
||
loop.run_until_complete(_dispose_async_resources())
|
||
loop.run_until_complete(loop.shutdown_asyncgens())
|
||
except Exception:
|
||
logger.debug("关闭 direct loop 前清理失败", exc_info=True)
|
||
finally:
|
||
loop.close()
|
||
_thread_local.loop = None
|
||
_thread_local.pid = None
|