celery异步恢复任务独立队列|扩增celery子进程连接池上限配置

video-gen-api/app/tasks/async_runner.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
 import asyncio
+import logging
 import os
 import threading
-from concurrent.futures import Future
+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()
@@ -18,6 +21,28 @@ _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"}:
@@ -46,16 +71,18 @@ def _get_or_create_thread_local_loop() -> asyncio.AbstractEventLoop:
 def _single_loop_worker(loop: asyncio.AbstractEventLoop, ready: threading.Event) -> None:
     asyncio.set_event_loop(loop)
     ready.set()
-    loop.run_forever()
+    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))
 
-    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()
+        loop.run_until_complete(_dispose_async_resources())
+        loop.run_until_complete(loop.shutdown_asyncgens())
+        loop.close()
 
 
 def _get_or_create_single_loop() -> asyncio.AbstractEventLoop:
@@ -92,13 +119,28 @@ def _get_or_create_single_loop() -> asyncio.AbstractEventLoop:
         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。
+    - 避免 got Future attached to a different loop；
+    - 当 Celery soft time limit 打断 future.result() 时，主动 cancel 后台协程并
+      释放连接池，避免 QueuePool 被残留任务长期占用。
 
     降级 direct 模式：
     - 兼容旧的线程本地 loop 方案；
@@ -106,7 +148,15 @@ def run_async(coro: Awaitable[T]) -> T:
     """
     if _runner_mode() == "direct":
         loop = _get_or_create_thread_local_loop()
-        return loop.run_until_complete(coro)
+        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:
@@ -118,7 +168,16 @@ def run_async(coro: Awaitable[T]) -> T:
         raise RuntimeError("run_async() 不能在 Celery async_runner 的事件循环内部被同步调用")
 
     future: Future[T] = asyncio.run_coroutine_threadsafe(coro, loop)
-    return future.result()
+    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:
@@ -130,8 +189,16 @@ def close_loop() -> None:
         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)
+            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
@@ -141,6 +208,17 @@ def close_loop() -> None:
     # 关闭 direct 降级模式的线程本地 loop。
     loop = getattr(_thread_local, "loop", None)
     if loop is not None and not loop.is_closed():
-        loop.close()
+        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