Record exception on context manager exit

This updates the tracer context manager to automatically record
exceptions as events on exit if an exception was raised within the
context manager's context.

Author: owais

instrumentation/opentelemetry-instrumentation-requests/src/opentelemetry/instrumentation/requests/__init__.py

@@ -50,7 +50,11 @@
 from opentelemetry.instrumentation.requests.version import __version__
 from opentelemetry.instrumentation.utils import http_status_to_canonical_code
 from opentelemetry.trace import SpanKind, get_tracer
-from opentelemetry.trace.status import Status, StatusCanonicalCode
+from opentelemetry.trace.status import (
+    EXCEPTION_STATUS_FIELD,
+    Status,
+    StatusCanonicalCode,
+)
 
 # A key to a context variable to avoid creating duplicate spans when instrumenting
 # both, Session.request and Session.send, since Session.request calls into Session.send
@@ -121,8 +125,6 @@ def _instrumented_requests_call(
         method = method.upper()
         span_name = "HTTP {}".format(method)
 
-        exception = None
-
         recorder = RequestsInstrumentor().metric_recorder
 
         labels = {}
@@ -132,6 +134,7 @@ def _instrumented_requests_call(
         with get_tracer(
             __name__, __version__, tracer_provider
         ).start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
+            exception = None
             with recorder.record_duration(labels):
                 if span.is_recording():
                     span.set_attribute("component", "http")
@@ -150,16 +153,16 @@ def _instrumented_requests_call(
                     result = call_wrapped()  # *** PROCEED
                 except Exception as exc:  # pylint: disable=W0703
                     exception = exc
+                    setattr(
+                        exception,
+                        EXCEPTION_STATUS_FIELD,
+                        _exception_to_canonical_code(exception),
+                    )
                     result = getattr(exc, "response", None)
+
                 finally:
                     context.detach(token)
 
-                if exception is not None and span.is_recording():
-                    span.set_status(
-                        Status(_exception_to_canonical_code(exception))
-                    )
-                    span.record_exception(exception)
-
                 if result is not None:
                     if span.is_recording():
                         span.set_attribute(
@@ -184,8 +187,8 @@ def _instrumented_requests_call(
                 if span_callback is not None:
                     span_callback(span, result)
 
-        if exception is not None:
-            raise exception.with_traceback(exception.__traceback__)
+            if exception is not None:
+                raise exception.with_traceback(exception.__traceback__)
 
         return result
 

opentelemetry-api/src/opentelemetry/trace/__init__.py

@@ -282,6 +282,7 @@ def start_as_current_span(
         kind: SpanKind = SpanKind.INTERNAL,
         attributes: types.Attributes = None,
         links: typing.Sequence[Link] = (),
+        record_exception: bool = True,
     ) -> typing.Iterator["Span"]:
         """Context manager for creating a new span and set it
         as the current span in this tracer's context.
@@ -310,7 +311,7 @@ def start_as_current_span(
         is equivalent to::
 
             span = tracer.start_span(name)
-            with tracer.use_span(span, end_on_exit=True):
+            with tracer.use_span(span, end_on_exit=True, record_exception=True):
                 do_work()
 
         Args:
@@ -320,6 +321,8 @@ def start_as_current_span(
                 meaningful even if there is no parent.
             attributes: The span's attributes.
             links: Links span to other spans
+            record_exception: Whether to record any exceptions raised within the
+                context as error event on the span.
 
         Yields:
             The newly-created span.
@@ -328,7 +331,10 @@ def start_as_current_span(
     @contextmanager  # type: ignore
     @abc.abstractmethod
     def use_span(
-        self, span: "Span", end_on_exit: bool = False
+        self,
+        span: "Span",
+        end_on_exit: bool = False,
+        record_exception: bool = True,
     ) -> typing.Iterator[None]:
         """Context manager for setting the passed span as the
         current span in the context, as well as resetting the
@@ -345,6 +351,8 @@ def use_span(
             span: The span to start and make current.
             end_on_exit: Whether to end the span automatically when leaving the
                 context manager.
+            record_exception: Whether to record any exceptions raised within the
+                context as error event on the span.
         """
 
 
@@ -375,13 +383,17 @@ def start_as_current_span(
         kind: SpanKind = SpanKind.INTERNAL,
         attributes: types.Attributes = None,
         links: typing.Sequence[Link] = (),
+        record_exception: bool = True,
     ) -> typing.Iterator["Span"]:
         # pylint: disable=unused-argument,no-self-use
         yield INVALID_SPAN
 
     @contextmanager  # type: ignore
     def use_span(
-        self, span: "Span", end_on_exit: bool = False
+        self,
+        span: "Span",
+        end_on_exit: bool = False,
+        record_exception: bool = True,
     ) -> typing.Iterator[None]:
         # pylint: disable=unused-argument,no-self-use
         yield

opentelemetry-api/src/opentelemetry/trace/status.py

@@ -19,6 +19,9 @@
 logger = logging.getLogger(__name__)
 
 
+EXCEPTION_STATUS_FIELD = "_otel_status_code"
+
+
 class StatusCanonicalCode(enum.Enum):
     """Represents the canonical set of status codes of a finished Span."""
 

opentelemetry-sdk/CHANGELOG.md

@@ -10,6 +10,7 @@
   ([#1105](https://github.com/open-telemetry/opentelemetry-python/pull/1120))
 - Allow for Custom Trace and Span IDs Generation - `IdsGenerator` for TracerProvider
   ([#1153](https://github.com/open-telemetry/opentelemetry-python/pull/1153))
+- `start_as_current_span` and `use_span` can now optionally auto-record any exceptions raised inside the context manager.
 
 ## Version 0.13b0
 

opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

@@ -45,7 +45,11 @@
 from opentelemetry.sdk.util.instrumentation import InstrumentationInfo
 from opentelemetry.trace import SpanContext
 from opentelemetry.trace.propagation import SPAN_KEY
-from opentelemetry.trace.status import Status, StatusCanonicalCode
+from opentelemetry.trace.status import (
+    EXCEPTION_STATUS_FIELD,
+    Status,
+    StatusCanonicalCode,
+)
 from opentelemetry.util import time_ns, types
 
 logger = logging.getLogger(__name__)
@@ -687,9 +691,12 @@ def start_as_current_span(
         kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
         attributes: types.Attributes = None,
         links: Sequence[trace_api.Link] = (),
+        record_exception: bool = True,
     ) -> Iterator[trace_api.Span]:
-        span = self.start_span(name, parent, kind, attributes, links)
-        return self.use_span(span, end_on_exit=True)
+        span = self.start_span(name, parent, kind, attributes, links,)
+        return self.use_span(
+            span, end_on_exit=True, record_exception=record_exception
+        )
 
     def start_span(  # pylint: disable=too-many-locals
         self,
@@ -768,7 +775,10 @@ def start_span(  # pylint: disable=too-many-locals
 
     @contextmanager
     def use_span(
-        self, span: trace_api.Span, end_on_exit: bool = False
+        self,
+        span: trace_api.Span,
+        end_on_exit: bool = False,
+        record_exception: bool = True,
     ) -> Iterator[trace_api.Span]:
         try:
             token = context_api.attach(context_api.set_value(SPAN_KEY, span))
@@ -778,20 +788,24 @@ def use_span(
                 context_api.detach(token)
 
         except Exception as error:  # pylint: disable=broad-except
-            if (
-                isinstance(span, Span)
-                and span.status is None
-                and span._set_status_on_exception  # pylint:disable=protected-access  # noqa
-            ):
-                span.set_status(
-                    Status(
-                        canonical_code=StatusCanonicalCode.UNKNOWN,
-                        description="{}: {}".format(
-                            type(error).__name__, error
-                        ),
+            # pylint:disable=protected-access
+            if isinstance(span, Span):
+                if record_exception:
+                    span.record_exception(error)
+
+                if span.status is None and span._set_status_on_exception:
+                    span.set_status(
+                        Status(
+                            canonical_code=getattr(
+                                error,
+                                EXCEPTION_STATUS_FIELD,
+                                StatusCanonicalCode.UNKNOWN,
+                            ),
+                            description="{}: {}".format(
+                                type(error).__name__, error
+                            ),
+                        )
                     )
-                )
-
             raise
 
         finally:

opentelemetry-sdk/tests/trace/test_trace.py

@@ -800,6 +800,38 @@ def test_record_exception(self):
             exception_event.attributes["exception.stacktrace"],
         )
 
+    def test_record_exception_context_manager(self):
+        try:
+            with self.tracer.start_as_current_span("span") as span:
+                raise RuntimeError("example error")
+        except RuntimeError:
+            pass
+        finally:
+            self.assertEqual(len(span.events), 1)
+            event = span.events[0]
+            self.assertEqual("exception", event.name)
+            self.assertEqual(
+                "RuntimeError", event.attributes["exception.type"]
+            )
+            self.assertEqual(
+                "example error", event.attributes["exception.message"]
+            )
+
+            stacktrace = """in test_record_exception_context_manager
+    raise RuntimeError("example error")
+RuntimeError: example error"""
+            self.assertIn(stacktrace, event.attributes["exception.stacktrace"])
+
+        try:
+            with self.tracer.start_as_current_span(
+                "span", record_exception=False
+            ) as span:
+                raise RuntimeError("example error")
+        except RuntimeError:
+            pass
+        finally:
+            self.assertEqual(len(span.events), 0)
+
 
 def span_event_start_fmt(span_processor_name, span_name):
     return span_processor_name + ":" + span_name + ":start"