ensure that the context used to detach the token is the same as what was used to attach it

Author: VegetarianOrc

temporalio/contrib/opentelemetry.py

@@ -423,21 +423,21 @@ async def handle_query(self, input: temporalio.worker.HandleQueryInput) -> Any:
         # We do not put anything that happens in a query handler on the workflow
         # span.
         context_header = input.headers.get(self.header_key)
-        context: opentelemetry.context.Context
+        attach_context: opentelemetry.context.Context
         link_context_carrier: Optional[_CarrierDict] = None
         if context_header:
             context_carrier = self.payload_converter.from_payloads([context_header])[0]
-            context = self.text_map_propagator.extract(context_carrier)
+            attach_context = self.text_map_propagator.extract(context_carrier)
             # If there is a workflow span, use it as the link
             link_context_carrier = self._load_workflow_context_carrier()
         else:
             # Use an empty context
-            context = opentelemetry.context.Context()
+            attach_context = opentelemetry.context.Context()
 
         # We need to put this interceptor on the context too
-        context = self._set_on_context(context)
+        attach_context = self._set_on_context(attach_context)
         # Run under context with new span
-        token = opentelemetry.context.attach(context)
+        token = opentelemetry.context.attach(attach_context)
         try:
             # This won't be created if there was no context header
             self._completed_span(
@@ -449,7 +449,9 @@ async def handle_query(self, input: temporalio.worker.HandleQueryInput) -> Any:
             )
             return await super().handle_query(input)
         finally:
-            opentelemetry.context.detach(token)
+            detach_context = opentelemetry.context.get_current()
+            if detach_context is attach_context:
+                opentelemetry.context.detach(token)
 
     def handle_update_validator(
         self, input: temporalio.worker.HandleUpdateInput
@@ -508,19 +510,20 @@ def _top_level_workflow_context(
     ) -> Iterator[None]:
         # Load context only if there is a carrier, otherwise use empty context
         context_carrier = self._load_workflow_context_carrier()
-        context: opentelemetry.context.Context
+        attach_context: opentelemetry.context.Context
         if context_carrier:
-            context = self.text_map_propagator.extract(context_carrier)
+            attach_context = self.text_map_propagator.extract(context_carrier)
         else:
-            context = opentelemetry.context.Context()
+            attach_context = opentelemetry.context.Context()
         # We need to put this interceptor on the context too
-        context = self._set_on_context(context)
+        attach_context = self._set_on_context(attach_context)
         # Need to know whether completed and whether there was a fail-workflow
         # exception
         success = False
         exception: Optional[Exception] = None
         # Run under this context
-        token = opentelemetry.context.attach(context)
+        token = opentelemetry.context.attach(attach_context)
+
         try:
             yield None
             success = True
@@ -537,8 +540,12 @@ def _top_level_workflow_context(
                     exception=exception,
                     kind=opentelemetry.trace.SpanKind.INTERNAL,
                 )
-            opentelemetry.context.detach(token)
 
+            detach_context = opentelemetry.context.get_current()
+            if detach_context is attach_context:
+                opentelemetry.context.detach(token)
+
+    #
     def _context_to_headers(
         self, headers: Mapping[str, temporalio.api.common.v1.Payload]
     ) -> Mapping[str, temporalio.api.common.v1.Payload]:

tests/contrib/test_opentelemetry.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+import sys
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
@@ -21,6 +22,11 @@
 from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
+from tests.worker.test_workflow import (
+    CacheEvictionTearDownWorkflow,
+    WaitForeverWorkflow,
+    wait_forever_activity,
+)
 
 # Passing through because Python 3.9 has an import bug at
 # https://github.com/python/cpython/issues/91351
@@ -321,7 +327,10 @@ def dump_spans(
                 span_links: List[str] = []
                 for link in span.links:
                     for link_span in spans:
-                        if link_span.context.span_id == link.context.span_id:
+                        if (
+                            link_span.context is not None
+                            and link_span.context.span_id == link.context.span_id
+                        ):
                             span_links.append(link_span.name)
                 span_str += f" (links: {', '.join(span_links)})"
             # Signals can duplicate in rare situations, so we make sure not to
@@ -331,7 +340,7 @@ def dump_spans(
             ret.append(span_str)
             ret += dump_spans(
                 spans,
-                parent_id=span.context.span_id,
+                parent_id=span.context.span_id if span.context else None,
                 with_attributes=with_attributes,
                 indent_depth=indent_depth + 1,
             )
@@ -448,3 +457,68 @@ async def test_opentelemetry_benign_exception(client: Client):
 # * workflow failure and wft failure
 # * signal with start
 # * signal failure and wft failure from signal
+
+
+async def test_opentelemetry_safe_detach(client: Client):
+    # This test simulates forcing eviction. This purposely raises GeneratorExit on
+    # GC which triggers the finally which could run on any thread Python
+    # chooses. When this occurs, we should not detach the token from the context
+    # b/c the context no longer exists
+
+    # Create a tracer that has an in-memory exporter
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer = get_tracer(__name__, tracer_provider=provider)
+
+    class _OtelLogSpy(logging.Handler):
+        def __init__(self, level: int | str = 0) -> None:
+            self.seenOtelFailedMessage = False
+            super().__init__(level)
+
+        def emit(self, record: logging.LogRecord) -> None:
+            if not self.seenOtelFailedMessage:
+                self.seenOtelFailedMessage = (
+                    record.levelno == logging.ERROR
+                    and record.name == "opentelemetry.context"
+                    and record.message == "Failed to detach context"
+                )
+
+    async with Worker(
+        client,
+        workflows=[CacheEvictionTearDownWorkflow, WaitForeverWorkflow],
+        activities=[wait_forever_activity],
+        max_cached_workflows=0,
+        task_queue=f"task_queue_{uuid.uuid4()}",
+        disable_safe_workflow_eviction=True,
+        interceptors=[TracingInterceptor(tracer)],
+    ) as worker:
+        # Put a hook to catch unraisable exceptions
+        old_hook = sys.unraisablehook
+        hook_calls: List[sys.UnraisableHookArgs] = []
+        sys.unraisablehook = hook_calls.append
+        log_spy = _OtelLogSpy()
+        logging.getLogger().addHandler(log_spy)
+        try:
+            handle = await client.start_workflow(
+                CacheEvictionTearDownWorkflow.run,
+                id=f"wf-{uuid.uuid4()}",
+                task_queue=worker.task_queue,
+            )
+
+            # CacheEvictionTearDownWorkflow requires 3 signals to be sent
+            await handle.signal(CacheEvictionTearDownWorkflow.signal)
+            await handle.signal(CacheEvictionTearDownWorkflow.signal)
+            await handle.signal(CacheEvictionTearDownWorkflow.signal)
+
+            await handle.result()
+        finally:
+            sys.unraisablehook = old_hook
+            logging.getLogger().removeHandler(log_spy)
+
+        # Confirm at least 1 exception
+        assert hook_calls
+
+        assert (
+            not log_spy.seenOtelFailedMessage
+        ), "Detach from context message should not be logged"