Limit to only handling in our own spans

Author: tconley1428

temporalio/contrib/opentelemetry.py

@@ -171,11 +171,30 @@ def _start_as_current_span(
         attributes: opentelemetry.util.types.Attributes,
         input: Optional[_InputWithHeaders] = None,
         kind: opentelemetry.trace.SpanKind,
+        context: Optional[Context] = None,
     ) -> Iterator[None]:
-        with self.tracer.start_as_current_span(name, attributes=attributes, kind=kind):
+        with self.tracer.start_as_current_span(name, attributes=attributes, kind=kind, context=context, set_status_on_exception=False) as span:
             if input:
                 input.headers = self._context_to_headers(input.headers)
-            yield None
+            try:
+                yield None
+            except ApplicationError as exc:
+                if exc.category != ApplicationErrorCategory.BENIGN:
+                    span.set_status(
+                        Status(
+                            status_code=StatusCode.ERROR,
+                            description=f"{type(exc).__name__}: {exc}",
+                        )
+                    )
+                raise
+            except Exception as exc:
+                span.set_status(
+                    Status(
+                        status_code=StatusCode.ERROR,
+                        description=f"{type(exc).__name__}: {exc}",
+                    )
+                )
+                raise
 
     def _completed_workflow_span(
         self, params: _CompletedWorkflowSpanParams
@@ -286,7 +305,7 @@ async def execute_activity(
         self, input: temporalio.worker.ExecuteActivityInput
     ) -> Any:
         info = temporalio.activity.info()
-        with self.root.tracer.start_as_current_span(
+        with self.root._start_as_current_span(
             f"RunActivity:{info.activity_type}",
             context=self.root._context_from_headers(input.headers),
             attributes={
@@ -665,50 +684,6 @@ 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

@@ -16,7 +16,7 @@
 from temporalio import activity, workflow
 from temporalio.client import Client
 from temporalio.common import RetryPolicy
-from temporalio.contrib.opentelemetry import TracingInterceptor, TemporalTracer
+from temporalio.contrib.opentelemetry import TracingInterceptor
 from temporalio.contrib.opentelemetry import workflow as otel_workflow
 from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
 from temporalio.testing import WorkflowEnvironment
@@ -394,7 +394,6 @@ def benign_activity() -> str:
     global attempted
     if attempted:
         return "done"
-    print("Raising")
     attempted = True
     raise ApplicationError(category=ApplicationErrorCategory.BENIGN, message="Benign Error")
 
@@ -404,12 +403,13 @@ class BenignWorkflow:
     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):
+async def test_opentelemetry_benign_exception(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))
+    tracer = get_tracer(__name__, tracer_provider=provider)
+
     # Create new client with tracer interceptor
     client_config = client.config()
     client_config["interceptors"] = [TracingInterceptor(tracer)]
@@ -429,6 +429,8 @@ async def test_opentelemetry_exception_tracing(client: Client):
             retry_policy=RetryPolicy(maximum_attempts=2, initial_interval=timedelta(milliseconds=10)),
         )
     spans = exporter.get_finished_spans()
+    for span in spans:
+        print(f"{span.name}: {span.status.status_code}")
     assert all(span.status.status_code == StatusCode.UNSET for span in spans)
 
 # TODO(cretz): Additional tests to write