更强的Python日志装饰器

Nisus
254 阅读5分钟

您将了解到：

  1. 如何自定义自己的装饰器，含无参和有参的装饰器
  2. 如何近似的在代码中嵌入自己的链路追踪设计，日志打印 trace id

希望:

  1. 一个注解，自动帮我们打印方法的入参、出参、异常情况、耗时情况
  2. 进一步，有类似链路追踪的 trace id、span id 之类的，方便在庞杂的日志中快速串联起来排查问题

trace

先上个辅助工具，近似生成 trace id 的工具方法及装饰器。

先直接看效果。

调用链路:

a()
	b()
	b2()
c()

如果我在每个方法定义上都加上注解 @trace

@trace(True)
def a():
    print(f'aa {get()}')
    b()
    b2()

@trace
def b():
    print(f'bb {get()}')

@trace
def b2():
    print(f'bb2 {get()}')

@trace
def c():
    print(f'cc {get()}')

# ---
a()
c()

执行后

aa Trace(tid='1699775240926635', ptid='1699775240926635', rtid='1699775240926635')
bb Trace(tid='1699775240926930', ptid='1699775240926635', rtid='1699775240926635')
bb2 Trace(tid='1699775240926198', ptid='1699775240926635', rtid='1699775240926635')
cc Trace(tid='1699775240926132', ptid='1699775240926132', rtid='1699775240926132')

1699775240926635      --a()
		1699775240926930  --b()
		1699775240926198  --b2()
1699775240926132      --c()

  • trace 装饰器源码

    import dataclasses
import random
import threading
import time

@dataclasses.dataclass
class Trace:
    # 当前 trace id
    tid:str
    # 父 trace id
    ptid:str
    # root trace id
    rtid:str

# 毫秒时间戳 + 3位随机数
# 对于日志简单的trace标识基本够用了
gen_trace_id = lambda : f'{int(time.time()*1000)}{random.randint(1, 999):03d}'

# 'trace' -> Trace(...)
_trace_context = threading.local()

def start():
    """
    开启一段 span : start() , 内嵌在已有 trace 之下
    """
    p_trace = get()
    if p_trace.tid == '':  # root
        tid = gen_trace_id()
        # 链路的起始位置, 认为父id和根id都是自己
        root_trace = Trace(tid, tid, tid)
        _trace_context.trace = root_trace
        return root_trace
    
    # 父 Trace 下嵌套当前 Trace, 当前 Trace 的 ptid 就是父 Trace 的 tid
    trace = Trace(gen_trace_id(), p_trace.tid, p_trace.rtid)
    _trace_context.trace = trace
    return trace

def complete():
    trace = get()
    if trace.tid == trace.rtid:
        _trace_context.trace = Trace('','','')
        return
    # 经过任意的嵌套调用后，当前 span 的 trace 信息会有丢失，如果还需要用，请用一开始 start 返回的 trace 数据
    _trace_context.trace = Trace(trace.ptid, '', trace.rtid)

def get():
    return getattr(_trace_context, 'trace', Trace('','',''))

def trace(func=None, /):
    if func and not hasattr(func, '__call__'):
        func = None

    def out_wrapper(func):
        def wrapper(*args, **kwargs):
            t = start()
            try:
                res = func(*args, **kwargs)
                return res
            finally:
                complete()

        return wrapper

    if func:
        return out_wrapper(func)
    return out_wrapper

logx 装饰器

先看效果

import logging
import time
from logx import logx

logging.basicConfig(level=logging.DEBUG)

@logx
def test_logx_simple():
    logging.info('-------- test_logx_simple')

test_logx_simple()

# 位置参数我想记录 1,2,4 的参数; 键值对参数, 记录 code 和 xx 键名的
@logx(in_args=(1,2,4), in_kwargs=('code', 'xx'), out=False)
def test_logx_args(a, b, c, *argsss, **kws):
    logging.info('-------- test_logx_args')
    # print(1/0)
    time.sleep(3)
    nest()

@logx
def nest():
    logging.info('-------- nest')

# test_logx_args("aaa", d="ddd", b="bbb")
test_logx_args("aaa", "b", "c", "d", "e", "f", "g", "h", name="字符", code=42345)

产生两条 trace，其中第二条嵌套一个 span

span

  • 装饰器源码

    """
装饰器兼容传参和不传参
"""
import logging
import time
from typing import Tuple
from trace import start,complete

Vector = list[float]

IntTuple = Tuple[int, ...] | None
StrTuple = Tuple[str, ...] | None

def logx(func=None, /, *, in_args: IntTuple = (), in_kwargs: StrTuple = (), out: bool = True):
    """
    示例:
        @logx  # 所有入参将被记录日志
        def test1(a,b):
            pass
        @logx(in_args=(1,2,4), in_kwargs=('code', 'aa'), out=False)
        def test2(a,b):
            pass
    :param func: 被装饰函数. 自动打进来, 属于 Special Variables
    :param in_args: 指定位置形参那些需要记录日志, None 表示不记录, () 表示所有都记录, 默认
    :param in_kwargs: 字典参数那些需要记录, 参数名, None 表示不记录, () 表示所有都记录, 默认
    :param out: 是否记录出参日志, 默认记录
    :return:
    """

    # 无参拿不到 func, 延迟才能拿到, 返回 out_wrapper
    # 有参拿得到 func, 调一下 out_wrapper() 返回 wrapper
    # Tip: 两个 ()(), 无括号装饰, 就补上一个, 有括号装饰, 就不用补

    def pass_log_ins(args, kwargs):
        log_in_args = None
        if in_args is not None:
            if len(in_args) == 0:
                log_in_args = list(args)  # 所有入参
            else:  # 部分入参
                log_in_args = []
                for ix, arg in enumerate(args):
                    if ix in in_args:
                        log_in_args.append(arg)
        log_in_kwargs = None
        if in_kwargs is not None:
            if len(in_kwargs) == 0:  # 全部字典参数
                log_in_kwargs = kwargs
            else:  # 部分字典参数
                log_in_kwargs = {}
                for k, v in kwargs.items():
                    if k in in_kwargs:
                        log_in_kwargs[k] = kwargs[k]
        return log_in_args, log_in_kwargs

    def out_wrapper(func):
        def wrapper(*args, **kwargs):
            traceinfo = start()
            # trace info 日志
            traceinfostr = f'{traceinfo.rtid}..{traceinfo.ptid}.{traceinfo.tid}'
            st = time.time()
            log_in_args, log_in_kwargs = pass_log_ins(args, kwargs)
            logging.info(f'[{traceinfostr}][BEGIN]{func.__name__}: args: {log_in_args}, kwargs: {log_in_kwargs}')
            try:
                res = func(*args, **kwargs)
                res_log = f', return: {res}' if out else ''
                logging.info(f'[{traceinfostr}][END]{func.__name__}: cost: {time.time() - st:.3f}S{res_log}')
                return res
            except Exception as e:
                logging.error(f'[{traceinfostr}][END]{func.__name__}: cost: {time.time() - st:.3f}S, exception: {e}')
                raise
            finally:
                complete()

        return wrapper

    if func:
        return out_wrapper(func)
    return out_wrapper

后续

可能的问题

  • 线程数据传递时，可能引起的脏数据问题
  • 编码失误或其他原因，导致 start 和 complete 未能如期的成对执行，导致 trace id 错误
  • 由于 trace 树数据简单处理的，经过任意的嵌套调用后，当前 span 的 trace 信息会有丢失，如果还需要用，请用一开始 start 返回的 trace 数据
avatar
文章
阅读
粉丝