✨ Add lifespan handling to tracing instrumentation and instrument rabbitmq rpc clients

packages/aws-library/requirements/_base.txt

@@ -145,6 +145,7 @@ opentelemetry-api==1.30.0
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-instrumentation
+    #   opentelemetry-instrumentation-aio-pika
     #   opentelemetry-instrumentation-botocore
     #   opentelemetry-instrumentation-logging
     #   opentelemetry-instrumentation-redis
@@ -164,10 +165,13 @@ opentelemetry-exporter-otlp-proto-http==1.30.0
     # via opentelemetry-exporter-otlp
 opentelemetry-instrumentation==0.51b0
     # via
+    #   opentelemetry-instrumentation-aio-pika
     #   opentelemetry-instrumentation-botocore
     #   opentelemetry-instrumentation-logging
     #   opentelemetry-instrumentation-redis
     #   opentelemetry-instrumentation-requests
+opentelemetry-instrumentation-aio-pika==0.51b0
+    # via -r requirements/../../../packages/service-library/requirements/_base.in
 opentelemetry-instrumentation-botocore==0.51b0
     # via -r requirements/_base.in
 opentelemetry-instrumentation-logging==0.51b0
@@ -431,6 +435,7 @@ wrapt==1.17.2
     #   aiobotocore
     #   deprecated
     #   opentelemetry-instrumentation
+    #   opentelemetry-instrumentation-aio-pika
     #   opentelemetry-instrumentation-redis
 yarl==1.18.3
     # via

packages/service-library/requirements/_base.in

@@ -18,6 +18,7 @@ arrow # date/time
 faststream
 opentelemetry-api
 opentelemetry-exporter-otlp
+opentelemetry-instrumentation-aio-pika
 opentelemetry-instrumentation-logging
 opentelemetry-instrumentation-redis
 opentelemetry-instrumentation-requests

packages/service-library/requirements/_base.txt

@@ -103,6 +103,7 @@ opentelemetry-api==1.30.0
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-exporter-otlp-proto-http
     #   opentelemetry-instrumentation
+    #   opentelemetry-instrumentation-aio-pika
     #   opentelemetry-instrumentation-logging
     #   opentelemetry-instrumentation-redis
     #   opentelemetry-instrumentation-requests
@@ -120,9 +121,12 @@ opentelemetry-exporter-otlp-proto-http==1.30.0
     # via opentelemetry-exporter-otlp
 opentelemetry-instrumentation==0.51b0
     # via
+    #   opentelemetry-instrumentation-aio-pika
     #   opentelemetry-instrumentation-logging
     #   opentelemetry-instrumentation-redis
     #   opentelemetry-instrumentation-requests
+opentelemetry-instrumentation-aio-pika==0.51b0
+    # via -r requirements/_base.in
 opentelemetry-instrumentation-logging==0.51b0
     # via -r requirements/_base.in
 opentelemetry-instrumentation-redis==0.51b0
@@ -297,6 +301,7 @@ wrapt==1.17.2
     # via
     #   deprecated
     #   opentelemetry-instrumentation
+    #   opentelemetry-instrumentation-aio-pika
     #   opentelemetry-instrumentation-redis
 yarl==1.18.3
     # via

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

@@ -1,8 +1,7 @@
-""" Adds aiohttp middleware for tracing using opentelemetry instrumentation.
-
-"""
+"""Adds aiohttp middleware for tracing using opentelemetry instrumentation."""
 
 import logging
+from collections.abc import AsyncIterator, Callable
 
 from aiohttp import web
 from opentelemetry import trace
@@ -44,8 +43,15 @@
 except ImportError:
     HAS_REQUESTS = False
 
+try:
+    from opentelemetry.instrumentation.aio_pika import AioPikaInstrumentor
+
+    HAS_AIO_PIKA = True
+except ImportError:
+    HAS_AIO_PIKA = False
+
 
-def setup_tracing(
+def _startup(
     app: web.Application,
     tracing_settings: TracingSettings,
     service_name: str,
@@ -74,7 +80,9 @@ def setup_tracing(
     trace.set_tracer_provider(TracerProvider(resource=resource))
     tracer_provider: trace.TracerProvider = trace.get_tracer_provider()
 
-    tracing_destination: str = f"{URL(opentelemetry_collector_endpoint).with_port(opentelemetry_collector_port).with_path('/v1/traces')}"
+    tracing_destination: str = (
+        f"{URL(opentelemetry_collector_endpoint).with_port(opentelemetry_collector_port).with_path('/v1/traces')}"
+    )
 
     _logger.info(
         "Trying to connect service %s to tracing collector at %s.",
@@ -128,3 +136,52 @@ def setup_tracing(
             msg="Attempting to add requests opentelemetry autoinstrumentation...",
         ):
             RequestsInstrumentor().instrument()
+
+    if HAS_AIO_PIKA:
+        with log_context(
+            _logger,
+            logging.INFO,
+            msg="Attempting to add aio_pika opentelemetry autoinstrumentation...",
+        ):
+            AioPikaInstrumentor().instrument()
+
+
+def _shutdown() -> None:
+    """Uninstruments all opentelemetry instrumentors that were instrumented."""
+    try:
+        AioHttpClientInstrumentor().uninstrument()
+    except Exception:  # pylint:disable=broad-exception-caught
+        _logger.exception("Failed to uninstrument AioHttpClientInstrumentor")
+    if HAS_AIOPG:
+        try:
+            AiopgInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument AiopgInstrumentor")
+    if HAS_BOTOCORE:
+        try:
+            BotocoreInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument BotocoreInstrumentor")
+    if HAS_REQUESTS:
+        try:
+            RequestsInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument RequestsInstrumentor")
+    if HAS_AIO_PIKA:
+        try:
+            AioPikaInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument AioPikaInstrumentor")
+
+
+def get_tracing_lifespan(
+    app: web.Application, tracing_settings: TracingSettings, service_name: str
+) -> Callable[[web.Application], AsyncIterator]:
+    _startup(app=app, tracing_settings=tracing_settings, service_name=service_name)
+
+    async def tracing_lifespan(app: web.Application):
+        assert app  # nosec
+        yield
+        _shutdown()
+
+    return tracing_lifespan

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

@@ -1,10 +1,10 @@
-""" Adds fastapi middleware for tracing using opentelemetry instrumentation.
-
-"""
+"""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 (
@@ -60,10 +60,17 @@
 except ImportError:
     HAS_REQUESTS = False
 
+try:
+    from opentelemetry.instrumentation.aio_pika.aio_pika_instrumentor import (
+        AioPikaInstrumentor,
+    )
 
-def initialize_tracing(
-    app: FastAPI, tracing_settings: TracingSettings, service_name: str
-) -> None:
+    HAS_AIOPIKA_INSTRUMENTOR = True
+except ImportError:
+    HAS_AIOPIKA_INSTRUMENTOR = False
+
+
+def _startup(tracing_settings: TracingSettings, service_name: str) -> None:
     if (
         not tracing_settings.TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT
         and not tracing_settings.TRACING_OPENTELEMETRY_COLLECTOR_PORT
@@ -80,7 +87,9 @@ def initialize_tracing(
         f"{tracing_settings.TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT}"
     )
 
-    tracing_destination: str = f"{URL(opentelemetry_collector_endpoint).with_port(tracing_settings.TRACING_OPENTELEMETRY_COLLECTOR_PORT).with_path('/v1/traces')}"
+    tracing_destination: str = (
+        f"{URL(opentelemetry_collector_endpoint).with_port(tracing_settings.TRACING_OPENTELEMETRY_COLLECTOR_PORT).with_path('/v1/traces')}"
+    )
 
     _logger.info(
         "Trying to connect service %s to opentelemetry tracing collector at %s.",
@@ -91,16 +100,22 @@ def initialize_tracing(
     otlp_exporter = OTLPSpanExporterHTTP(endpoint=tracing_destination)
     span_processor = BatchSpanProcessor(otlp_exporter)
     global_tracer_provider.add_span_processor(span_processor)
-    # Instrument FastAPI
-    FastAPIInstrumentor().instrument_app(app)
 
+    FastAPIInstrumentor().instrument()
     if HAS_AIOPG:
         with log_context(
             _logger,
             logging.INFO,
             msg="Attempting to add asyncpg opentelemetry autoinstrumentation...",
         ):
             AiopgInstrumentor().instrument()
+    if HAS_AIOPIKA_INSTRUMENTOR:
+        with log_context(
+            _logger,
+            logging.INFO,
+            msg="Attempting to add aio_pika opentelemetry autoinstrumentation...",
+        ):
+            AioPikaInstrumentor().instrument()
     if HAS_ASYNCPG:
         with log_context(
             _logger,
@@ -131,5 +146,70 @@ def initialize_tracing(
             RequestsInstrumentor().instrument()
 
 
+def _shutdown() -> None:
+    """Uninstruments all opentelemetry instrumentors that were instrumented."""
+    FastAPIInstrumentor().uninstrument()
+    if HAS_AIOPG:
+        try:
+            AiopgInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument AiopgInstrumentor")
+    if HAS_AIOPIKA_INSTRUMENTOR:
+        try:
+            AioPikaInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument AioPikaInstrumentor")
+    if HAS_ASYNCPG:
+        try:
+            AsyncPGInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument AsyncPGInstrumentor")
+    if HAS_REDIS:
+        try:
+            RedisInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument RedisInstrumentor")
+    if HAS_BOTOCORE:
+        try:
+            BotocoreInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _logger.exception("Failed to uninstrument BotocoreInstrumentor")
+    if HAS_REQUESTS:
+        try:
+            RequestsInstrumentor().uninstrument()
+        except Exception:  # pylint:disable=broad-exception-caught
+            _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(tracing_settings=tracing_settings, service_name=service_name)
+
+    def _on_shutdown() -> None:
+        _shutdown()
+
+    app.add_event_handler("shutdown", _on_shutdown)
+
+
+def get_tracing_instrumentation_lifespan(
+    tracing_settings: TracingSettings, service_name: str
+):
+
+    _startup(tracing_settings=tracing_settings, service_name=service_name)
+
+    async def tracing_instrumentation_lifespan(
+        app: FastAPI,
+    ) -> AsyncIterator[State]:
+        assert app  # nosec
+
+        yield {}
+
+        _shutdown()
+
+    return tracing_instrumentation_lifespan

packages/service-library/src/servicelib/rabbitmq/_client_base.py

@@ -70,5 +70,4 @@ async def ping(self) -> bool:
         return False
 
     @abstractmethod
-    async def close(self) -> None:
-        ...
+    async def close(self) -> None: ...

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

@@ -11,7 +11,7 @@
 from aiohttp import web
 from aiohttp.test_utils import TestClient
 from pydantic import ValidationError
-from servicelib.aiohttp.tracing import setup_tracing
+from servicelib.aiohttp.tracing import get_tracing_lifespan
 from settings_library.tracing import TracingSettings
 
 
@@ -24,14 +24,23 @@ def tracing_settings_in(request):
 def set_and_clean_settings_env_vars(
     monkeypatch: pytest.MonkeyPatch, tracing_settings_in
 ):
+    endpoint_mocked = False
     if tracing_settings_in[0]:
+        endpoint_mocked = True
         monkeypatch.setenv(
             "TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT", f"{tracing_settings_in[0]}"
         )
+    port_mocked = False
     if tracing_settings_in[1]:
+        port_mocked = True
         monkeypatch.setenv(
             "TRACING_OPENTELEMETRY_COLLECTOR_PORT", f"{tracing_settings_in[1]}"
         )
+    yield
+    if endpoint_mocked:
+        monkeypatch.delenv("TRACING_OPENTELEMETRY_COLLECTOR_ENDPOINT")
+    if port_mocked:
+        monkeypatch.delenv("TRACING_OPENTELEMETRY_COLLECTOR_PORT")
 
 
 @pytest.mark.parametrize(
@@ -50,11 +59,10 @@ async def test_valid_tracing_settings(
     app = web.Application()
     service_name = "simcore_service_webserver"
     tracing_settings = TracingSettings()
-    setup_tracing(
-        app,
-        service_name=service_name,
-        tracing_settings=tracing_settings,
-    )
+    async for _ in get_tracing_lifespan(
+        app, service_name=service_name, tracing_settings=tracing_settings
+    )(app):
+        pass
 
 
 @pytest.mark.parametrize(
@@ -128,14 +136,15 @@ async def test_tracing_setup_package_detection(
     app = web.Application()
     service_name = "simcore_service_webserver"
     tracing_settings = TracingSettings()
-    setup_tracing(
-        app,
-        service_name=service_name,
-        tracing_settings=tracing_settings,
-    )
-    # idempotency
-    setup_tracing(
+    async for _ in get_tracing_lifespan(
         app,
         service_name=service_name,
         tracing_settings=tracing_settings,
-    )
+    )(app):
+        # idempotency
+        async for _ in get_tracing_lifespan(
+            app,
+            service_name=service_name,
+            tracing_settings=tracing_settings,
+        )(app):
+            pass