Fix tracing if method raise exception (#7133)

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

@@ -62,12 +62,14 @@ def wrapper_use_tracer(*args, **kwargs):
             child = parent_span.span(name=name)
             child.start()
             common.set_span_contexts(child)
-            ans = func(*args, **kwargs) # type: ignore
-            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)
+            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

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

@@ -62,12 +62,14 @@ async def wrapper_use_tracer(*args, **kwargs):
             child = parent_span.span(name=name)
             child.start()
             common.set_span_contexts(child)
-            ans = await func(*args, **kwargs)   # type: ignore
-            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)
+            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

sdk/core/azure-core/tests/test_tracing_decorator.py

@@ -67,6 +67,10 @@ def get_foo(self):
     def check_name_is_different(self):
         time.sleep(0.001)
 
+    @distributed_trace
+    def raising_exception(self):
+        raise ValueError("Something went horribly wrong here")
+
 
 def random_function():
     pass
@@ -188,3 +192,24 @@ def test_span_with_opencensus_complicated(self, value):
             assert parent.children[2].children[0].span_data.name == "MockClient.make_request"
             assert parent.children[3].span_data.name == "MockClient.make_request"
             assert not parent.children[3].children
+
+    def test_span_with_exception(self):
+        """Assert that if an exception is raised, the next sibling method is actually a sibling span.
+        """
+        with ContextHelper():
+            exporter = MockExporter()
+            trace = tracer_module.Tracer(sampler=AlwaysOnSampler(), exporter=exporter)
+            with trace.span("overall"):
+                client = MockClient()
+                try:
+                    client.raising_exception()
+                except:
+                    pass
+                client.get_foo()
+            trace.finish()
+            exporter.build_tree()
+            parent = exporter.root
+            assert len(parent.children) == 3
+            assert parent.children[0].span_data.name == "MockClient.__init__"
+            assert parent.children[1].span_data.name == "MockClient.raising_exception"
+            assert parent.children[2].span_data.name == "MockClient.get_foo"