add osparc-trace-id response header to api-server and webserver

Author: bisgaard-itis

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

@@ -18,6 +18,7 @@
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from servicelib.logging_utils import log_context
+from servicelib.tracing import get_trace_id_header
 from settings_library.tracing import TracingSettings
 from yarl import URL
 
@@ -55,6 +56,7 @@ def _startup(
     app: web.Application,
     tracing_settings: TracingSettings,
     service_name: str,
+    add_response_trace_id_header: bool = False,
 ) -> None:
     """
     Sets up this service for a distributed tracing system (opentelemetry)
@@ -106,6 +108,8 @@ def _startup(
     #
     # Since the code that is provided (monkeypatched) in the __init__ that the opentelemetry-autoinstrumentation-library provides is only 4 lines,
     # just adding a middleware, we are free to simply execute this "missed call" [since we can't call the monkeypatch'ed __init__()] in this following line:
+    if add_response_trace_id_header:
+        app.middlewares.insert(0, ResponseTraceIdHeaderMiddleware)
     app.middlewares.insert(0, aiohttp_server_opentelemetry_middleware)
     # Code of the aiohttp server instrumentation: github.com/open-telemetry/opentelemetry-python-contrib/blob/eccb05c808a7d797ef5b6ecefed3590664426fbf/instrumentation/opentelemetry-instrumentation-aiohttp-server/src/opentelemetry/instrumentation/aiohttp_server/__init__.py#L246
     # For reference, the above statement was written for:
@@ -146,6 +150,15 @@ def _startup(
             AioPikaInstrumentor().instrument()
 
 
+@web.middleware
+async def ResponseTraceIdHeaderMiddleware(request: web.Request, handler):
+    response = await handler(request)
+    trace_id_header = get_trace_id_header()
+    if trace_id_header:
+        response.headers.update(trace_id_header)
+    return response
+
+
 def _shutdown() -> None:
     """Uninstruments all opentelemetry instrumentors that were instrumented."""
     try:
@@ -175,9 +188,17 @@ def _shutdown() -> None:
 
 
 def get_tracing_lifespan(
-    app: web.Application, tracing_settings: TracingSettings, service_name: str
+    app: web.Application,
+    tracing_settings: TracingSettings,
+    service_name: str,
+    add_response_trace_id_header: bool = False,
 ) -> Callable[[web.Application], AsyncIterator]:
-    _startup(app=app, tracing_settings=tracing_settings, service_name=service_name)
+    _startup(
+        app=app,
+        tracing_settings=tracing_settings,
+        service_name=service_name,
+        add_response_trace_id_header=add_response_trace_id_header,
+    )
 
     async def tracing_lifespan(app: web.Application):
         assert app  # nosec

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

@@ -3,7 +3,7 @@
 import logging
 from collections.abc import AsyncIterator
 
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 from fastapi_lifespan_manager import State
 from httpx import AsyncClient, Client
 from opentelemetry import trace
@@ -16,7 +16,9 @@
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from servicelib.logging_utils import log_context
+from servicelib.tracing import get_trace_id_header
 from settings_library.tracing import TracingSettings
+from starlette.middleware.base import BaseHTTPMiddleware
 from yarl import URL
 
 _logger = logging.getLogger(__name__)
@@ -180,7 +182,11 @@ def _shutdown() -> None:
             _logger.exception("Failed to uninstrument RequestsInstrumentor")
 
 
-def initialize_fastapi_app_tracing(app: FastAPI):
+def initialize_fastapi_app_tracing(
+    app: FastAPI, *, add_response_trace_id_header: bool = False
+):
+    if add_response_trace_id_header:
+        app.add_middleware(ResponseTraceIdHeaderMiddleware)
     FastAPIInstrumentor.instrument_app(app)
 
 
@@ -216,3 +222,13 @@ async def tracing_instrumentation_lifespan(
         _shutdown()
 
     return tracing_instrumentation_lifespan
+
+
+class ResponseTraceIdHeaderMiddleware(BaseHTTPMiddleware):
+
+    async def dispatch(self, request: Request, call_next):
+        response = await call_next(request)
+        trace_id_header = get_trace_id_header()
+        if trace_id_header:
+            response.headers.update(trace_id_header)
+        return response

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

@@ -8,6 +8,8 @@
 
 TracingContext: TypeAlias = otcontext.Context | None
 
+_OSPARC_TRACE_ID_HEADER = "X-Osparc-Trace-Id"
+
 
 def _is_tracing() -> bool:
     return trace.get_current_span().is_recording()
@@ -34,3 +36,15 @@ def use_tracing_context(context: TracingContext):
 def setup_log_tracing(tracing_settings: TracingSettings):
     _ = tracing_settings
     LoggingInstrumentor().instrument(set_logging_format=False)
+
+
+def get_trace_id_header() -> dict[str, str] | None:
+    """Generates a dictionary containing the trace ID header if tracing is active."""
+    span = trace.get_current_span()
+    if span.is_recording():
+        trace_id = span.get_span_context().trace_id
+        trace_id_hex = format(
+            trace_id, "032x"
+        )  # Convert trace_id to 32-character hex string
+        return {_OSPARC_TRACE_ID_HEADER: trace_id_hex}
+    return None

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

@@ -97,7 +97,7 @@ def init_app(settings: ApplicationSettings | None = None) -> FastAPI:
         setup_prometheus_instrumentation(app)
 
     if settings.API_SERVER_TRACING:
-        initialize_fastapi_app_tracing(app)
+        initialize_fastapi_app_tracing(app, add_response_trace_id_header=True)
 
     if settings.API_SERVER_WEBSERVER:
         webserver.setup(

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

@@ -28,5 +28,6 @@ def setup_app_tracing(app: web.Application):
             app,
             tracing_settings=tracing_settings,
             service_name=APP_NAME,
+            add_response_trace_id_header=True,
         )
     )