Merge branch 'main' into baggage-propagation-fix

Author: romank0

.gitignore

@@ -10,3 +10,4 @@ temporalio/bridge/temporal_sdk_bridge*
 /sdk-python.iml
 /.zed
 *.DS_Store
+tags

temporalio/contrib/opentelemetry.py

@@ -17,6 +17,16 @@
     cast,
 )
 
+import temporalio.activity
+import temporalio.api.common.v1
+import temporalio.client
+import temporalio.converter
+import temporalio.exceptions
+import temporalio.worker
+import temporalio.workflow
+from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
+from typing_extensions import Protocol, TypeAlias, TypedDict
+
 import opentelemetry.baggage.propagation
 import opentelemetry.context
 import opentelemetry.context.context
@@ -26,18 +36,7 @@
 import opentelemetry.trace.propagation.tracecontext
 import opentelemetry.util.types
 from opentelemetry.context import Context
-from opentelemetry.trace import Span, SpanKind, Status, StatusCode, _Links
-from opentelemetry.util import types
-from typing_extensions import Protocol, TypeAlias, TypedDict
-
-import temporalio.activity
-import temporalio.api.common.v1
-import temporalio.client
-import temporalio.converter
-import temporalio.exceptions
-import temporalio.worker
-import temporalio.workflow
-from temporalio.exceptions import ApplicationError, ApplicationErrorCategory
+from opentelemetry.trace import Status, StatusCode
 
 # OpenTelemetry dynamically, lazily chooses its context implementation at
 # runtime. When first accessed, they use pkg_resources.iter_entry_points + load.
@@ -173,29 +172,37 @@ def _start_as_current_span(
         kind: opentelemetry.trace.SpanKind,
         context: Optional[Context] = None,
     ) -> Iterator[None]:
-        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)
-            try:
-                yield None
-            except Exception as exc:
-                if (
-                    not isinstance(exc, ApplicationError)
-                    or exc.category != ApplicationErrorCategory.BENIGN
-                ):
-                    span.set_status(
-                        Status(
-                            status_code=StatusCode.ERROR,
-                            description=f"{type(exc).__name__}: {exc}",
+        if context:
+            token = opentelemetry.context.attach(context)
+        else:
+            token = None
+        try:
+            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)
+                try:
+                    yield None
+                except Exception as exc:
+                    if (
+                        not isinstance(exc, ApplicationError)
+                        or exc.category != ApplicationErrorCategory.BENIGN
+                    ):
+                        span.set_status(
+                            Status(
+                                status_code=StatusCode.ERROR,
+                                description=f"{type(exc).__name__}: {exc}",
+                            )
                         )
-                    )
-                raise
+                    raise
+        finally:
+            if token:
+                opentelemetry.context.detach(token)
 
     def _completed_workflow_span(
         self, params: _CompletedWorkflowSpanParams