Skip to content
This repository was archived by the owner on Apr 26, 2024. It is now read-only.

Commit b21fc7c

Browse files
committed
Allow use of both @trace and @tag_args stacked on the same function
```py @trace @tag_args async def get_oldest_event_ids_with_depth_in_room(...) ... ```
1 parent a648a06 commit b21fc7c

File tree

1 file changed

+91
-55
lines changed

1 file changed

+91
-55
lines changed

synapse/logging/opentracing.py

Lines changed: 91 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,7 @@ def set_fates(clotho, lachesis, atropos, father="Zues", mother="Themis"):
173173
Any,
174174
Callable,
175175
Collection,
176+
ContextManager,
176177
Dict,
177178
Generator,
178179
Iterable,
@@ -823,75 +824,110 @@ def extract_text_map(carrier: Dict[str, str]) -> Optional["opentracing.SpanConte
823824
# Tracing decorators
824825

825826

826-
def trace_with_opname(opname: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
827+
def _create_decorator(
828+
func: Callable[P, R],
829+
# TODO: What is the correct type for these `Any`? `P.args, P.kwargs` isn't allowed here
830+
wrapping_logic: Callable[[Callable[P, R], Any, Any], ContextManager[None]],
831+
) -> Callable[P, R]:
827832
"""
828-
Decorator to trace a function with a custom opname.
829-
830-
See the module's doc string for usage examples.
831-
833+
Creates a decorator that is able to handle sync functions, async functions
834+
(coroutines), and inlineDeferred from Twisted.
835+
Example usage:
836+
```py
837+
# Decorator to time the function and log it out
838+
def duration(func: Callable[P, R]) -> Callable[P, R]:
839+
@contextlib.contextmanager
840+
def _wrapping_logic(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs):
841+
start_ts = time.time()
842+
yield
843+
end_ts = time.time()
844+
duration = end_ts - start_ts
845+
logger.info("%s took %s seconds", func.__name__, duration)
846+
return _create_decorator(func, _wrapping_logic)
847+
```
848+
Args:
849+
func: The function to be decorated
850+
wrapping_logic: The business logic of your custom decorator.
851+
This should be a ContextManager so you are able to run your logic
852+
before/after the function as desired.
832853
"""
833854

834-
def decorator(func: Callable[P, R]) -> Callable[P, R]:
835-
if opentracing is None:
836-
return func # type: ignore[unreachable]
837-
855+
@wraps(func)
856+
async def _wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
838857
if inspect.iscoroutinefunction(func):
839-
840-
@wraps(func)
841-
async def _trace_inner(*args: P.args, **kwargs: P.kwargs) -> R:
842-
with start_active_span(opname):
843-
return await func(*args, **kwargs) # type: ignore[misc]
844-
858+
with wrapping_logic(func, *args, **kwargs):
859+
return await func(*args, **kwargs)
845860
else:
846861
# The other case here handles both sync functions and those
847862
# decorated with inlineDeferred.
848-
@wraps(func)
849-
def _trace_inner(*args: P.args, **kwargs: P.kwargs) -> R:
850-
scope = start_active_span(opname)
851-
scope.__enter__()
852-
853-
try:
854-
result = func(*args, **kwargs)
855-
if isinstance(result, defer.Deferred):
856-
857-
def call_back(result: R) -> R:
858-
scope.__exit__(None, None, None)
859-
return result
860-
861-
def err_back(result: R) -> R:
862-
scope.__exit__(None, None, None)
863-
return result
864-
865-
result.addCallbacks(call_back, err_back)
866-
867-
else:
868-
if inspect.isawaitable(result):
869-
logger.error(
870-
"@trace may not have wrapped %s correctly! "
871-
"The function is not async but returned a %s.",
872-
func.__qualname__,
873-
type(result).__name__,
874-
)
863+
scope = wrapping_logic(func, *args, **kwargs)
864+
scope.__enter__()
875865

866+
try:
867+
result = func(*args, **kwargs)
868+
if isinstance(result, defer.Deferred):
869+
870+
def call_back(result: R) -> R:
876871
scope.__exit__(None, None, None)
872+
return result
877873

878-
return result
874+
def err_back(result: R) -> R:
875+
scope.__exit__(None, None, None)
876+
return result
879877

880-
except Exception as e:
881-
scope.__exit__(type(e), None, e.__traceback__)
882-
raise
878+
result.addCallbacks(call_back, err_back)
883879

884-
return _trace_inner # type: ignore[return-value]
880+
else:
881+
if inspect.isawaitable(result):
882+
logger.error(
883+
"@trace may not have wrapped %s correctly! "
884+
"The function is not async but returned a %s.",
885+
func.__qualname__,
886+
type(result).__name__,
887+
)
885888

886-
return decorator
889+
scope.__exit__(None, None, None)
890+
891+
return result
892+
893+
except Exception as e:
894+
scope.__exit__(type(e), None, e.__traceback__)
895+
raise
896+
897+
return _wrapper # type: ignore[return-value]
898+
899+
900+
def trace_with_opname(opname: str) -> Callable[[Callable[P, R]], Callable[P, R]]:
901+
"""
902+
Decorator to trace a function with a custom opname.
903+
See the module's doc string for usage examples.
904+
"""
905+
906+
@contextlib.contextmanager
907+
def _wrapping_logic(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs):
908+
if opentracing is None:
909+
return None
910+
911+
scope = start_active_span(opname)
912+
scope.__enter__()
913+
try:
914+
yield
915+
except Exception as e:
916+
scope.__exit__(type(e), None, e.__traceback__)
917+
raise
918+
finally:
919+
scope.__exit__(None, None, None)
920+
921+
def _decorator(func: Callable[P, R]):
922+
return _create_decorator(func, _wrapping_logic)
923+
924+
return _decorator
887925

888926

889927
def trace(func: Callable[P, R]) -> Callable[P, R]:
890928
"""
891929
Decorator to trace a function.
892-
893930
Sets the operation name to that of the function's name.
894-
895931
See the module's doc string for usage examples.
896932
"""
897933

@@ -900,22 +936,22 @@ def trace(func: Callable[P, R]) -> Callable[P, R]:
900936

901937
def tag_args(func: Callable[P, R]) -> Callable[P, R]:
902938
"""
903-
Tags all of the args to the active span.
939+
Decorator to tag all of the args to the active span.
904940
"""
905941

906942
if not opentracing:
907943
return func
908944

909-
@wraps(func)
910-
def _tag_args_inner(*args: P.args, **kwargs: P.kwargs) -> R:
945+
@contextlib.contextmanager
946+
def _wrapping_logic(func: Callable[P, R], *args: P.args, **kwargs: P.kwargs):
911947
argspec = inspect.getfullargspec(func)
912948
for i, arg in enumerate(argspec.args[1:]):
913949
set_tag("ARG_" + arg, str(args[i])) # type: ignore[index]
914950
set_tag("args", str(args[len(argspec.args) :])) # type: ignore[index]
915951
set_tag("kwargs", str(kwargs))
916-
return func(*args, **kwargs)
952+
yield
917953

918-
return _tag_args_inner
954+
return _create_decorator(func, _wrapping_logic)
919955

920956

921957
@contextlib.contextmanager

0 commit comments

Comments
 (0)