33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
from __future__ import annotations
|
||
|
||
from contextvars import ContextVar, Token
|
||
|
||
# 当前消息链路的 trace_id：
|
||
# 1. 使用 ContextVar 而不是全局变量，避免并发消息之间互相覆盖；
|
||
# 2. asyncio Task 会自动继承上下文，因此插件里再起的协程也能拿到同一个 trace_id；
|
||
# 3. 对不在消息主链路里的后台任务，该值默认为空字符串，不会影响原有逻辑。
|
||
_CURRENT_TRACE_ID: ContextVar[str] = ContextVar("current_trace_id", default="")
|
||
|
||
|
||
def set_current_trace_id(trace_id: str) -> Token:
|
||
"""写入当前上下文的 trace_id，并返回可回滚 token。"""
|
||
return _CURRENT_TRACE_ID.set(str(trace_id or "").strip())
|
||
|
||
|
||
def reset_current_trace_id(token: Token) -> None:
|
||
"""根据 token 回滚 trace_id，避免上下文泄漏到后续无关任务。"""
|
||
_CURRENT_TRACE_ID.reset(token)
|
||
|
||
|
||
def get_current_trace_id() -> str:
|
||
"""读取当前上下文中的 trace_id。"""
|
||
return str(_CURRENT_TRACE_ID.get("") or "").strip()
|
||
|
||
|
||
def format_trace_prefix(trace_id: str = "") -> str:
|
||
"""统一生成日志前缀，避免各模块手写格式不一致。"""
|
||
resolved_trace_id = str(trace_id or get_current_trace_id() or "").strip()
|
||
if not resolved_trace_id:
|
||
return ""
|
||
return f"[trace_id={resolved_trace_id}] "
|