add lifespan to tracing in fastapi case

Author: bisgaard-itis

packages/service-library/src/servicelib/fastapi/tracing.py

@@ -1,8 +1,10 @@
 """Adds fastapi middleware for tracing using opentelemetry instrumentation."""
 
 import logging
+from collections.abc import AsyncIterator
 
 from fastapi import FastAPI
+from fastapi_lifespan_manager import State
 from httpx import AsyncClient, Client
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
@@ -68,7 +70,7 @@
     HAS_AIOPIKA_INSTRUMENTOR = False
 
 
-def initialize_tracing(
+def _startup(
     app: FastAPI, tracing_settings: TracingSettings, service_name: str
 ) -> None:
     if (
@@ -147,5 +149,66 @@ def initialize_tracing(
             RequestsInstrumentor().instrument()
 
 
+def _shutdown() -> None:
+    """Uninstruments all opentelemetry instrumentors that were instrumented."""
+    if HAS_AIOPG:
+        try:
+            AiopgInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument AiopgInstrumentor")
+    if HAS_AIOPIKA_INSTRUMENTOR:
+        try:
+            AioPikaInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument AioPikaInstrumentor")
+    if HAS_ASYNCPG:
+        try:
+            AsyncPGInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument AsyncPGInstrumentor")
+    if HAS_REDIS:
+        try:
+            RedisInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument RedisInstrumentor")
+    if HAS_BOTOCORE:
+        try:
+            BotocoreInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument BotocoreInstrumentor")
+    if HAS_REQUESTS:
+        try:
+            RequestsInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument RequestsInstrumentor")
+
+
 def setup_httpx_client_tracing(client: AsyncClient | Client):
     HTTPXClientInstrumentor.instrument_client(client)
+
+
+def setup_tracing(
+    app: FastAPI, tracing_settings: TracingSettings, service_name: str
+) -> None:
+
+    _startup(app, tracing_settings, service_name)
+
+    def _on_shutdown() -> None:
+        _shutdown()
+
+    app.add_event_handler("shutdown", _on_shutdown)
+
+
+async def tracing_instrumentation_lifespan(
+    *,
+    app: FastAPI,
+    state: State,
+    tracing_settings: TracingSettings,
+    service_name: str,
+) -> AsyncIterator[State]:
+
+    _startup(app, tracing_settings, service_name)
+
+    yield state
+
+    _shutdown()

packages/service-library/tests/fastapi/test_tracing.py

@@ -11,7 +11,7 @@
 import pytest
 from fastapi import FastAPI
 from pydantic import ValidationError
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import tracing_instrumentation_lifespan
 from settings_library.tracing import TracingSettings
 
 
@@ -63,17 +63,19 @@ async def test_valid_tracing_settings(
     uninstrument_opentelemetry: Iterator[None],
 ):
     tracing_settings = TracingSettings()
-    initialize_tracing(
-        mocked_app,
+    async for state in tracing_instrumentation_lifespan(
+        app=mocked_app,
+        state={},
         tracing_settings=tracing_settings,
         service_name="Mock-Openetlemetry-Pytest",
-    )
-    # idempotency
-    initialize_tracing(
-        mocked_app,
-        tracing_settings=tracing_settings,
-        service_name="Mock-Openetlemetry-Pytest",
-    )
+    ):
+        async for _ in tracing_instrumentation_lifespan(
+            app=mocked_app,
+            state=state,
+            tracing_settings=tracing_settings,
+            service_name="Mock-Openetlemetry-Pytest",
+        ):
+            pass
 
 
 @pytest.mark.parametrize(
@@ -101,11 +103,13 @@ async def test_invalid_tracing_settings(
     app = mocked_app
     with pytest.raises((BaseException, ValidationError, TypeError)):  # noqa: PT012
         tracing_settings = TracingSettings()
-        initialize_tracing(
-            app,
+        async for _ in tracing_instrumentation_lifespan(
+            app=app,
+            state={},
             tracing_settings=tracing_settings,
             service_name="Mock-Openetlemetry-Pytest",
-        )
+        ):
+            pass
 
 
 def install_package(package):
@@ -156,16 +160,19 @@ async def test_tracing_setup_package_detection(
 ):
     package_name = manage_package
     importlib.import_module(package_name)
-    #
     tracing_settings = TracingSettings()
-    initialize_tracing(
-        mocked_app,
+    # Use tracing_instrumentation_lifespan instead of _startup
+    async for _ in tracing_instrumentation_lifespan(
+        app=mocked_app,
+        state={},
         tracing_settings=tracing_settings,
         service_name="Mock-Openetlemetry-Pytest",
-    )
-    # idempotency
-    initialize_tracing(
-        mocked_app,
-        tracing_settings=tracing_settings,
-        service_name="Mock-Openetlemetry-Pytest",
-    )
+    ):
+        # idempotency: call again
+        async for _ in tracing_instrumentation_lifespan(
+            app=mocked_app,
+            state={},
+            tracing_settings=tracing_settings,
+            service_name="Mock-Openetlemetry-Pytest",
+        ):
+            pass

services/agent/src/simcore_service_agent/core/application.py

@@ -5,7 +5,7 @@
     get_common_oas_options,
     override_fastapi_openapi_method,
 )
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 from servicelib.logging_utils import config_all_loggers
 
 from .._meta import (
@@ -55,6 +55,9 @@ def create_app() -> FastAPI:
     override_fastapi_openapi_method(app)
     app.state.settings = settings
 
+    if settings.AGENT_TRACING:
+        setup_tracing(app, settings.AGENT_TRACING, APP_NAME)
+
     setup_instrumentation(app)
 
     setup_rabbitmq(app)
@@ -63,9 +66,6 @@ def create_app() -> FastAPI:
     setup_rest_api(app)
     setup_rpc_api_routes(app)
 
-    if settings.AGENT_TRACING:
-        initialize_tracing(app, settings.AGENT_TRACING, APP_NAME)
-
     async def _on_startup() -> None:
         print(APP_STARTED_BANNER_MSG, flush=True)  # noqa: T201
 

services/api-server/src/simcore_service_api_server/core/application.py

@@ -5,7 +5,7 @@
 from models_library.basic_types import BootModeEnum
 from packaging.version import Version
 from servicelib.fastapi.profiler import initialize_profiler
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 from servicelib.logging_utils import config_all_loggers
 
 from .. import exceptions
@@ -82,6 +82,9 @@ def init_app(settings: ApplicationSettings | None = None) -> FastAPI:
 
     app.state.settings = settings
 
+    if settings.API_SERVER_TRACING:
+        setup_tracing(app, settings.API_SERVER_TRACING, APP_NAME)
+
     if settings.API_SERVER_POSTGRES:
         setup_postgres(app)
 
@@ -90,9 +93,6 @@ def init_app(settings: ApplicationSettings | None = None) -> FastAPI:
     if app.state.settings.API_SERVER_PROMETHEUS_INSTRUMENTATION_ENABLED:
         setup_prometheus_instrumentation(app)
 
-    if settings.API_SERVER_TRACING:
-        initialize_tracing(app, settings.API_SERVER_TRACING, APP_NAME)
-
     if settings.API_SERVER_WEBSERVER:
         webserver.setup(
             app,

services/autoscaling/src/simcore_service_autoscaling/core/application.py

@@ -1,7 +1,7 @@
 import logging
 
 from fastapi import FastAPI
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 
 from .._meta import (
     API_VERSION,
@@ -60,6 +60,9 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     assert app.state.settings.API_VERSION == API_VERSION  # nosec
 
     # PLUGINS SETUP
+    if app.state.settings.AUTOSCALING_TRACING:
+        setup_tracing(app, app.state.settings.AUTOSCALING_TRACING, APP_NAME)
+
     setup_instrumentation(app)
     setup_api_routes(app)
     setup_docker(app)
@@ -70,8 +73,6 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
 
     setup_auto_scaler_background_task(app)
     setup_buffer_machines_pool_task(app)
-    if app.state.settings.AUTOSCALING_TRACING:
-        initialize_tracing(app, app.state.settings.AUTOSCALING_TRACING, APP_NAME)
 
     # ERROR HANDLERS
 

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

@@ -5,10 +5,10 @@
 from models_library.basic_types import BootModeEnum
 from servicelib.fastapi import timing_middleware
 from servicelib.fastapi.monitoring import (
-    initialize_prometheus_instrumentation,
+    setup_prometheus_instrumentation,
 )
 from servicelib.fastapi.openapi import override_fastapi_openapi_method
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from .._meta import (
@@ -63,10 +63,10 @@ def create_app() -> FastAPI:
     app.state.settings = settings
 
     # MIDDLEWARES
-    if settings.CATALOG_PROMETHEUS_INSTRUMENTATION_ENABLED:
-        initialize_prometheus_instrumentation(app)
     if settings.CATALOG_TRACING:
-        initialize_tracing(app, settings.CATALOG_TRACING, APP_NAME)
+        setup_tracing(app, settings.CATALOG_TRACING, APP_NAME)
+    if settings.CATALOG_PROMETHEUS_INSTRUMENTATION_ENABLED:
+        setup_prometheus_instrumentation(app)
 
     if settings.SC_BOOT_MODE != BootModeEnum.PRODUCTION:
         # middleware to time requests (ONLY for development)

services/clusters-keeper/src/simcore_service_clusters_keeper/core/application.py

@@ -4,7 +4,7 @@
 from servicelib.fastapi.monitoring import (
     setup_prometheus_instrumentation,
 )
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 
 from .._meta import (
     API_VERSION,
@@ -58,14 +58,14 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     app.state.settings = settings
     assert app.state.settings.API_VERSION == API_VERSION  # nosec
 
-    if app.state.settings.CLUSTERS_KEEPER_PROMETHEUS_INSTRUMENTATION_ENABLED:
-        setup_prometheus_instrumentation(app)
     if app.state.settings.CLUSTERS_KEEPER_TRACING:
-        initialize_tracing(
+        setup_tracing(
             app,
             app.state.settings.CLUSTERS_KEEPER_TRACING,
             APP_NAME,
         )
+    if app.state.settings.CLUSTERS_KEEPER_PROMETHEUS_INSTRUMENTATION_ENABLED:
+        setup_prometheus_instrumentation(app)
 
     # PLUGINS SETUP
     setup_api_routes(app)

services/datcore-adapter/src/simcore_service_datcore_adapter/core/application.py

@@ -10,7 +10,7 @@
     setup_prometheus_instrumentation,
 )
 from servicelib.fastapi.openapi import override_fastapi_openapi_method
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 from starlette.middleware.base import BaseHTTPMiddleware
 
 from .._meta import API_VERSION, API_VTAG, APP_NAME
@@ -60,14 +60,14 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
 
     app.state.settings = settings
 
-    if app.state.settings.DATCORE_ADAPTER_PROMETHEUS_INSTRUMENTATION_ENABLED:
-        setup_prometheus_instrumentation(app)
     if app.state.settings.DATCORE_ADAPTER_TRACING:
-        initialize_tracing(
+        setup_tracing(
             app,
             app.state.settings.DATCORE_ADAPTER_TRACING,
             APP_NAME,
         )
+    if app.state.settings.DATCORE_ADAPTER_PROMETHEUS_INSTRUMENTATION_ENABLED:
+        setup_prometheus_instrumentation(app)
 
     if settings.SC_BOOT_MODE != BootModeEnum.PRODUCTION:
         # middleware to time requests (ONLY for development)

services/director-v2/src/simcore_service_director_v2/core/application.py

@@ -7,7 +7,7 @@
     override_fastapi_openapi_method,
 )
 from servicelib.fastapi.profiler import initialize_profiler
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 from servicelib.logging_utils import config_all_loggers
 
 from .._meta import API_VERSION, API_VTAG, APP_NAME, PROJECT_NAME, SUMMARY
@@ -145,12 +145,12 @@ def init_app(settings: AppSettings | None = None) -> FastAPI:
 
     substitutions.setup(app)
 
+    if settings.DIRECTOR_V2_TRACING:
+        setup_tracing(app, settings.DIRECTOR_V2_TRACING, APP_NAME)
+
     if settings.DIRECTOR_V2_PROMETHEUS_INSTRUMENTATION_ENABLED:
         instrumentation.setup(app)
 
-    if settings.DIRECTOR_V2_TRACING:
-        initialize_tracing(app, settings.DIRECTOR_V2_TRACING, APP_NAME)
-
     if settings.DIRECTOR_V0.DIRECTOR_ENABLED:
         director_v0.setup(
             app,

services/director/src/simcore_service_director/core/application.py

@@ -4,7 +4,7 @@
 from fastapi import FastAPI
 from servicelib.async_utils import cancel_sequential_workers
 from servicelib.fastapi.client_session import setup_client_session
-from servicelib.fastapi.tracing import initialize_tracing
+from servicelib.fastapi.tracing import setup_tracing
 
 from .._meta import (
     API_VERSION,
@@ -48,13 +48,13 @@ def create_app(settings: ApplicationSettings) -> FastAPI:
     assert app.state.settings.API_VERSION == API_VERSION  # nosec
 
     # PLUGINS SETUP
+    if app.state.settings.DIRECTOR_TRACING:
+        setup_tracing(app, app.state.settings.DIRECTOR_TRACING, APP_NAME)
+
     setup_api_routes(app)
 
     setup_instrumentation(app)
 
-    if app.state.settings.DIRECTOR_TRACING:
-        initialize_tracing(app, app.state.settings.DIRECTOR_TRACING, APP_NAME)
-
     setup_client_session(
         app,
         max_keepalive_connections=settings.DIRECTOR_REGISTRY_CLIENT_MAX_KEEPALIVE_CONNECTIONS,