From ee0f3b748bf951131a2281f901224ad9686087d0 Mon Sep 17 00:00:00 2001
From: Laurent Mazuel <laurent.mazuel@gmail.com>
Date: Tue, 10 Sep 2019 09:10:08 -0700
Subject: [PATCH] Fix tracing if method raise exception (#7133)

---
 .../azure/core/tracing/decorator.py           | 14 ++++++-----
 .../azure/core/tracing/decorator_async.py     | 14 ++++++-----
 .../tests/test_tracing_decorator.py           | 25 +++++++++++++++++++
 3 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/sdk/core/azure-core/azure/core/tracing/decorator.py b/sdk/core/azure-core/azure/core/tracing/decorator.py
index bd60bf86b875..cdc97d61fd2b 100644
--- a/sdk/core/azure-core/azure/core/tracing/decorator.py
+++ b/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
diff --git a/sdk/core/azure-core/azure/core/tracing/decorator_async.py b/sdk/core/azure-core/azure/core/tracing/decorator_async.py
index 875fbe34f06d..97eed2abfb83 100644
--- a/sdk/core/azure-core/azure/core/tracing/decorator_async.py
+++ b/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
diff --git a/sdk/core/azure-core/tests/test_tracing_decorator.py b/sdk/core/azure-core/tests/test_tracing_decorator.py
index 5c5a1df7b9e6..7f06f6e61610 100644
--- a/sdk/core/azure-core/tests/test_tracing_decorator.py
+++ b/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"