✨ Refactor: Update service setup functions to use async iterators for improved lifecycle management

Author: pcrespov

services/catalog/src/simcore_service_catalog/core/application.py

@@ -22,12 +22,7 @@
 from ..api.rest.routes import setup_rest_api_routes
 from ..api.rpc.routes import setup_rpc_api_routes
 from ..exceptions.handlers import setup_exception_handlers
-from ..services.function_services import setup_function_services
-from ..services.rabbitmq import setup_rabbitmq
 from . import events
-from .events import (
-    _create_on_shutdown,
-)
 from .settings import ApplicationSettings
 
 _logger = logging.getLogger(__name__)
@@ -73,14 +68,10 @@ def create_app() -> FastAPI:
     if settings.CATALOG_TRACING:
         initialize_tracing(app, settings.CATALOG_TRACING, APP_NAME)
 
-    # PLUGIN SETUP
-    setup_function_services(app)
-    setup_rabbitmq(app)
-
+    # MIDDLEWARES
     if settings.CATALOG_PROMETHEUS_INSTRUMENTATION_ENABLED:
         setup_prometheus_instrumentation(app)
 
-    # MIDDLEWARES
     if settings.CATALOG_PROFILING:
         initialize_profiler(app)
 
@@ -96,9 +87,6 @@ def create_app() -> FastAPI:
     setup_rest_api_routes(app, vtag=API_VTAG)
     setup_rpc_api_routes(app)
 
-    # SHUTDOWN-EVENT
-    app.add_event_handler("shutdown", _create_on_shutdown(app))
-
     # EXCEPTIONS
     setup_exception_handlers(app)
 

services/catalog/src/simcore_service_catalog/core/background_tasks.py

@@ -11,11 +11,13 @@
 
 import asyncio
 import logging
+from collections.abc import AsyncIterator
 from contextlib import suppress
 from pprint import pformat
 from typing import Final
 
 from fastapi import FastAPI, HTTPException
+from fastapi_lifespan_manager import State
 from models_library.services import ServiceMetaDataPublished
 from models_library.services_types import ServiceKey, ServiceVersion
 from packaging.version import Version
@@ -115,9 +117,9 @@ async def _ensure_registry_and_database_are_synced(app: FastAPI) -> None:
     director_api = get_director_api(app)
     services_in_manifest_map = await manifest.get_services_map(director_api)
 
-    services_in_db: set[
-        tuple[ServiceKey, ServiceVersion]
-    ] = await _list_services_in_database(app.state.engine)
+    services_in_db: set[tuple[ServiceKey, ServiceVersion]] = (
+        await _list_services_in_database(app.state.engine)
+    )
 
     # check that the db has all the services at least once
     missing_services_in_db = set(services_in_manifest_map.keys()) - services_in_db
@@ -232,3 +234,11 @@ async def stop_registry_sync_task(app: FastAPI) -> None:
             await task
         app.state.registry_sync_task = None
     _logger.info("registry syncing task stopped")
+
+
+async def setup_background_task(app: FastAPI) -> AsyncIterator[State]:
+    await start_registry_sync_task(app)
+    try:
+        yield {}
+    finally:
+        await stop_registry_sync_task(app)

services/catalog/src/simcore_service_catalog/core/events.py

@@ -1,28 +1,24 @@
-import contextlib
 import logging
-from collections.abc import AsyncIterator, Awaitable, Callable
-from typing import TypeAlias
+from collections.abc import AsyncIterator
 
 from fastapi import FastAPI
 from fastapi_lifespan_manager import LifespanManager, State
 from servicelib.fastapi.postgres_lifespan import (
     PostgresLifespanStateKeys,
     postgres_lifespan,
 )
-from servicelib.logging_utils import log_context
 
 from .._meta import APP_FINISHED_BANNER_MSG, APP_STARTED_BANNER_MSG
 from ..db.events import setup_database
-from ..services.director import close_director, setup_director
-from .background_tasks import start_registry_sync_task, stop_registry_sync_task
+from ..services.director import setup_director
+from ..services.function_services import setup_function_services
+from ..services.rabbitmq import setup_rabbitmq
+from .background_tasks import setup_background_task
 from .settings import ApplicationSettings
 
 _logger = logging.getLogger(__name__)
 
 
-EventCallable: TypeAlias = Callable[[], Awaitable[None]]
-
-
 def flush_started_banner() -> None:
     # WARNING: this function is spied in the tests
     print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
@@ -32,7 +28,7 @@ def flush_finished_banner() -> None:
     print(APP_FINISHED_BANNER_MSG, flush=True)  # noqa: T201
 
 
-async def _main_setup(app: FastAPI) -> AsyncIterator[State]:
+async def _setup_app(app: FastAPI) -> AsyncIterator[State]:
     flush_started_banner()
 
     settings: ApplicationSettings = app.state.settings
@@ -44,55 +40,25 @@ async def _main_setup(app: FastAPI) -> AsyncIterator[State]:
     flush_finished_banner()
 
 
-def _create_on_startup(app: FastAPI) -> EventCallable:
-    async def _() -> None:
-
-        if app.state.settings.CATALOG_DIRECTOR:
-            # setup connection to director
-            await setup_director(app)
-
-            # FIXME: check director service is in place and ready. Hand-shake??
-            # SEE https://github.com/ITISFoundation/osparc-simcore/issues/1728
-            await start_registry_sync_task(app)
-
-        _logger.info("Application started")
-
-    return _
-
-
-def _create_on_shutdown(app: FastAPI) -> EventCallable:
-    async def _() -> None:
-
-        with log_context(_logger, logging.INFO, "Application shutdown"):
-            if app.state.settings.CATALOG_DIRECTOR:
-                try:
-                    await stop_registry_sync_task(app)
-                    await close_director(app)
-                except Exception:  # pylint: disable=broad-except
-                    _logger.exception("Unexpected error while closing application")
-
-    return _
-
-
-@contextlib.asynccontextmanager
-async def _other_setup(app: FastAPI) -> AsyncIterator[State]:
-
-    await _create_on_startup(app)()
-
-    yield {}
-
-    await _create_on_shutdown(app)()
-
-
 def create_app_lifespan():
     # app lifespan
     app_lifespan = LifespanManager()
-    app_lifespan.add(_main_setup)
+    app_lifespan.add(_setup_app)
 
     # - postgres lifespan
     postgres_lifespan.add(setup_database)
     app_lifespan.include(postgres_lifespan)
 
-    app_lifespan.add(_other_setup)
+    # - rabbitmq lifespan
+    app_lifespan.add(setup_rabbitmq)
+
+    # - director lifespan
+    app_lifespan.add(setup_director)
+
+    # - function services lifespan
+    app_lifespan.add(setup_function_services)
+
+    # - background task lifespan
+    app_lifespan.add(setup_background_task)
 
     return app_lifespan

services/catalog/src/simcore_service_catalog/services/director.py

@@ -3,14 +3,15 @@
 import json
 import logging
 import urllib.parse
-from collections.abc import Awaitable, Callable
+from collections.abc import AsyncIterator, Awaitable, Callable
 from contextlib import suppress
 from pprint import pformat
 from typing import Any, Final
 
 import httpx
 from common_library.json_serialization import json_dumps
 from fastapi import FastAPI, HTTPException
+from fastapi_lifespan_manager import State
 from models_library.api_schemas_directorv2.services import ServiceExtras
 from models_library.services_metadata_published import ServiceMetaDataPublished
 from models_library.services_types import ServiceKey, ServiceVersion
@@ -65,7 +66,7 @@ def _validate_kind(entry_to_validate: dict[str, Any], kind_name: str):
 
 
 def _return_data_or_raise_error(
-    request_func: Callable[..., Awaitable[httpx.Response]]
+    request_func: Callable[..., Awaitable[httpx.Response]],
 ) -> Callable[..., Awaitable[list[Any] | dict[str, Any]]]:
     """
     Creates a context for safe inter-process communication (IPC)
@@ -288,13 +289,14 @@ async def get_service_extras(
         return TypeAdapter(ServiceExtras).validate_python(result)
 
 
-async def setup_director(app: FastAPI) -> None:
+async def setup_director(app: FastAPI) -> AsyncIterator[State]:
+    client: DirectorApi | None = None
+
     if settings := app.state.settings.CATALOG_DIRECTOR:
         with log_context(
             _logger, logging.DEBUG, "Setup director at %s", f"{settings.base_url=}"
         ):
             async for attempt in AsyncRetrying(**_director_startup_retry_policy):
-                client = DirectorApi(base_url=settings.base_url, app=app)
                 with attempt:
                     client = DirectorApi(base_url=settings.base_url, app=app)
                     if not await client.is_responsive():
@@ -303,17 +305,16 @@ async def setup_director(app: FastAPI) -> None:
                         raise DirectorUnresponsiveError
 
                 _logger.info(
-                    "Connection to director-v0 succeded [%s]",
+                    "Connection to director-v0 succeeded [%s]",
                     json_dumps(attempt.retry_state.retry_object.statistics),
                 )
 
             # set when connected
             app.state.director_api = client
 
-
-async def close_director(app: FastAPI) -> None:
-    client: DirectorApi | None
-    if client := app.state.director_api:
-        await client.close()
-
-    _logger.debug("Director client closed successfully")
+    try:
+        yield {}
+    finally:
+        if client:
+            await client.close()
+        _logger.debug("Director client closed successfully")

services/catalog/src/simcore_service_catalog/services/function_services.py

@@ -1,9 +1,12 @@
+from collections.abc import AsyncIterator
+
 # mypy: disable-error-code=truthy-function
 from typing import Any
 
 from fastapi import status
 from fastapi.applications import FastAPI
 from fastapi.exceptions import HTTPException
+from fastapi_lifespan_manager import State
 from models_library.function_services_catalog import (
     is_function_service,
     iter_service_docker_data,
@@ -31,12 +34,15 @@ def get_function_service(key, version) -> ServiceMetaDataPublished:
         ) from err
 
 
-def setup_function_services(app: FastAPI):
-    def _on_startup() -> None:
-        catalog = [_as_dict(metadata) for metadata in iter_service_docker_data()]
-        app.state.frontend_services_catalog = catalog
+async def setup_function_services(app: FastAPI) -> AsyncIterator[State]:
+    app.state.frontend_services_catalog = [
+        _as_dict(metadata) for metadata in iter_service_docker_data()
+    ]
 
-    app.add_event_handler("startup", _on_startup)
+    try:
+        yield {}
+    finally:
+        app.state.frontend_services_catalog = None
 
 
 __all__: tuple[str, ...] = (

services/catalog/src/simcore_service_catalog/services/rabbitmq.py

@@ -1,7 +1,9 @@
 import logging
+from collections.abc import AsyncIterator
 from typing import cast
 
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 from servicelib.rabbitmq import RabbitMQRPCClient, wait_till_rabbitmq_responsive
 from settings_library.rabbit import RabbitSettings
 
@@ -15,25 +17,21 @@ def get_rabbitmq_settings(app: FastAPI) -> RabbitSettings:
     return settings
 
 
-def setup_rabbitmq(app: FastAPI) -> None:
+async def setup_rabbitmq(app: FastAPI) -> AsyncIterator[State]:
     settings: RabbitSettings = get_rabbitmq_settings(app)
-    app.state.rabbitmq_rpc_server = None
+    await wait_till_rabbitmq_responsive(settings.dsn)
 
-    async def _on_startup() -> None:
-        await wait_till_rabbitmq_responsive(settings.dsn)
+    app.state.rabbitmq_rpc_server = await RabbitMQRPCClient.create(
+        client_name=f"{PROJECT_NAME}_rpc_server", settings=settings
+    )
 
-        app.state.rabbitmq_rpc_server = await RabbitMQRPCClient.create(
-            client_name=f"{PROJECT_NAME}_rpc_server", settings=settings
-        )
-
-    async def _on_shutdown() -> None:
+    try:
+        yield {}
+    finally:
         if app.state.rabbitmq_rpc_server:
             await app.state.rabbitmq_rpc_server.close()
             app.state.rabbitmq_rpc_server = None
 
-    app.add_event_handler("startup", _on_startup)
-    app.add_event_handler("shutdown", _on_shutdown)
-
 
 def get_rabbitmq_rpc_server(app: FastAPI) -> RabbitMQRPCClient:
     assert app.state.rabbitmq_rpc_server  # nosec