Add basic tracing middleware and global control

Author: davidigandan

src/workflows/services/common_service.py

@@ -12,6 +12,13 @@
 import workflows
 import workflows.logging
 
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from workflows.transport.middleware.otel_tracing import OTELTracingMiddleware
+from opentelemetry.sdk.resources import Resource, SERVICE_NAME
+
 
 class Status(enum.Enum):
     """
@@ -185,6 +192,31 @@ def start_transport(self):
             self.transport.subscription_callback_set_intercept(
                 self._transport_interceptor
             )
+            
+            # Configure OTELTracing
+            resource = Resource.create({
+                SERVICE_NAME: self._service_name,
+            })
+
+            self.log.debug("Configuring OTELTracing")
+            provider = TracerProvider(resource=resource)
+            trace.set_tracer_provider(provider)
+
+            # Configure BatchProcessor and OTLPSpanExporter to point to OTELCollector
+            otlp_exporter = OTLPSpanExporter(
+                endpoint="https://otel.tracing.diamond.ac.uk:4318/v1/traces",
+                timeout=10
+            )
+            span_processor = BatchSpanProcessor(otlp_exporter)
+            provider.add_span_processor(span_processor)
+
+            # Add OTELTracingMiddleware to the transport layer
+            tracer = trace.get_tracer(__name__)
+            otel_middleware = OTELTracingMiddleware(tracer, service_name=self._service_name)
+            self._transport.add_middleware(otel_middleware)
+
+            self.log.debug("OTELTracingMiddleware added to transport layer of %s", self._service_name)
+
             metrics = self._environment.get("metrics")
             if metrics:
                 import prometheus_client

src/workflows/transport/middleware/otel_tracing.py

@@ -0,0 +1,42 @@
+from opentelemetry import trace
+from workflows.transport.middleware import BaseTransportMiddleware
+from collections.abc import Callable
+import functools
+from opentelemetry.propagate import inject
+
+class OTELTracingMiddleware(BaseTransportMiddleware):
+    def __init__(self, tracer: trace.Tracer, service_name: str):
+        """
+        Initialize the OpenTelemetry Tracing Middleware.
+
+        :param tracer: An OpenTelemetry tracer instance used to create spans.
+        """
+        self.tracer = tracer
+        self.service_name = service_name
+
+
+    def send(self, call_next: Callable, destination, message, **kwargs):
+        """
+        Middleware for tracing the `send` operation
+        
+        :param call_next: The next middleware or the original `send` method.
+        :param destination: The destination service to which the message is being sent. 
+        :param message: The message being sent.
+        :param kwargs: Additional arguments for the `send` method. 
+        """
+
+        # Start a new span for the `send` operation
+        with self.tracer.start_as_current_span("transport.send") as span:
+            # Attributes we're interested in
+            span.set_attribute("service_name", self.service_name)
+            span.set_attribute("destination", destination)
+            span.set_attribute("message", str(message))
+            
+            # Inject trace context into message headers
+            headers = kwargs.setdefault("headers", {})
+            inject(headers)
+            kwargs["headers"] = headers
+
+            # Call the next middleware or the original `send` method
+            return call_next(destination, message, **kwargs)
+