Instrument on subscribe and add dcid to span attributes

Author: davidigandan

src/workflows/recipe/__init__.py

@@ -3,6 +3,7 @@
 import functools
 import logging
 from collections.abc import Callable
+from opentelemetry import trace
 from typing import Any
 
 from workflows.recipe.recipe import Recipe
@@ -69,6 +70,39 @@ def unwrap_recipe(header, message):
             message = mangle_for_receiving(message)
         if header.get("workflows-recipe") in {True, "True", "true", 1}:
             rw = RecipeWrapper(message=message, transport=transport_layer)
+            print(rw)
+            logger.log(1, rw)
+
+            # Extract and set DCID on the current span
+            span = trace.get_current_span()
+            dcid = None
+
+            # Try multiple locations where DCID might be stored
+            top_level_params = {}
+            if isinstance(message, dict):
+                # Direct parameters (top-level or in recipe)
+                top_level_params = message.get("parameters", {})
+            
+            # Payload parameters (most common location)
+            payload = message.get("payload", {})
+            payload_params = {}
+            if isinstance(payload, dict):
+                payload_params = payload.get("parameters", {})
+            
+            # Try all common locations
+            dcid = (
+                top_level_params.get("ispyb_dcid") or
+                top_level_params.get("dcid") or
+                payload_params.get("ispyb_dcid") or
+                payload_params.get("dcid") or
+                payload.get("ispyb_dcid") or  
+                payload.get("dcid")
+            )
+
+            if dcid:
+                span.set_attribute("dcid", dcid)
+                span.add_event("recipe.dcid_extracted", attributes={"dcid": dcid})
+        
             if log_extender and rw.environment and rw.environment.get("ID"):
                 with log_extender("recipe_ID", rw.environment["ID"]):
                     return callback(rw, header, message.get("payload"))

src/workflows/transport/middleware/otel_tracing.py

@@ -2,7 +2,7 @@
 from workflows.transport.middleware import BaseTransportMiddleware
 from collections.abc import Callable
 import functools
-from opentelemetry.propagate import inject
+from opentelemetry.propagate import inject, extract
 
 class OTELTracingMiddleware(BaseTransportMiddleware):
     def __init__(self, tracer: trace.Tracer, service_name: str):
@@ -14,29 +14,23 @@ def __init__(self, tracer: trace.Tracer, service_name: str):
         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))
+    def subscribe(self, call_next: Callable, channel, callback, **kwargs) -> int:
+        @functools.wraps(callback)
+        def wrapped_callback(header, message):
+            # Extract trace context from message headers
+            ctx = extract(header) if header else None
 
-            # 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)
-
+            # Start a new span with the extracted context
+            with self.tracer.start_as_current_span(
+                "transport.subscribe",
+                context=ctx
+            ) as span:
+                span.set_attribute("service_name", self.service_name)
+                span.set_attribute("channel", channel)
+        
+                
+                # Call the original callback
+                return callback(header, message)
+        
+        # Call the next middleware with the wrapped callback
+        return call_next(channel, wrapped_callback, **kwargs)