48 lines
1.3 KiB
Python
48 lines
1.3 KiB
Python
import asyncio
|
||
import os
|
||
import threading
|
||
from typing import Awaitable, TypeVar
|
||
|
||
T = TypeVar("T")
|
||
|
||
_thread_local = threading.local()
|
||
|
||
|
||
def _get_or_create_loop() -> asyncio.AbstractEventLoop:
|
||
"""
|
||
给当前进程/线程维护一个长期 event loop。
|
||
|
||
Linux prefork：
|
||
每个 Celery 子进程通常单线程跑任务，这里相当于每个子进程一个长期 loop。
|
||
|
||
Windows -P threads：
|
||
每个线程一个 loop，但注意 asyncpg pool 仍不适合跨线程共享；
|
||
Windows threads 模式建议继续用 NullPool 或只做本地调试。
|
||
"""
|
||
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 run_async(coro: Awaitable[T]) -> T:
|
||
"""
|
||
Celery 同步 task 调用异步协程的统一入口。
|
||
不使用 asyncio.run()，避免每个 task 结束时关闭 event loop。
|
||
"""
|
||
loop = _get_or_create_loop()
|
||
return loop.run_until_complete(coro)
|
||
|
||
|
||
def close_loop() -> None:
|
||
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 |