Tracing update

sdk/core/azure-core/azure/core/pipeline/policies/distributed_tracing.py

@@ -24,11 +24,11 @@
 #
 # --------------------------------------------------------------------------
 """Traces network calls using the implementation library from the settings."""
-
+import sys
 from six.moves import urllib
 
 from azure.core.tracing.context import tracing_context
-from azure.core.tracing.common import set_span_contexts
+from azure.core.tracing.common import set_span_contexts, get_parent_span
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 from azure.core.settings import settings
 
@@ -64,17 +64,16 @@ def set_header(self, request, span):  # pylint: disable=no-self-use
 
     def on_request(self, request):
         # type: (PipelineRequest) -> None
-        parent_span = tracing_context.current_span.get()
-        wrapper_class = settings.tracing_implementation()
-        original_context = [parent_span, None]
-        if parent_span is None and wrapper_class is not None:
-            current_span_instance = wrapper_class.get_current_span()
-            original_context[1] = current_span_instance
-            parent_span = wrapper_class(current_span_instance)
-
-        if parent_span is None:
+        span_impl_type = settings.tracing_implementation()
+        if span_impl_type is None:
             return
 
+        original_context = (
+            tracing_context.current_span.get(),
+            span_impl_type.get_current_span()
+        )
+        parent_span = get_parent_span()
+
         path = urllib.parse.urlparse(request.http_request.url).path  # type: ignore
         if not path:
             path = "/"
@@ -85,7 +84,7 @@ def on_request(self, request):
         self.parent_span_dict[child] = original_context
         self.set_header(request, child)
 
-    def end_span(self, request, response=None):
+    def end_span(self, request, response=None, exc_info=None):
         # type: (HttpRequest, Optional[HttpResponse]) -> None
         """Ends the span that is tracing the network and updates its status."""
         span = tracing_context.current_span.get()  # type: AbstractSpan
@@ -96,7 +95,11 @@ def end_span(self, request, response=None):
                 span.add_attribute(self._request_id, request_id)
             if response and self._response_id in response.headers:
                 span.add_attribute(self._response_id, response.headers[self._response_id])
-            span.finish()
+            if exc_info:
+                span.__exit__(*exc_info)
+            else:
+                span.finish()
+
             original_context = self.parent_span_dict.pop(span, None)
             if original_context:
                 set_span_contexts(original_context[0], original_context[1])
@@ -107,5 +110,5 @@ def on_response(self, request, response):
 
     def on_exception(self, _request):  # pylint: disable=unused-argument
         # type: (PipelineRequest) -> bool
-        self.end_span(_request.http_request)  # type: ignore
+        self.end_span(_request.http_request, exc_info=sys.exc_info())  # type: ignore
         return False

sdk/core/azure-core/azure/core/tracing/__init__.py

@@ -2,8 +2,8 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
-from azure.core.tracing.abstract_span import AbstractSpan
+from azure.core.tracing.abstract_span import AbstractSpan, SpanKind
 
 __all__ = [
-    "AbstractSpan",
+    "AbstractSpan", "SpanKind"
 ]

sdk/core/azure-core/azure/core/tracing/abstract_span.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # ------------------------------------
 """Protocol that defines what functions wrappers of tracing libraries should implement."""
+from enum import Enum
 
 try:
     from typing import TYPE_CHECKING
@@ -20,6 +21,15 @@
     Protocol = object  # type: ignore
 
 
+class SpanKind(Enum):
+    UNSPECIFIED = 1
+    SERVER = 2
+    CLIENT = 3
+    PRODUCER = 4
+    CONSUMER = 5
+    INTERNAL = 6
+
+
 class AbstractSpan(Protocol):
     """Wraps a span from a distributed tracing implementation."""
 
@@ -37,6 +47,22 @@ def span(self, name="child_span"):
         The child span must be wrapped by an implementation of AbstractSpan
         """
 
+    @property
+    def kind(self):
+        # type: () -> SpanKind
+        """Get the span kind of this span."""
+
+    @kind.setter
+    def kind(self, value):
+        # type: (SpanKind) -> None
+        """Set the span kind of this span."""
+
+    def __enter__(self):
+        """Start a span."""
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        """Finish a span."""
+
     def start(self):
         # type: () -> None
         """Set the start time for a span."""
@@ -73,6 +99,21 @@ def set_http_attributes(self, request, response=None):
         :type response: HttpResponse
         """
 
+    def get_trace_parent(self):
+        # type: () -> str
+        """Return traceparent string as defined in W3C trace context specification.
+
+        Example:
+        Value = 00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01
+        base16(version) = 00
+        base16(trace-id) = 4bf92f3577b34da6a3ce929d0e0e4736
+        base16(parent-id) = 00f067aa0ba902b7
+        base16(trace-flags) = 01  // sampled
+
+        :return: a traceparent string
+        :rtype: str
+        """
+
     @property
     def span_instance(self):
         # type: () -> Any
@@ -81,7 +122,17 @@ def span_instance(self):
         """
 
     @classmethod
-    def link(cls, headers):
+    def link(cls, traceparent):
+        # type: (Dict[str, str]) -> None
+        """
+        Given a traceparent, extracts the context and links the context to the current tracer.
+
+        :param traceparent: A string representing a traceparent
+        :type headers: str
+        """
+
+    @classmethod
+    def link_from_headers(cls, headers):
         # type: (Dict[str, str]) -> None
         """
         Given a dictionary, extracts the context and links the context to the current tracer.

sdk/core/azure-core/azure/core/tracing/common.py

@@ -62,43 +62,58 @@ def set_span_contexts(wrapped_span, span_instance=None):
     # type: (Union[AbstractSpan, None], Optional[AbstractSpan]) -> None
     """
     Set the sdk context and the implementation context. `span_instance` will be used to set the implementation context
-    if passed in else will use `wrapped_span.span_instance`.
+    ONLY if wrapped_span is None. Otherwise, will use internal implementation span.
 
     :param wrapped_span: The `AbstractSpan` to set as the sdk context
     :type wrapped_span: `azure.core.tracing.abstract_span.AbstractSpan`
-    :param span_instance: The span to set as the current span for the implementation context
+    :param span_instance: The span to set as the current span for the implementation context if wrapped_span is None.
     """
+    span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]
+    if span_impl_type is None:
+        return
+
+    # Store the current wrapped span into our SDK context
     tracing_context.current_span.set(wrapped_span)
-    impl_wrapper = settings.tracing_implementation()
-    if wrapped_span is not None:
-        span_instance = wrapped_span.span_instance
-    if impl_wrapper is not None:
-        impl_wrapper.set_current_span(span_instance)
+    if wrapped_span is None:
+        span_impl_type.set_current_span(span_instance)
+    else:
+        span_impl_type.set_current_span(wrapped_span.span_instance)
 
 
-def get_parent_span(parent_span):
+def get_parent_span(parent_span=None):
     # type: (Any) -> Optional[AbstractSpan]
     """
     Returns the current span so that the function's span will be its child. It will create a new span if there is
     no current span in any of the context.
 
+    The only possiblity to get None is if there is no tracing plugin available.
+
+    Algorithm is:
+    - Return a SDK span if parent_span exists
+    - ELSE return the SDK current span if exists
+    - ELSE return the a new SDK span based on implementation if exists
+    - ELSE creates a new implementation and SDK span and return it
+
+    If this method creates a span, it will NOT store it in SDK context.
+
     :param parent_span: The parent_span arg that the user passes into the top level function
     :returns: the parent_span of the function to be traced
     :rtype: `azure.core.tracing.abstract_span.AbstractSpan`
     """
-    wrapper_class = settings.tracing_implementation()
-    if wrapper_class is None:
+    span_impl_type = settings.tracing_implementation()  # type: Type[AbstractSpan]
+    if span_impl_type is None:
         return None
 
-    orig_wrapped_span = tracing_context.current_span.get()
-    # parent span is given, get from my context, get from the implementation context or make our own
-    parent_span = orig_wrapped_span if parent_span is None else wrapper_class(parent_span)
-    if parent_span is None:
-        current_span = wrapper_class.get_current_span()
-        parent_span = (
-            wrapper_class(span=current_span)
-            if current_span
-            else wrapper_class(name="azure-sdk-for-python-first_parent_span")
-        )
-
-    return parent_span
+    if parent_span is not None:
+        return span_impl_type(parent_span)  # This span is NOT stored in SDK context yet.
+
+    orig_wrapped_span = tracing_context.current_span.get()  # type: AbstractSpan
+    if orig_wrapped_span is not None:
+        return orig_wrapped_span
+
+    current_span = span_impl_type.get_current_span()
+    if current_span is not None:
+        return span_impl_type(current_span)
+
+    # Everything is None, SDK has no span and customer didn't create one. Create the base one for the customer
+    return span_impl_type(name="azure-sdk-for-python-first_parent_span")

sdk/core/azure-core/azure/core/tracing/decorator.py

@@ -48,30 +48,35 @@ def distributed_trace(func=None, name_of_span=None):
     @functools.wraps(func)
     def wrapper_use_tracer(*args, **kwargs):
         # type: (Any, Any) -> Any
+        merge_span = kwargs.pop('merge_span', False)
         passed_in_parent = kwargs.pop("parent_span", None)
-        orig_wrapped_span = tracing_context.current_span.get()
-        wrapper_class = settings.tracing_implementation()
-        original_span_instance = None
-        if wrapper_class is not None:
-            original_span_instance = wrapper_class.get_current_span()
+
+        span_impl_type = settings.tracing_implementation()
+        if span_impl_type is None:
+            return func(*args, **kwargs) # type: ignore
+
+        # Merge span is parameter is set, but only if no explicit parent are passed
+        if merge_span and not passed_in_parent:
+            return func(*args, **kwargs) # type: ignore
+
+        # Get original context
+        original_wrapped_span = tracing_context.current_span.get()  # type: AbstractSpan
+        original_span_instance = span_impl_type.get_current_span()
+
         parent_span = common.get_parent_span(passed_in_parent)
-        ans = None
-        if parent_span is not None and orig_wrapped_span is None:
+        common.set_span_contexts(parent_span)
+
+        name = name_of_span or common.get_function_and_class_name(func, *args)  # type: ignore
+        child = parent_span.span(name=name)
+        try:
+            with child:
+                common.set_span_contexts(child)
+                return func(*args, **kwargs) # type: ignore
+        finally:
             common.set_span_contexts(parent_span)
-            name = name_of_span or common.get_function_and_class_name(func, *args)  # type: ignore
-            child = parent_span.span(name=name)
-            child.start()
-            common.set_span_contexts(child)
-            try:
-                ans = func(*args, **kwargs) # type: ignore
-            finally:
-                child.finish()
-                common.set_span_contexts(parent_span)
-                if orig_wrapped_span is None and passed_in_parent is None and original_span_instance is None:
-                    parent_span.finish()
-                common.set_span_contexts(orig_wrapped_span, span_instance=original_span_instance)
-        else:
-            ans = func(*args, **kwargs) # type: ignore
-        return ans
+            # This test means "get_parent" created the span, so I need to finish it.
+            if original_wrapped_span is None and passed_in_parent is None and original_span_instance is None:
+                parent_span.finish()
+            common.set_span_contexts(original_wrapped_span, span_instance=original_span_instance)
 
     return wrapper_use_tracer

sdk/core/azure-core/azure/core/tracing/decorator_async.py

@@ -48,30 +48,35 @@ def distributed_trace_async(func=None, name_of_span=None):
     @functools.wraps(func)
     async def wrapper_use_tracer(*args, **kwargs):
         # type: (Any, Any) -> Any
+        merge_span = kwargs.pop('merge_span', False)
         passed_in_parent = kwargs.pop("parent_span", None)
-        orig_wrapped_span = tracing_context.current_span.get()
-        wrapper_class = settings.tracing_implementation()
-        original_span_instance = None
-        if wrapper_class is not None:
-            original_span_instance = wrapper_class.get_current_span()
+
+        span_impl_type = settings.tracing_implementation()
+        if span_impl_type is None:
+            return await func(*args, **kwargs) # type: ignore
+
+        # Merge span is parameter is set, but only if no explicit parent are passed
+        if merge_span and not passed_in_parent:
+            return await func(*args, **kwargs) # type: ignore
+
+        # Get original context
+        original_wrapped_span = tracing_context.current_span.get()  # type: AbstractSpan
+        original_span_instance = span_impl_type.get_current_span()
+
         parent_span = common.get_parent_span(passed_in_parent)
-        ans = None
-        if parent_span is not None and orig_wrapped_span is None:
+        common.set_span_contexts(parent_span)
+
+        name = name_of_span or common.get_function_and_class_name(func, *args)  # type: ignore
+        child = parent_span.span(name=name)
+        try:
+            with child:
+                common.set_span_contexts(child)
+                return await func(*args, **kwargs)  # type: ignore
+        finally:
             common.set_span_contexts(parent_span)
-            name = name_of_span or common.get_function_and_class_name(func, *args)  # type: ignore
-            child = parent_span.span(name=name)
-            child.start()
-            common.set_span_contexts(child)
-            try:
-                ans = await func(*args, **kwargs)   # type: ignore
-            finally:
-                child.finish()
-                common.set_span_contexts(parent_span)
-                if orig_wrapped_span is None and passed_in_parent is None and original_span_instance is None:
-                    parent_span.finish()
-                common.set_span_contexts(orig_wrapped_span, span_instance=original_span_instance)
-        else:
-            ans = await func(*args, **kwargs)   # type: ignore
-        return ans
+            # This test means "get_parent" created the span, so I need to finish it.
+            if original_wrapped_span is None and passed_in_parent is None and original_span_instance is None:
+                parent_span.finish()
+            common.set_span_contexts(original_wrapped_span, span_instance=original_span_instance)
 
     return wrapper_use_tracer