修复trace和asyncio共用时的问题

taskmgr · taskmgr · commit ce2972937f3e · 2020-11-30T19:46:40.000+08:00
diff --git a/skywalking/decorators.py b/skywalking/decorators.py
@@ -17,12 +17,14 @@
 
 from functools import wraps
 from typing import List
+import inspect
 
 from skywalking import Layer, Component
 from skywalking.trace.context import get_context
 from skywalking.trace.tags import Tag
 
 
+
 def trace(
         op: str = None,
         layer: Layer = Layer.Unknown,
@@ -31,7 +33,7 @@ def trace(
 ):
     def decorator(func):
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        async def wrapper(*args, **kwargs):
             _op = op or func.__name__
             context = get_context()
             with context.new_local_span(op=_op) as span:
@@ -40,6 +42,8 @@ def wrapper(*args, **kwargs):
                 [span.tag(tag) for tag in tags or []]
                 try:
                     result = func(*args, **kwargs)
+                    if inspect.isawaitable(result):
+                        return async await result
                     return result
                 except Exception:
                     span.raised()
