Try catch wrapping to make sure the span is finished

1 files changed, 34 insertions, 3 deletions

diff --git a/synapse/logging/opentracing.py b/synapse/logging/opentracing.py
index 846452e617..e185f8cf97 100644
--- a/synapse/logging/opentracing.py
+++ b/synapse/logging/opentracing.py
@@ -45,6 +45,7 @@ from twisted.internet import defer
 from twisted.internet import defer
 
 logger = logging.getLogger(__name__)
+import inspect
 
 
 class _DumTagNames(object):
@@ -405,9 +406,10 @@ def trace_function(func):
     @wraps(func)
     def f(self, *args, **kwargs):
         TracerUtil.start_active_span(func.__name__)
-        result = func(self, *args, **kwargs)
-        TracerUtil.close_active_span()
-        return result
+        try:
+            return func(self, *args, **kwargs)
+        finally:
+            TracerUtil.close_active_span()
 
     return f
 
@@ -420,3 +422,32 @@ def tag_args(func):
         return func(self, *args, **kwargs)
 
     return f
+
+
+def wrap_in_span(func):
+    """Its purpose is to wrap a function that is being passed into a context
+    which is a complete break from the current logcontext. This function creates
+    a non active span from the current context and closes it after the function
+    executes."""
+
+    # I haven't use this function yet
+
+    if not TracerUtil._opentracing:
+        return func
+
+    span = TracerUtil._opentracing.tracer.start_span(
+        func.__name__, child_of=TracerUtil._opentracing.tracer.active_span
+    )
+
+    @wraps(func)
+    def f(self, *args, **kwargs):
+        try:
+            return func(self, *args, **kwargs)
+        except Exception as e:
+            span.set_tag("error", True)
+            span.log_kv({"exception", e})
+            raise
+        finally:
+            span.finish()
+
+    return f