deregister log streamer via starlette background task

Author: bisgaard-itis

services/api-server/src/simcore_service_api_server/api/routes/solvers_jobs_getters.py

@@ -3,6 +3,7 @@
 import logging
 from collections import deque
 from collections.abc import Callable
+from functools import partial
 from typing import Annotated, Any, Union
 from uuid import UUID
 
@@ -16,10 +17,10 @@
 from models_library.wallets import ZERO_CREDITS
 from pydantic import HttpUrl, NonNegativeInt
 from pydantic.types import PositiveInt
-from servicelib.fastapi.requests_decorators import cancel_on_disconnect
 from servicelib.logging_utils import log_context
 from simcore_service_api_server.models.api_resources import parse_resources_ids
 from sqlalchemy.ext.asyncio import AsyncEngine
+from starlette.background import BackgroundTask
 
 from ..._service_solvers import SolverService
 from ...exceptions.custom_errors import InsufficientCreditsError, MissingWalletError
@@ -526,7 +527,6 @@ async def get_job_pricing_unit(
     response_class=LogStreamingResponse,
     responses=_LOGSTREAM_STATUS_CODES,
 )
-@cancel_on_disconnect
 async def get_log_stream(
     request: Request,
     solver_key: SolverKeyId,
@@ -553,6 +553,8 @@ async def get_log_stream(
             log_distributor=log_distributor,
             log_check_timeout=log_check_timeout,
         )
+        await log_distributor.register(job_id, log_streamer.queue)
         return LogStreamingResponse(
             log_streamer.log_generator(),
+            background=BackgroundTask(partial(log_distributor.deregister, job_id)),
         )

services/api-server/src/simcore_service_api_server/services_http/log_streaming.py

@@ -69,6 +69,7 @@ async def _distribute_logs(self, data: bytes):
         return False
 
     async def register(self, job_id: JobID, queue: Queue[JobLog]):
+        _logger.debug("Registering log streamer for job_id=%s", job_id)
         if job_id in self._log_streamers:
             raise LogStreamerRegistrationConflictError(job_id=job_id)
         self._log_streamers[job_id] = queue
@@ -77,6 +78,7 @@ async def register(self, job_id: JobID, queue: Queue[JobLog]):
         )
 
     async def deregister(self, job_id: JobID):
+        _logger.debug("Deregistering log streamer for job_id=%s", job_id)
         if job_id not in self._log_streamers:
             msg = f"No stream was connected to {job_id}."
             raise LogStreamerNotRegisteredError(details=msg, job_id=job_id)
@@ -103,7 +105,7 @@ def __init__(
     ):
         self._user_id = user_id
         self._director2_api = director2_api
-        self._queue: Queue[JobLog] = Queue()
+        self.queue: Queue[JobLog] = Queue()
         self._job_id: JobID = job_id
         self._log_distributor: LogDistributor = log_distributor
         self._log_check_timeout: NonNegativeInt = log_check_timeout
@@ -116,12 +118,11 @@ async def _project_done(self) -> bool:
 
     async def log_generator(self) -> AsyncIterable[str]:
         try:
-            await self._log_distributor.register(self._job_id, self._queue)
             done: bool = False
             while not done:
                 try:
                     log: JobLog = await asyncio.wait_for(
-                        self._queue.get(), timeout=self._log_check_timeout
+                        self.queue.get(), timeout=self._log_check_timeout
                     )
                     yield log.model_dump_json() + _NEW_LINE
                 except TimeoutError:
@@ -145,6 +146,3 @@ async def log_generator(self) -> AsyncIterable[str]:
                 )
             )
             yield ErrorGet(errors=[error_msg]).model_dump_json() + _NEW_LINE
-
-        finally:
-            await self._log_distributor.deregister(self._job_id)

services/api-server/tests/unit/test_services_rabbitmq.py

@@ -451,7 +451,7 @@ async def test_log_generator(mocker: MockFixture, faker: Faker):
         msg = faker.text()
         published_logs.append(msg)
         job_log.messages = [msg]
-        await log_streamer._queue.put(job_log)
+        await log_streamer.queue.put(job_log)
 
     collected_logs: list[str] = []
     async for log in log_streamer.log_generator():