fix: await takeover polling jobs

This commit is contained in:
stefanfeng
2026-08-19 14:32:10 +08:00
parent 4ab7732da9
commit c5bfa47a23
2 changed files with 68 additions and 10 deletions
+40 -6
View File
@@ -4,7 +4,7 @@ from fastapi.middleware.cors import CORSMiddleware
import os import os
import logging import logging
from apscheduler.schedulers.background import BackgroundScheduler from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger from apscheduler.triggers.interval import IntervalTrigger
import redis as redis_lib import redis as redis_lib
@@ -23,6 +23,9 @@ from responses import ok
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
takeover_scheduler = None
takeover_db = None
app = FastAPI(title="会会数字分身 API", version="1.0.0") app = FastAPI(title="会会数字分身 API", version="1.0.0")
app.add_middleware( app.add_middleware(
@@ -110,9 +113,14 @@ def seed():
@app.on_event("startup") @app.on_event("startup")
def on_startup(): def on_startup():
global takeover_scheduler, takeover_db
init_db() init_db()
seed() seed()
# Release stale resources when startup is invoked again by a reload/test.
stop_takeover_scheduler()
# --- Takeover scheduler --- # --- Takeover scheduler ---
try: try:
# Initialize Redis (optional) # Initialize Redis (optional)
@@ -137,16 +145,42 @@ def on_startup():
# Initialize takeover service # Initialize takeover service
from services.takeover_service import TakeoverService from services.takeover_service import TakeoverService
takeover_service = TakeoverService(SessionLocal(), boxim_client, redis_client) takeover_db = SessionLocal()
takeover_service = TakeoverService(takeover_db, boxim_client, redis_client)
# Start periodic polling job # AsyncIOScheduler awaits the service coroutine instead of dropping it.
scheduler = BackgroundScheduler() takeover_scheduler = AsyncIOScheduler()
scheduler.add_job( takeover_scheduler.add_job(
takeover_service.poll_and_process_messages, takeover_service.poll_and_process_messages,
trigger=IntervalTrigger(seconds=10), trigger=IntervalTrigger(seconds=10),
id="takeover_message_poll", id="takeover_message_poll",
) )
scheduler.start() takeover_scheduler.start()
logger.info("Takeover message polling scheduler started (interval=10s)") logger.info("Takeover message polling scheduler started (interval=10s)")
except Exception as e: except Exception as e:
stop_takeover_scheduler()
logger.warning(f"Failed to initialize takeover scheduler, app will continue without it: {e}") logger.warning(f"Failed to initialize takeover scheduler, app will continue without it: {e}")
def stop_takeover_scheduler():
global takeover_scheduler, takeover_db
if takeover_scheduler is not None:
try:
if takeover_scheduler.running:
takeover_scheduler.shutdown(wait=False)
except Exception as e:
logger.warning(f"Failed to stop takeover scheduler cleanly: {e}")
finally:
takeover_scheduler = None
if takeover_db is not None:
try:
takeover_db.close()
finally:
takeover_db = None
@app.on_event("shutdown")
def on_shutdown():
stop_takeover_scheduler()
@@ -14,7 +14,7 @@ def test_app_has_startup_event():
@patch("services.takeover_service.TakeoverService") @patch("services.takeover_service.TakeoverService")
@patch("services.boxim_client.BoxIMClient") @patch("services.boxim_client.BoxIMClient")
@patch("main.redis_lib.from_url") @patch("main.redis_lib.from_url")
@patch("main.BackgroundScheduler") @patch("main.AsyncIOScheduler")
def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_url, mock_boxim_cls, mock_takeover_cls): def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_url, mock_boxim_cls, mock_takeover_cls):
"""Verify scheduler is initialized when Redis is available.""" """Verify scheduler is initialized when Redis is available."""
mock_redis = MagicMock() mock_redis = MagicMock()
@@ -25,6 +25,7 @@ def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_
mock_boxim_cls.return_value = mock_boxim mock_boxim_cls.return_value = mock_boxim
mock_takeover = MagicMock() mock_takeover = MagicMock()
mock_takeover.poll_and_process_messages = AsyncMock()
mock_takeover_cls.return_value = mock_takeover mock_takeover_cls.return_value = mock_takeover
with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://localhost:6379"}): with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://localhost:6379"}):
@@ -32,19 +33,23 @@ def test_scheduler_initialized_with_redis(mock_scheduler_class, mock_redis_from_
on_startup() on_startup()
mock_scheduler_class.return_value.add_job.assert_called_once() mock_scheduler_class.return_value.add_job.assert_called_once()
scheduled_callable = mock_scheduler_class.return_value.add_job.call_args.args[0]
call_kwargs = mock_scheduler_class.return_value.add_job.call_args[1] call_kwargs = mock_scheduler_class.return_value.add_job.call_args[1]
assert scheduled_callable is mock_takeover.poll_and_process_messages
assert call_kwargs["id"] == "takeover_message_poll" assert call_kwargs["id"] == "takeover_message_poll"
mock_scheduler_class.return_value.start.assert_called_once_with()
@patch("services.takeover_service.TakeoverService") @patch("services.takeover_service.TakeoverService")
@patch("services.boxim_client.BoxIMClient") @patch("services.boxim_client.BoxIMClient")
@patch("main.BackgroundScheduler") @patch("main.AsyncIOScheduler")
def test_scheduler_starts_without_redis(mock_scheduler_class, mock_boxim_cls, mock_takeover_cls): def test_scheduler_starts_without_redis(mock_scheduler_class, mock_boxim_cls, mock_takeover_cls):
"""App should start even when REDIS_URL is not set.""" """App should start even when REDIS_URL is not set."""
mock_boxim = MagicMock() mock_boxim = MagicMock()
mock_boxim_cls.return_value = mock_boxim mock_boxim_cls.return_value = mock_boxim
mock_takeover = MagicMock() mock_takeover = MagicMock()
mock_takeover.poll_and_process_messages = AsyncMock()
mock_takeover_cls.return_value = mock_takeover mock_takeover_cls.return_value = mock_takeover
with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": ""}, clear=False): with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": ""}, clear=False):
@@ -57,7 +62,7 @@ def test_scheduler_starts_without_redis(mock_scheduler_class, mock_boxim_cls, mo
@patch("services.takeover_service.TakeoverService") @patch("services.takeover_service.TakeoverService")
@patch("services.boxim_client.BoxIMClient") @patch("services.boxim_client.BoxIMClient")
@patch("main.redis_lib.from_url") @patch("main.redis_lib.from_url")
@patch("main.BackgroundScheduler") @patch("main.AsyncIOScheduler")
def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from_url, mock_boxim_cls, mock_takeover_cls): def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from_url, mock_boxim_cls, mock_takeover_cls):
"""App should start even when Redis ping fails.""" """App should start even when Redis ping fails."""
mock_redis_from_url.side_effect = ConnectionError("Connection refused") mock_redis_from_url.side_effect = ConnectionError("Connection refused")
@@ -66,6 +71,7 @@ def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from
mock_boxim_cls.return_value = mock_boxim mock_boxim_cls.return_value = mock_boxim
mock_takeover = MagicMock() mock_takeover = MagicMock()
mock_takeover.poll_and_process_messages = AsyncMock()
mock_takeover_cls.return_value = mock_takeover mock_takeover_cls.return_value = mock_takeover
with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://badhost:6379"}): with patch("main.init_db"), patch("main.seed"), patch.dict("os.environ", {"REDIS_URL": "redis://badhost:6379"}):
@@ -75,7 +81,7 @@ def test_scheduler_starts_when_redis_fails(mock_scheduler_class, mock_redis_from
mock_scheduler_class.return_value.add_job.assert_called_once() mock_scheduler_class.return_value.add_job.assert_called_once()
@patch("main.BackgroundScheduler") @patch("main.AsyncIOScheduler")
def test_scheduler_fails_gracefully(mock_scheduler_class): def test_scheduler_fails_gracefully(mock_scheduler_class):
"""If scheduler init raises, the app should still start (exception caught).""" """If scheduler init raises, the app should still start (exception caught)."""
mock_scheduler_class.side_effect = RuntimeError("Scheduler crash") mock_scheduler_class.side_effect = RuntimeError("Scheduler crash")
@@ -87,6 +93,24 @@ def test_scheduler_fails_gracefully(mock_scheduler_class):
# No exception should propagate # No exception should propagate
def test_scheduler_shutdown_releases_resources():
"""Shutdown should stop polling and close its dedicated database session."""
import main
mock_scheduler = MagicMock()
mock_scheduler.running = True
mock_db = MagicMock()
main.takeover_scheduler = mock_scheduler
main.takeover_db = mock_db
main.on_shutdown()
mock_scheduler.shutdown.assert_called_once_with(wait=False)
mock_db.close.assert_called_once_with()
assert main.takeover_scheduler is None
assert main.takeover_db is None
# --- poll_and_process_messages --- # --- poll_and_process_messages ---