Make use_span more flexible (closes #147). (#154)

Co-Authored-By: Reiley Yang <reyang@microsoft.com>
Co-Authored-By: Chris Kleinknecht <libc@google.com>

Author: 3 people

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

@@ -415,7 +415,8 @@ def start_span(
         is equivalent to::
 
             span = tracer.create_span(name)
-            with tracer.use_span(span):
+            span.start()
+            with tracer.use_span(span, end_on_exit=True):
                 do_work()
 
         Args:
@@ -472,17 +473,22 @@ def create_span(
         return INVALID_SPAN
 
     @contextmanager  # type: ignore
-    def use_span(self, span: "Span") -> typing.Iterator[None]:
+    def use_span(
+        self, span: "Span", end_on_exit: bool = False
+    ) -> typing.Iterator[None]:
         """Context manager for controlling a span's lifetime.
 
-        Start the given span and set it as the current span in this tracer's
-        context.
+        Set the given span as the current span in this tracer's context.
 
-        On exiting the context manager stop the span and set its parent as the
-        current span.
+        On exiting the context manager set the span that was previously active
+        as the current span (this is usually but not necessarily the parent of
+        the given span). If ``end_on_exit`` is ``True``, then the span is also
+        ended when exiting the context manager.
 
         Args:
             span: The span to start and make current.
+            end_on_exit: Whether to end the span automatically when leaving the
+                context manager.
         """
         # pylint: disable=unused-argument,no-self-use
         yield

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

@@ -436,7 +436,10 @@ def start_span(
         kind: trace_api.SpanKind = trace_api.SpanKind.INTERNAL,
     ) -> typing.Iterator["Span"]:
         """See `opentelemetry.trace.Tracer.start_span`."""
-        with self.use_span(self.create_span(name, parent, kind)) as span:
+
+        span = self.create_span(name, parent, kind)
+        span.start()
+        with self.use_span(span, end_on_exit=True):
             yield span
 
     def create_span(
@@ -473,16 +476,20 @@ def create_span(
         )
 
     @contextmanager
-    def use_span(self, span: "Span") -> typing.Iterator["Span"]:
+    def use_span(
+        self, span: Span, end_on_exit: bool = False
+    ) -> typing.Iterator[Span]:
         """See `opentelemetry.trace.Tracer.use_span`."""
-        span.start()
-        span_snapshot = self._current_span_slot.get()
-        self._current_span_slot.set(span)
         try:
-            yield span
+            span_snapshot = self._current_span_slot.get()
+            self._current_span_slot.set(span)
+            try:
+                yield span
+            finally:
+                self._current_span_slot.set(span_snapshot)
         finally:
-            self._current_span_slot.set(span_snapshot)
-            span.end()
+            if end_on_exit:
+                span.end()
 
     def add_span_processor(self, span_processor: SpanProcessor) -> None:
         """Registers a new :class:`SpanProcessor` for this `Tracer`.