feat: move redis client lifecycle to app server's one

Author: giancarloromeo

packages/celery-library/src/celery_library/signals.py

@@ -8,7 +8,6 @@
 from servicelib.logging_utils import log_context
 from settings_library.celery import CelerySettings
 
-from .common import create_task_manager
 from .utils import get_app_server, set_app_server
 
 _logger = logging.getLogger(__name__)
@@ -26,20 +25,15 @@ def _init(startup_complete_event: threading.Event) -> None:
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
 
-        async def _setup_task_manager():
+        async def _setup():
             assert sender.app  # nosec
             assert isinstance(sender.app, Celery)  # nosec
 
-            app_server.task_manager = await create_task_manager(
-                sender.app,
-                celery_settings,
-            )
-
             set_app_server(sender.app, app_server)
 
         app_server.event_loop = loop
 
-        loop.run_until_complete(_setup_task_manager())
+        loop.run_until_complete(_setup())
         loop.run_until_complete(app_server.lifespan(startup_complete_event))
 
     thread = threading.Thread(

packages/service-library/src/servicelib/celery/models.py

@@ -58,7 +58,7 @@ async def get_task_metadata(self, task_id: TaskID) -> TaskMetadata | None: ...
 
     async def get_task_progress(self, task_id: TaskID) -> ProgressReport | None: ...
 
-    async def list_tasks(self, task_context: TaskFilter) -> list[Task]: ...
+    async def list_tasks(self, task_filter: TaskFilter) -> list[Task]: ...
 
     async def remove_task(self, task_id: TaskID) -> None: ...
 

services/storage/src/simcore_service_storage/core/application.py

@@ -36,7 +36,7 @@
 from ..dsm import setup_dsm
 from ..dsm_cleaner import setup_dsm_cleaner
 from ..exceptions.handlers import set_exception_handlers
-from ..modules.celery import setup_task_manager
+from ..modules.celery import setup_celery
 from ..modules.db import setup_db
 from ..modules.long_running_tasks import setup_rest_api_long_running_tasks_for_uploads
 from ..modules.rabbitmq import setup as setup_rabbitmq
@@ -71,12 +71,13 @@ def create_app(settings: ApplicationSettings) -> FastAPI:  # noqa: C901
     setup_s3(app)
     setup_client_session(app, tracing_settings=settings.STORAGE_TRACING)
 
-    if settings.STORAGE_CELERY and not settings.STORAGE_WORKER_MODE:
-        setup_rabbitmq(app)
+    if settings.STORAGE_CELERY:
+        setup_celery(app, settings=settings.STORAGE_CELERY)
 
-        setup_task_manager(app, celery_settings=settings.STORAGE_CELERY)
+        if not settings.STORAGE_WORKER_MODE:
+            setup_rabbitmq(app)
+            setup_rpc_routes(app)
 
-        setup_rpc_routes(app)
     setup_rest_api_long_running_tasks_for_uploads(app)
     setup_rest_api_routes(app, API_VTAG)
     set_exception_handlers(app)

services/storage/src/simcore_service_storage/modules/celery/__init__.py

@@ -1,26 +1,55 @@
-from celery_library.common import create_app, create_task_manager
+import logging
+
+from celery_library.backends._redis import RedisTaskInfoStore
+from celery_library.common import create_app
 from celery_library.task_manager import CeleryTaskManager
 from celery_library.types import register_celery_types, register_pydantic_types
 from fastapi import FastAPI
 from models_library.api_schemas_storage.storage_schemas import (
     FileUploadCompletionBody,
     FoldersBody,
 )
+from servicelib.logging_utils import log_context
+from servicelib.redis import RedisClientSDK
 from settings_library.celery import CelerySettings
+from settings_library.redis import RedisDatabase
 
 from ...models import FileMetaData
 
+_logger = logging.getLogger(__name__)
+
 
-def setup_task_manager(app: FastAPI, celery_settings: CelerySettings) -> None:
+def setup_celery(app: FastAPI, settings: CelerySettings) -> None:
     async def on_startup() -> None:
-        app.state.task_manager = await create_task_manager(
-            create_app(celery_settings), celery_settings
-        )
+        with log_context(_logger, logging.INFO, "Setting up Celery"):
+            redis_client_sdk = RedisClientSDK(
+                settings.CELERY_REDIS_RESULT_BACKEND.build_redis_dsn(
+                    RedisDatabase.CELERY_TASKS
+                ),
+                client_name="celery_tasks",
+            )
+            app.state.celery_tasks_redis_client_sdk = redis_client_sdk
+            await redis_client_sdk.setup()
+
+            app.state.task_manager = CeleryTaskManager(
+                create_app(settings),
+                settings,
+                RedisTaskInfoStore(redis_client_sdk),
+            )
+
+            register_celery_types()
+            register_pydantic_types(FileUploadCompletionBody, FileMetaData, FoldersBody)
 
-        register_celery_types()
-        register_pydantic_types(FileUploadCompletionBody, FileMetaData, FoldersBody)
+    async def on_shutdown() -> None:
+        with log_context(_logger, logging.INFO, "Shutting down Celery"):
+            redis_client_sdk: RedisClientSDK | None = (
+                app.state.celery_tasks_redis_client_sdk
+            )
+            if redis_client_sdk:
+                await redis_client_sdk.shutdown()
 
     app.add_event_handler("startup", on_startup)
+    app.add_event_handler("shutdown", on_shutdown)
 
 
 def get_task_manager_from_app(app: FastAPI) -> CeleryTaskManager: