Allow the passing of operation_name to trace

synapse/logging/opentracing.py

@@ -85,14 +85,14 @@ def interesting_function(*args, **kwargs):
        return something_usual_and_useful
 
 
-Operation names can be explicitly set for functions by using
-``trace_using_operation_name``
+Operation names can be explicitly set for a function by passing the
+operation name to ``trace``
 
 .. code-block:: python
 
-   from synapse.logging.opentracing import trace_using_operation_name
+   from synapse.logging.opentracing import trace
 
-   @trace_using_operation_name("A *much* better operation name")
+   @trace(operation_name="A *much* better operation name")
    def interesting_badly_named_function(*args, **kwargs):
        # Does all kinds of cool and expected things
        return something_usual_and_useful
@@ -641,57 +641,16 @@ def extract_text_map(carrier):
 # Tracing decorators
 
 
-def trace(func):
+def trace(func=None, operation_name=None):
     """
     Decorator to trace a function.
     Sets the operation name to that of the function's.
     """
-    if opentracing is None:
-        return func
-
-    @wraps(func)
-    def _trace_inner(self, *args, **kwargs):
-        if opentracing is None:
-            return func(self, *args, **kwargs)
-
-        scope = start_active_span(func.__name__)
-        scope.__enter__()
-
-        try:
-            result = func(self, *args, **kwargs)
-            if isinstance(result, defer.Deferred):
-
-                def call_back(result):
-                    scope.__exit__(None, None, None)
-                    return result
-
-                def err_back(result):
-                    scope.span.set_tag(tags.ERROR, True)
-                    scope.__exit__(None, None, None)
-                    return result
 
-                result.addCallbacks(call_back, err_back)
+    if func and not operation_name:
+        operation_name = func.__name__
 
-            else:
-                scope.__exit__(None, None, None)
-
-            return result
-
-        except Exception as e:
-            scope.__exit__(type(e), None, e.__traceback__)
-            raise
-
-    return _trace_inner
-
-
-def trace_using_operation_name(operation_name):
-    """Decorator to trace a function. Explicitely sets the operation_name."""
-
-    def trace(func):
-        """
-        Decorator to trace a function.
-        Sets the operation name to that of the function's.
-        """
+    def decorator(func):
         if opentracing is None:
             return func
 
@@ -717,6 +676,7 @@ def err_back(result):
                         return result
 
                     result.addCallbacks(call_back, err_back)
+
                 else:
                     scope.__exit__(None, None, None)
 
@@ -728,7 +688,10 @@ def err_back(result):
 
         return _trace_inner
 
-    return trace
+    if func:
+        return decorator(func)
+    else:
+        return decorator
 
 
 def tag_args(func):

synapse/rest/client/v2_alpha/sendtodevice.py

@@ -19,7 +19,7 @@
 
 from synapse.http import servlet
 from synapse.http.servlet import parse_json_object_from_request
-from synapse.logging.opentracing import set_tag, trace_using_operation_name
+from synapse.logging.opentracing import set_tag, trace
 from synapse.rest.client.transactions import HttpTransactionCache
 
 from ._base import client_patterns
@@ -43,7 +43,7 @@ def __init__(self, hs):
         self.txns = HttpTransactionCache(hs)
         self.device_message_handler = hs.get_device_message_handler()
 
-    @trace_using_operation_name("sendToDevice")
+    @trace(operation_name="sendToDevice")
     def on_PUT(self, request, message_type, txn_id):
         set_tag("message_type", message_type)
         set_tag("txn_id", txn_id)