Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 39 additions & 7 deletions backend/app/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import logging
import time
from collections.abc import Awaitable, Callable
from contextlib import asynccontextmanager, suppress
from contextlib import asynccontextmanager

from fastapi import Depends, FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
Expand Down Expand Up @@ -243,6 +243,43 @@ def _prewarm_jieba() -> None:
logger.debug("jieba prewarm skipped", exc_info=True)


async def _shutdown_prewarm_tasks(
cache_task: asyncio.Task | None,
jieba_task: asyncio.Task | None,
*,
jieba_timeout: float = 10.0,
) -> None:
"""停机第一步:回收预热任务。异常隔离 + jieba 超时(issue #73)。

- 缓存预热任务以非 CancelledError 异常结束时只记 warning,
不再中断余下的调度器停机 / drain / DuckDB 关闭 / engine.dispose;
- jieba 预热底层是 to_thread(不可真正取消),必须 wait_for 加超时,
否则词典加载卡住会挂死整个优雅停机。
"""
if cache_task is not None:
if not cache_task.done():
cache_task.cancel()
try:
# 已完成但带异常的任务同样要 await 取回(否则异常静默丢失);
# 非 CancelledError 只记 warning,不中断余下停机步骤。
await cache_task
except asyncio.CancelledError:
pass
except Exception:
logger.warning("Cache warmup task ended with error during shutdown", exc_info=True)

if jieba_task is None:
return
try:
await asyncio.wait_for(jieba_task, timeout=jieba_timeout)
except TimeoutError:
logger.warning("Jieba prewarm did not finish within %.0fs; proceeding with shutdown", jieba_timeout)
except asyncio.CancelledError:
pass
except Exception:
logger.warning("Jieba prewarm task failed during shutdown", exc_info=True)


@asynccontextmanager
async def lifespan(app: FastAPI):
global _cache_warmup_task
Expand Down Expand Up @@ -416,12 +453,7 @@ async def _seed_model_catalog() -> None:
yield

# Shutdown: stop scheduler, close connections, dispose engine
if _cache_warmup_task and not _cache_warmup_task.done():
_cache_warmup_task.cancel()
with suppress(asyncio.CancelledError):
await _cache_warmup_task
with suppress(Exception):
await _jieba_prewarm_task
await _shutdown_prewarm_tasks(_cache_warmup_task, _jieba_prewarm_task)
shutdown_scheduler()

# 回收受管后台任务(调度器启动 rescan/恢复、日报后台生成等),
Expand Down
60 changes: 60 additions & 0 deletions backend/tests/test_shutdown_prewarm.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
"""停机预热任务回收回归(#73 / D-10)。

钉死 _shutdown_prewarm_tasks 的三条契约:
- 缓存预热任务带非 CancelledError 异常结束时:记 warning、不向调用方抛出
(否则会跳过调度器停机 / DuckDB 关闭 / engine.dispose 全部余下清理);
- jieba 预热超时:wait_for 兜底返回,不挂死停机;
- 正常/已取消路径静默通过。
"""

from __future__ import annotations

import asyncio
import logging

import pytest

from app.main import _shutdown_prewarm_tasks


@pytest.mark.asyncio
async def test_cache_task_error_does_not_propagate(caplog):
async def boom():
raise ValueError("warmup exploded")

cache = asyncio.create_task(boom())
jieba = asyncio.create_task(asyncio.sleep(0))
await asyncio.wait({cache}) # 让异常先发生(未启动即 cancel 只会得到 CancelledError)

with caplog.at_level(logging.WARNING):
await _shutdown_prewarm_tasks(cache, jieba) # 不应抛出

assert any("Cache warmup task ended with error" in r.getMessage() for r in caplog.records)


@pytest.mark.asyncio
async def test_jieba_timeout_does_not_hang():
async def stuck():
await asyncio.sleep(60)

cache = None
jieba = asyncio.create_task(stuck())

# 0.05s 超时:若 wait_for 缺失或失效,本用例会挂到全局超时
await asyncio.wait_for(
_shutdown_prewarm_tasks(cache, jieba, jieba_timeout=0.05),
timeout=5.0,
)


@pytest.mark.asyncio
async def test_normal_and_cancelled_paths_pass_quietly():
done_cache = asyncio.create_task(asyncio.sleep(0))
await done_cache
cancelled_jieba = asyncio.create_task(asyncio.sleep(60))
cancelled_jieba.cancel()
with pytest.raises(asyncio.CancelledError):
await cancelled_jieba

await _shutdown_prewarm_tasks(done_cache, cancelled_jieba) # 不抛即通过
await _shutdown_prewarm_tasks(None, None) # None 容忍
Loading