Experiment with new tracer

Author: tconley1428

temporalio/contrib/opentelemetry.py

@@ -25,6 +25,9 @@
 import opentelemetry.trace
 import opentelemetry.trace.propagation.tracecontext
 import opentelemetry.util.types
+from opentelemetry.context import Context
+from opentelemetry.trace import SpanKind, _Links, Span, StatusCode, Status
+from opentelemetry.util import types
 from typing_extensions import Protocol, TypeAlias, TypedDict
 
 import temporalio.activity
@@ -34,6 +37,7 @@
 import temporalio.exceptions
 import temporalio.worker
 import temporalio.workflow
+from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
 
 # OpenTelemetry dynamically, lazily chooses its context implementation at
 # runtime. When first accessed, they use pkg_resources.iter_entry_points + load.
@@ -661,6 +665,50 @@ def start_local_activity(
         return super().start_local_activity(input)
 
 
+class TemporalTracer(opentelemetry.trace.Tracer):
+    def __init__(self, tracer: opentelemetry.trace.Tracer):
+        self._tracer = tracer
+        super().__init__()
+
+    def start_span(self, name: str, context: Optional[Context] = None, kind: SpanKind = SpanKind.INTERNAL,
+                   attributes: types.Attributes = None, links: _Links = None, start_time: Optional[int] = None,
+                   record_exception: bool = True, set_status_on_exception: bool = True) -> Span:
+        return self._tracer.start_span(name, context, kind, attributes, links, start_time, record_exception, set_status_on_exception)
+
+    @staticmethod
+    def handle_exception(exc: Exception, span: Span, record_exception: bool, set_status_on_exception: bool) -> None:
+        if record_exception:
+            span.record_exception(exc)
+
+        # Set status in case exception was raised
+        if set_status_on_exception:
+            span.set_status(
+                Status(
+                    status_code=StatusCode.ERROR,
+                    description=f"{type(exc).__name__}: {exc}",
+                )
+            )
+
+    @contextmanager
+    def start_as_current_span(self, name: str, context: Optional[Context] = None, kind: SpanKind = SpanKind.INTERNAL,
+                              attributes: types.Attributes = None, links: _Links = None,
+                              start_time: Optional[int] = None, record_exception: bool = True,
+                              set_status_on_exception: bool = True, end_on_exit: bool = True) -> Iterator[Span]:
+
+        with self._tracer.start_as_current_span(name, context, kind, attributes, links, start_time, False, False, end_on_exit) as span:
+            try:
+                yield span
+
+            # TODO: Catch base exception and handle cancellation errors
+            except ApplicationError as exc:
+                if exc.category != ApplicationErrorCategory.BENIGN:
+                    self.handle_exception(exc, span, record_exception, set_status_on_exception)
+                raise
+            except Exception as exc:
+                self.handle_exception(exc, span, record_exception, set_status_on_exception)
+                raise
+
+
 class workflow:
     """Contains static methods that are safe to call from within a workflow.
 

tests/contrib/test_opentelemetry.py

@@ -3,20 +3,22 @@
 import asyncio
 import logging
 import uuid
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Iterable, List, Optional
 
 from opentelemetry.sdk.trace import ReadableSpan, TracerProvider
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
-from opentelemetry.trace import get_tracer
+from opentelemetry.trace import get_tracer, StatusCode
 
 from temporalio import activity, workflow
 from temporalio.client import Client
 from temporalio.common import RetryPolicy
-from temporalio.contrib.opentelemetry import TracingInterceptor
+from temporalio.contrib.opentelemetry import TracingInterceptor, TemporalTracer
 from temporalio.contrib.opentelemetry import workflow as otel_workflow
+from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
 from temporalio.testing import WorkflowEnvironment
 from temporalio.worker import UnsandboxedWorkflowRunner, Worker
 
@@ -386,6 +388,49 @@ async def test_opentelemetry_always_create_workflow_spans(client: Client):
     assert spans[0].name == "RunWorkflow:SimpleWorkflow"
 
 
+attempted = False
+@activity.defn
+def benign_activity() -> str:
+    global attempted
+    if attempted:
+        return "done"
+    print("Raising")
+    attempted = True
+    raise ApplicationError(category=ApplicationErrorCategory.BENIGN, message="Benign Error")
+
+@workflow.defn
+class BenignWorkflow:
+    @workflow.run
+    async def run(self) -> str:
+       return await workflow.execute_activity(benign_activity, schedule_to_close_timeout=timedelta(seconds=1))
+
+async def test_opentelemetry_exception_tracing(client: Client):
+    # Create a tracer that has an in-memory exporter
+    exporter = InMemorySpanExporter()
+    provider = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    tracer = TemporalTracer(get_tracer(__name__, tracer_provider=provider))
+    # Create new client with tracer interceptor
+    client_config = client.config()
+    client_config["interceptors"] = [TracingInterceptor(tracer)]
+    client = Client(**client_config)
+
+    async with Worker(
+        client,
+        task_queue=f"task_queue_{uuid.uuid4()}",
+        workflows=[BenignWorkflow],
+        activities=[benign_activity],
+        activity_executor=ThreadPoolExecutor(max_workers=1),
+    ) as worker:
+        assert "done" == await client.execute_workflow(
+            BenignWorkflow.run,
+            id=f"workflow_{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            retry_policy=RetryPolicy(maximum_attempts=2, initial_interval=timedelta(milliseconds=10)),
+        )
+    spans = exporter.get_finished_spans()
+    assert all(span.status.status_code == StatusCode.UNSET for span in spans)
+
 # TODO(cretz): Additional tests to write
 # * query without interceptor (no headers)
 # * workflow without interceptor (no headers) but query with interceptor (headers)