From 16f2e6ecee01a6e9c87b34d62f352aa2a35c50af Mon Sep 17 00:00:00 2001
From: Innokentii Konstantinov <innokenty.konstantinov@grafana.com>
Date: Mon, 4 Mar 2024 17:25:39 +0800
Subject: [PATCH] Use uwsgi instrumentation and log request headers (#3997)

# What this PR does
Use uwsgi instead of TracingMiddleware to not to mess up with span
attrs. Also it's an attempt to fix different trace IDs between
cortex-gw, oncall and labels plugin. Last, but not least - add a
middleware to log request headers if OTEL is enabled. It's needed to
debug tracing headers
---
 engine/engine/middlewares.py | 20 ++++++++++++--------
 engine/engine/wsgi.py        |  3 +++
 engine/requirements.in       |  1 +
 engine/requirements.txt      | 13 +++++++++++--
 engine/settings/base.py      |  6 +++---
 5 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/engine/engine/middlewares.py b/engine/engine/middlewares.py
index 82fc64cd..b3c406a6 100644
--- a/engine/engine/middlewares.py
+++ b/engine/engine/middlewares.py
@@ -5,7 +5,6 @@ from django.conf import settings
 from django.core.exceptions import PermissionDenied
 from django.db import OperationalError
 from django.utils.deprecation import MiddlewareMixin
-from opentelemetry import trace
 
 logger = logging.getLogger(__name__)
 
@@ -92,14 +91,19 @@ class BanAlertConsumptionBasedOnSettingsMiddleware(MiddlewareMixin):
             raise PermissionDenied()
 
 
-class TracingMiddleware:
+class LogRequestHeadersMiddleware:
+    """
+    Middleware to log the request headers.
+    Introduced to debug tracing issues.
+    """
+
     def __init__(self, get_response):
         self.get_response = get_response
-        self.tracer = trace.get_tracer(__name__)
 
     def __call__(self, request):
-        with self.tracer.start_as_current_span("middleware"):
-            # TODO: add span attrs
-            # Process the request
-            response = self.get_response(request)
-            return response
+        # Log request headers
+        logger.info("Request Headers: %s", request.headers)
+
+        response = self.get_response(request)
+
+        return response
diff --git a/engine/engine/wsgi.py b/engine/engine/wsgi.py
index 65c8da59..7c202c94 100644
--- a/engine/engine/wsgi.py
+++ b/engine/engine/wsgi.py
@@ -14,6 +14,7 @@ from django.core.wsgi import get_wsgi_application
 from opentelemetry import trace
 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
 from opentelemetry.instrumentation.logging import LoggingInstrumentor
+from opentelemetry.instrumentation.wsgi import OpenTelemetryMiddleware
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from whitenoise import WhiteNoise
@@ -30,6 +31,8 @@ if settings.OTEL_TRACING_ENABLED and settings.OTEL_EXPORTER_OTLP_ENDPOINT:
     try:
         from uwsgidecorators import postfork
 
+        application = OpenTelemetryMiddleware(application)
+
         @postfork
         def init_tracing():
             trace.set_tracer_provider(TracerProvider())
diff --git a/engine/requirements.in b/engine/requirements.in
index a2796fb4..f9ae3e2d 100644
--- a/engine/requirements.in
+++ b/engine/requirements.in
@@ -38,6 +38,7 @@ opentelemetry-sdk==1.23.0
 opentelemetry-api==1.23.0
 opentelemetry-exporter-otlp-proto-grpc==1.15.0
 opentelemetry-instrumentation-logging==0.44b0
+opentelemetry-instrumentation-wsgi==0.44b0
 phonenumbers==8.10.0
 prometheus_client==0.16.0
 psutil==5.9.4
diff --git a/engine/requirements.txt b/engine/requirements.txt
index 0c1ecf1b..97fb6a42 100644
--- a/engine/requirements.txt
+++ b/engine/requirements.txt
@@ -265,13 +265,18 @@ opentelemetry-api==1.23.0
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-instrumentation
     #   opentelemetry-instrumentation-logging
+    #   opentelemetry-instrumentation-wsgi
     #   opentelemetry-sdk
 opentelemetry-exporter-otlp-proto-grpc==1.15.0
     # via -r ./engine/requirements.in
 opentelemetry-instrumentation==0.44b0
-    # via opentelemetry-instrumentation-logging
+    # via
+    #   opentelemetry-instrumentation-logging
+    #   opentelemetry-instrumentation-wsgi
 opentelemetry-instrumentation-logging==0.44b0
     # via -r ./engine/requirements.in
+opentelemetry-instrumentation-wsgi==0.44b0
+    # via -r ./engine/requirements.in
 opentelemetry-proto==1.15.0
     # via opentelemetry-exporter-otlp-proto-grpc
 opentelemetry-sdk==1.23.0
@@ -279,7 +284,11 @@ opentelemetry-sdk==1.23.0
     #   -r ./engine/requirements.in
     #   opentelemetry-exporter-otlp-proto-grpc
 opentelemetry-semantic-conventions==0.44b0
-    # via opentelemetry-sdk
+    # via
+    #   opentelemetry-instrumentation-wsgi
+    #   opentelemetry-sdk
+opentelemetry-util-http==0.44b0
+    # via opentelemetry-instrumentation-wsgi
 pem==23.1.0
     # via django-sns-view
 phonenumbers==8.10.0
diff --git a/engine/settings/base.py b/engine/settings/base.py
index ba848139..b481eaef 100644
--- a/engine/settings/base.py
+++ b/engine/settings/base.py
@@ -345,7 +345,7 @@ MIDDLEWARE = [
 ]
 
 if OTEL_TRACING_ENABLED:
-    MIDDLEWARE.insert(0, "engine.middlewares.TracingMiddleware")
+    MIDDLEWARE.insert(0, "engine.middlewares.LogRequestHeadersMiddleware")
 
 LOG_REQUEST_ID_HEADER = "HTTP_X_CLOUD_TRACE_CONTEXT"
 
@@ -354,8 +354,8 @@ log_fmt = "source=engine:app google_trace_id=%(request_id)s logger=%(name)s %(me
 
 if OTEL_TRACING_ENABLED:
     log_fmt = (
-        "source=engine:app google_trace_id=%(request_id)s logger=%(name)s trace_id=%(otelTraceID)s span_id=%("
-        "otelSpanID)s trace_sampled=%(otelTraceSampled)s %(message)s"
+        "source=engine:app trace_id=%(otelTraceID)s span_id=%("
+        "otelSpanID)s trace_sampled=%(otelTraceSampled)s google_trace_id=%(request_id)s logger=%(name)s %(message)s"
     )
 LOGGING = {
     "version": 1,