diff --git a/opentelemetry-api/src/opentelemetry/trace/__init__.py b/opentelemetry-api/src/opentelemetry/trace/__init__.py index b79cdeb4df5..33c62e9d3b6 100644 --- a/opentelemetry-api/src/opentelemetry/trace/__init__.py +++ b/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 diff --git a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py b/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py index c839fe4b616..544906e6fb4 100644 --- a/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py +++ b/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`.