🐛 Deregister log streamer via starlette background task (#7626)

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,9 +17,9 @@
 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 sqlalchemy.ext.asyncio import AsyncEngine
+from starlette.background import BackgroundTask
 
 from ..._service_solvers import SolverService
 from ...exceptions.custom_errors import InsufficientCreditsError, MissingWalletError
@@ -478,7 +479,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,
@@ -505,8 +505,10 @@ 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,13 +78,14 @@ 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)
         await self._rabbit_client.remove_topics(
             LoggerRabbitMessage.get_channel_name(), topics=[f"{job_id}.*"]
         )
-        del self._log_streamers[job_id]
+        self._log_streamers.pop(job_id)
 
     @property
     def iter_log_queue_sizes(self) -> Iterator[tuple[JobID, int]]:
@@ -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

@@ -361,6 +361,7 @@ async def test_log_streamer_with_distributor(
     project_id: ProjectID,
     node_id: NodeID,
     produce_logs: Callable,
+    log_distributor: LogDistributor,
     log_streamer_with_distributor: LogStreamer,
     faker: Faker,
     computation_done: Callable[[], bool],
@@ -375,12 +376,21 @@ async def _log_publisher():
 
     publish_task = asyncio.create_task(_log_publisher())
 
+    @asynccontextmanager
+    async def registered_log_streamer():
+        await log_distributor.register(project_id, log_streamer_with_distributor.queue)
+        try:
+            yield
+        finally:
+            await log_distributor.deregister(project_id)
+
     collected_messages: list[str] = []
-    async for log in log_streamer_with_distributor.log_generator():
-        job_log: JobLog = JobLog.model_validate_json(log)
-        assert len(job_log.messages) == 1
-        assert job_log.job_id == project_id
-        collected_messages.append(job_log.messages[0])
+    async with registered_log_streamer():
+        async for log in log_streamer_with_distributor.log_generator():
+            job_log: JobLog = JobLog.model_validate_json(log)
+            assert len(job_log.messages) == 1
+            assert job_log.job_id == project_id
+            collected_messages.append(job_log.messages[0])
 
     if not publish_task.done():
         publish_task.cancel()
@@ -451,7 +461,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():