From 2006f7ae94e80e758727177911dd646e37589b09 Mon Sep 17 00:00:00 2001
From: Jorik Schellekens <joriks@matrix.org>
Date: Fri, 28 Jun 2019 10:30:56 +0100
Subject: [PATCH] Add decerators for tracing functions

---
 synapse/logging/opentracing.py | 60 ++++++++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/synapse/logging/opentracing.py b/synapse/logging/opentracing.py
index f0ceea2a64..846452e617 100644
--- a/synapse/logging/opentracing.py
+++ b/synapse/logging/opentracing.py
@@ -39,6 +39,8 @@ import contextlib
 import logging
 import re
 from functools import wraps
+from twisted.internet import defer
+
 
 from twisted.internet import defer
 
@@ -360,3 +362,61 @@ def trace_servlet(servlet_name, func):
         defer.returnValue(result)
 
     return _trace_servlet_inner
+
+
+def trace_defered_function(func):
+    @wraps(func)
+    @defer.inlineCallbacks
+    def f(self, *args, **kwargs):
+        # Start scope
+        TracerUtil.start_active_span(func.__name__)
+        try:
+            r = yield func(self, *args, **kwargs)
+        except:
+            raise
+        finally:
+            TracerUtil.close_active_span()
+        defer.returnValue(r)
+
+    return f
+
+
+def trace_defered_function_using_operation_name(name):
+    def trace_defered_function(func):
+        @wraps(func)
+        @defer.inlineCallbacks
+        def f(self, *args, **kwargs):
+            # Start scope
+            TracerUtil.start_active_span(name)
+            try:
+                r = yield func(self, *args, **kwargs)
+            except:
+                raise
+            finally:
+                TracerUtil.close_active_span()
+            defer.returnValue(r)
+
+        return f
+
+    return trace_defered_function
+
+
+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
+
+    return f
+
+
+def tag_args(func):
+    @wraps(func)
+    def f(self, *args, **kwargs):
+        TracerUtil.set_tag("args", args)
+        TracerUtil.set_tag("kwargs", kwargs)
+        return func(self, *args, **kwargs)
+
+    return f