diff --git a/CHANGELOG.md b/CHANGELOG.md index aefae053b..ddcfc8ba5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ ENHANCEMENTS: * Update API, CLI, and UI dependencies to address high-severity Dependabot alerts, including `PyJWT`, `Vite`, `lodash`, `fast-uri`, `flatted`, `immutable`, and `minimatch`. BUG FIXES: +* Prevent silent garbage collection of background API task workers by implementing app-scoped lifecycle management. ([#4923](https://github.com/microsoft/AzureTRE/issues/4923)) * Fix API timeout and name collision failures on workspace creation by checking storage account name availability and improved logging. ([#4946](https://github.com/microsoft/AzureTRE/pull/4946)) * Fix error handling in airlock processor ([#4929](https://github.com/microsoft/AzureTRE/pull/4929)) * Fix dependabot high severity alerts for packages fast-uri, lodash, picomatch, immutable, minimatch, flatted and PyJWT diff --git a/api_app/_version.py b/api_app/_version.py index eda2a21a4..7df7ad28d 100644 --- a/api_app/_version.py +++ b/api_app/_version.py @@ -1 +1 @@ -__version__ = "0.25.26" +__version__ = "0.25.27" diff --git a/api_app/main.py b/api_app/main.py index 0bdc76914..6b7a3f893 100644 --- a/api_app/main.py +++ b/api_app/main.py @@ -1,4 +1,5 @@ import asyncio +from typing import Any import uvicorn from fastapi import FastAPI @@ -22,8 +23,25 @@ from service_bus.airlock_request_status_update import AirlockStatusUpdater +class BackgroundTaskManager: + def __init__(self) -> None: + self._tasks: set[asyncio.Task[Any]] = set() + self.is_shutting_down: bool = False + + def add(self, task: asyncio.Task[Any]) -> None: + self._tasks.add(task) + + def discard(self, task: asyncio.Task[Any]) -> None: + self._tasks.discard(task) + + def get_tasks(self) -> list[asyncio.Task[Any]]: + return list(self._tasks) + + @asynccontextmanager async def lifespan(app: FastAPI): + app.state.background_tasks = BackgroundTaskManager() + while not await bootstrap_database(): await asyncio.sleep(5) logger.warning("Database connection could not be established") @@ -34,9 +52,66 @@ async def lifespan(app: FastAPI): airlockStatusUpdater = AirlockStatusUpdater() await airlockStatusUpdater.init_repos() - asyncio.create_task(deploymentStatusUpdater.receive_messages()) - asyncio.create_task(airlockStatusUpdater.receive_messages()) - yield + def track(task: asyncio.Task[Any]) -> None: + def _done_callback(task: asyncio.Task[Any]) -> None: + app.state.background_tasks.discard(task) + if app.state.background_tasks.is_shutting_down: + return + + if task.cancelled(): + return + + try: + exception = task.exception() + except asyncio.CancelledError: + return + + if exception is not None: + logger.error( + f"Background task {task.get_name()} failed", + exc_info=(type(exception), exception, exception.__traceback__) + ) + + app.state.background_tasks.add(task) + task.add_done_callback(_done_callback) + + track(asyncio.create_task( + deploymentStatusUpdater.receive_messages(), + name="deployment-status-updater" + )) + track(asyncio.create_task( + airlockStatusUpdater.receive_messages(), + name="airlock-status-updater" + )) + + try: + yield + finally: + app.state.background_tasks.is_shutting_down = True + tasks = app.state.background_tasks.get_tasks() + logger.info(f"Cancelling {len(tasks)} background tasks") + + for task in tasks: + logger.debug(f"Cancelling task {task.get_name()}") + task.cancel() + + if tasks: + try: + results = await asyncio.wait_for( + asyncio.gather(*tasks, return_exceptions=True), + timeout=10.0 + ) + for task, result in zip(tasks, results): + if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError): + logger.warning( + f"Background task {task.get_name()} raised exception during shutdown: {result}", + exc_info=(type(result), result, result.__traceback__) + ) + except asyncio.TimeoutError: + logger.error("Timeout waiting for background tasks to shutdown") + pending = [t for t in tasks if not t.done()] + for t in pending: + logger.warning(f"Task {t.get_name()} did not terminate in time during shutdown") def get_application() -> FastAPI: @@ -74,4 +149,4 @@ def get_application() -> FastAPI: FastAPIInstrumentor.instrument_app(app) if __name__ == "__main__": - uvicorn.run(app, host="0.0.0.0", port=8000, loop="asyncio") + uvicorn.run(app, host="0.0.0.0", port=8000, loop="asyncio") # nosec B104: intentional bind to all interfaces diff --git a/api_app/tests_ma/test_main_lifecycle.py b/api_app/tests_ma/test_main_lifecycle.py new file mode 100644 index 000000000..8b365522b --- /dev/null +++ b/api_app/tests_ma/test_main_lifecycle.py @@ -0,0 +1,127 @@ +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch +import pytest +from fastapi import FastAPI + + +@pytest.mark.asyncio +async def test_background_task_manager(): + from main import BackgroundTaskManager + + manager = BackgroundTaskManager() + assert manager.is_shutting_down is False + + async def dummy(): + pass + + task = asyncio.create_task(dummy()) + manager.add(task) + assert task in manager.get_tasks() + + manager.discard(task) + assert task not in manager.get_tasks() + + await task + + +@pytest.mark.asyncio +@patch("main.bootstrap_database", return_value=True) +@patch("main.DeploymentStatusUpdater") +@patch("main.AirlockStatusUpdater") +async def test_lifespan_lifecycle_normal(mock_airlock, mock_deployment, mock_bootstrap): + from main import BackgroundTaskManager, lifespan + + mock_deployment_instance = MagicMock() + mock_deployment_instance.init_repos = AsyncMock() + mock_deployment_instance.receive_messages = AsyncMock() + mock_deployment.return_value = mock_deployment_instance + + mock_airlock_instance = MagicMock() + mock_airlock_instance.init_repos = AsyncMock() + mock_airlock_instance.receive_messages = AsyncMock() + mock_airlock.return_value = mock_airlock_instance + + app = FastAPI() + + async with lifespan(app): + assert hasattr(app.state, "background_tasks") + assert isinstance(app.state.background_tasks, BackgroundTaskManager) + tasks = app.state.background_tasks.get_tasks() + assert len(tasks) == 2 + names = {t.get_name() for t in tasks} + assert "deployment-status-updater" in names + assert "airlock-status-updater" in names + + assert app.state.background_tasks.is_shutting_down is True + + +@pytest.mark.asyncio +@patch("main.bootstrap_database", return_value=True) +@patch("main.DeploymentStatusUpdater") +@patch("main.AirlockStatusUpdater") +@patch("main.logger") +async def test_lifespan_task_failure_during_runtime(mock_logger, mock_airlock, mock_deployment, mock_bootstrap): + from main import lifespan + + mock_deployment_instance = MagicMock() + mock_deployment_instance.init_repos = AsyncMock() + + async def fail_immediately(): + raise ValueError("Runtime failure") + + mock_deployment_instance.receive_messages = fail_immediately + mock_deployment.return_value = mock_deployment_instance + + mock_airlock_instance = MagicMock() + mock_airlock_instance.init_repos = AsyncMock() + mock_airlock_instance.receive_messages = AsyncMock() + mock_airlock.return_value = mock_airlock_instance + + app = FastAPI() + + async with lifespan(app): + await asyncio.sleep(0.1) + + # Verify that the runtime failure was logged as an error + error_logs = [call for call in mock_logger.error.call_args_list if "Background task deployment-status-updater failed" in call[0][0]] + assert len(error_logs) == 1 + + +@pytest.mark.asyncio +@patch("main.bootstrap_database", return_value=True) +@patch("main.DeploymentStatusUpdater") +@patch("main.AirlockStatusUpdater") +@patch("main.logger") +async def test_lifespan_task_failure_during_shutdown(mock_logger, mock_airlock, mock_deployment, mock_bootstrap): + from main import lifespan + + mock_deployment_instance = MagicMock() + mock_deployment_instance.init_repos = AsyncMock() + + async def raise_teardown_exception(): + try: + await asyncio.sleep(5) + except asyncio.CancelledError: + raise RuntimeError("Database connection closed abruptly") + + mock_deployment_instance.receive_messages = raise_teardown_exception + mock_deployment.return_value = mock_deployment_instance + + mock_airlock_instance = MagicMock() + mock_airlock_instance.init_repos = AsyncMock() + mock_airlock_instance.receive_messages = AsyncMock() + mock_airlock.return_value = mock_airlock_instance + + app = FastAPI() + + async with lifespan(app): + await asyncio.sleep(0.1) + + # The error should NOT be logged as error in _done_callback because it was gated on is_shutting_down + for call in mock_logger.error.call_args_list: + assert "Background task deployment-status-updater failed" not in call[0][0] + + # Instead, it should be caught during shutdown gather and logged as warning + warning_logs = [call for call in mock_logger.warning.call_args_list if "raised exception during shutdown" in call[0][0]] + assert len(warning_logs) == 1 + assert "deployment-status-updater" in warning_logs[0][0][0]