Improve handling of FASTAPI Asyncio in API

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

api_app/_version.py

@@ -1 +1 @@
-__version__ = "0.25.26"
+__version__ = "0.25.27"

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

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]