Implement ExitStack() to manage multiple context managers and clean them up

Author: davidigandan

src/workflows/recipe/__init__.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-
+from contextlib import ExitStack
 import functools
 import logging
 from collections.abc import Callable
@@ -70,6 +70,7 @@ def unwrap_recipe(header, message):
         if mangle_for_receiving:
             message = mangle_for_receiving(message)
         if header.get("workflows-recipe") in {True, "True", "true", 1}:
+            otel_logs = None
             rw = RecipeWrapper(message=message, transport=transport_layer)
 
             # Extract recipe ID from environment and add to current span
@@ -79,6 +80,8 @@ def unwrap_recipe(header, message):
             if recipe_id:
                 span.set_attribute("recipe_id", recipe_id)
 
+            
+
             # Extract span_id and trace_id for logging
             span_context = span.get_span_context()
             if span_context and span_context.is_valid:
@@ -92,13 +95,17 @@ def unwrap_recipe(header, message):
 
                 if recipe_id:
                     otel_logs["recipe_id"] = recipe_id
-            else:
-                otel_logs = "No OTEL related logs available"
-
-            if log_extender and rw.environment and rw.environment.get("ID"):
-                with log_extender("recipe_ID", rw.environment["ID"]), log_extender("otel_logs", otel_logs):
-                    return callback(rw, header, message.get("payload"))
-            return callback(rw, header, message.get("payload"))
+        
+            with ExitStack() as stack:
+                # Configure the context depending on if service is emitting spans
+                if otel_logs and log_extender and rw.environment and rw.environment.get("ID"):
+                    stack.enter_context(log_extender('recipe_ID', rw.environment.get("ID")))
+                    stack.enter_context(log_extender('otel_logs', otel_logs))
+                elif log_extender and rw.environment and rw.environment.get("ID"):
+                    stack.enter_context(log_extender('recipe_ID', rw.environment.get("ID")))
+
+                return callback(rw, header, message.get("payload"))
+    
 
         if allow_non_recipe_messages:
             return callback(None, header, message)

src/workflows/services/common_service.py

@@ -197,6 +197,13 @@ def start_transport(self):
             otel_config = (
                 self.config._opentelemetry if self.config and hasattr(self.config, "opentelemetry") else None
             )
+            # debugging
+            with open("/scratch/logs.txt", 'w+') as file:
+                if otel_config:
+                    import json
+                    json.dump(otel_config, file, indent=4)
+                else:
+                    file.write("otel config was not truthy")
 
             if otel_config and "timeout" not in otel_config:
                 self.log.warning("Missing optional OTEL configuration field `timeout`. Will default to 10 seconds. ")

src/workflows/util/zocalo/configuration.py

@@ -33,6 +33,8 @@ def activate(configuration):
         OTEL.config["endpoint"] = endpoint
         OTEL.config["timeout"] = configuration.get("timeout", 10)
 
+        return OTEL.config
+
 
 class Stomp:
     """A Zocalo configuration plugin to pre-populate StompTransport config defaults"""