add lifespan management to tracing instrumentation in aiohttp

Author: bisgaard-itis

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

@@ -1,6 +1,7 @@
 """Adds aiohttp middleware for tracing using opentelemetry instrumentation."""
 
 import logging
+from collections.abc import AsyncIterator, Callable
 
 from aiohttp import web
 from opentelemetry import trace
@@ -50,7 +51,7 @@
     HAS_AIO_PIKA = False
 
 
-def setup_tracing(
+def _startup(
     app: web.Application,
     tracing_settings: TracingSettings,
     service_name: str,
@@ -143,3 +144,43 @@ def setup_tracing(
             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:
+        _logger.exception("Failed to uninstrument AioHttpClientInstrumentor")
+    if HAS_AIOPG:
+        try:
+            AiopgInstrumentor().uninstrument()
+        except Exception:
+            _logger.exception("Failed to uninstrument AiopgInstrumentor")
+    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")
+    if HAS_AIO_PIKA:
+        try:
+            AioPikaInstrumentor().uninstrument()
+        except Exception:
+            _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):
+        yield
+        _shutdown()
+
+    return tracing_lifespan

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
 
 
@@ -59,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(
@@ -137,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

services/web/server/src/simcore_service_webserver/tracing.py

@@ -2,7 +2,7 @@
 
 from aiohttp import web
 from servicelib.aiohttp.application_setup import ModuleCategory, app_module_setup
-from servicelib.aiohttp.tracing import setup_tracing
+from servicelib.aiohttp.tracing import get_tracing_lifespan
 from settings_library.tracing import TracingSettings
 
 from ._meta import APP_NAME
@@ -23,8 +23,10 @@ def get_plugin_settings(app: web.Application) -> TracingSettings:
 )
 def setup_app_tracing(app: web.Application):
     tracing_settings: TracingSettings = get_plugin_settings(app)
-    setup_tracing(
-        app,
-        tracing_settings=tracing_settings,
-        service_name=APP_NAME,
+    app.cleanup_ctx.append(
+        get_tracing_lifespan(
+            app,
+            tracing_settings=tracing_settings,
+            service_name=APP_NAME,
+        )
     )