147 lines
4.8 KiB
Python
147 lines
4.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import threading
|
|
from concurrent.futures import Future
|
|
from typing import Awaitable, TypeVar
|
|
|
|
from app.config import settings
|
|
|
|
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
|
|
|
|
|
|
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()
|
|
loop.run_forever()
|
|
|
|
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(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 run_async(coro: Awaitable[T]) -> T:
|
|
"""Celery 同步 task 调用异步协程的统一入口。
|
|
|
|
默认 single_loop 模式：
|
|
- 一个 Celery 子进程只有一个专用 event loop；
|
|
- 所有 asyncpg / redis.asyncio 操作都在这个 loop 内创建和使用；
|
|
- 避免 got Future attached to a different loop。
|
|
|
|
降级 direct 模式：
|
|
- 兼容旧的线程本地 loop 方案；
|
|
- 如果使用 direct，建议同时开启 CELERY_DB_USE_NULLPOOL=true。
|
|
"""
|
|
if _runner_mode() == "direct":
|
|
loop = _get_or_create_thread_local_loop()
|
|
return loop.run_until_complete(coro)
|
|
|
|
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)
|
|
return future.result()
|
|
|
|
|
|
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():
|
|
loop.call_soon_threadsafe(loop.stop)
|
|
thread.join(timeout=5)
|
|
|
|
_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():
|
|
loop.close()
|
|
_thread_local.loop = None
|
|
_thread_local.pid = None
|