diff --git a/document-parser/services/analysis_service.py b/document-parser/services/analysis_service.py index b9b9113..2059506 100644 --- a/document-parser/services/analysis_service.py +++ b/document-parser/services/analysis_service.py @@ -100,16 +100,26 @@ class AnalysisService: await _mark_failed(job_id, str(e)) +_background_tasks: set[asyncio.Task] = set() + + def _on_task_done(task: asyncio.Task, *, job_id: str) -> None: """Log unhandled exceptions from background analysis tasks and mark job as FAILED.""" if task.cancelled(): logger.warning("Analysis task was cancelled: %s", job_id) - asyncio.ensure_future(_mark_failed(job_id, "Task was cancelled")) + _schedule_mark_failed(job_id, "Task was cancelled") return exc = task.exception() if exc: logger.error("Unhandled exception in analysis task %s: %s", job_id, exc, exc_info=True) - asyncio.ensure_future(_mark_failed(job_id, str(exc))) + _schedule_mark_failed(job_id, str(exc)) + + +def _schedule_mark_failed(job_id: str, error: str) -> None: + """Schedule _mark_failed as a tracked background task.""" + t = asyncio.ensure_future(_mark_failed(job_id, error)) + _background_tasks.add(t) + t.add_done_callback(_background_tasks.discard) async def _mark_failed(job_id: str, error: str) -> None: diff --git a/document-parser/tests/test_analysis_service.py b/document-parser/tests/test_analysis_service.py index 7a9cbaa..ec26641 100644 --- a/document-parser/tests/test_analysis_service.py +++ b/document-parser/tests/test_analysis_service.py @@ -40,12 +40,12 @@ class TestOnTaskDone: async def slow_task(): await asyncio.sleep(999) + import contextlib + task = asyncio.create_task(slow_task()) task.cancel() - try: + with contextlib.suppress(asyncio.CancelledError): await task - except asyncio.CancelledError: - pass with patch("services.analysis_service._mark_failed", new_callable=AsyncMock) as mock_mark: _on_task_done(task, job_id=job_id)