175 lines
5.7 KiB
Python
175 lines
5.7 KiB
Python
"""
|
||
Qihu-AIFlow 统一可观测基座 (Phase 2.1)
|
||
=====================================
|
||
提供贯穿一次请求的全局链路 ID（trace_id）与结构化日志能力，
|
||
为后续重构（上下文引擎、媒体/LLM 成本归因）建立回归基线。
|
||
|
||
设计原则（Strangler Fig 增量）：
|
||
- 纯增量埋点，不改任何既有接口的返回结构。
|
||
- 向后兼容现有 `logging.getLogger(__name__)` 用法：trace_id 由 root logger 的
|
||
TraceIdFilter 自动注入，调用方无需改动。
|
||
- 不引入第三方依赖（structlog 未安装），自带可选的 JSON formatter。
|
||
|
||
使用方式：
|
||
from app.core.observability import set_trace_id, get_trace_id, span
|
||
|
||
# 在中间件里为每个请求注入：
|
||
set_trace_id(request_id)
|
||
|
||
# 在关键调用处打计时 span（自动带 trace_id）：
|
||
async with span("chat.send_message", conv_id=conv_id):
|
||
return await _send_message_core(...)
|
||
"""
|
||
import logging
|
||
import time
|
||
import uuid
|
||
import json
|
||
from contextvars import ContextVar
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# trace_id：贯穿一次请求（含其内所有子调用/协程）的全局链路 ID
|
||
# ---------------------------------------------------------------------------
|
||
_trace_id_var: ContextVar[str] = ContextVar("aiflow_trace_id", default="-")
|
||
|
||
|
||
def new_trace_id() -> str:
|
||
"""生成一个崭新的链路 ID（32 位 hex）。"""
|
||
return uuid.uuid4().hex
|
||
|
||
|
||
def set_trace_id(tid: str | None = None) -> str:
|
||
"""设置当前上下文的 trace_id；不传则自动生成。返回最终使用的 ID。"""
|
||
tid = tid or new_trace_id()
|
||
_trace_id_var.set(tid)
|
||
return tid
|
||
|
||
|
||
def get_trace_id() -> str:
|
||
"""读取当前上下文的 trace_id；未设置时返回占位 '-'。"""
|
||
return _trace_id_var.get()
|
||
|
||
|
||
class TraceIdFilter(logging.Filter):
|
||
"""为每一条日志记录注入 trace_id 属性（不存在则写为 '-'）。
|
||
|
||
挂载到 root logger 后，所有 `getLogger(__name__)` 的子 logger 都会继承，
|
||
从而让全平台既有日志自动带上链路 ID，零侵入。
|
||
"""
|
||
|
||
def filter(self, record: logging.LogRecord) -> bool:
|
||
if not hasattr(record, "trace_id"):
|
||
try:
|
||
record.trace_id = get_trace_id()
|
||
except Exception:
|
||
record.trace_id = "-"
|
||
return True
|
||
|
||
|
||
# 标准 logging 字段，JSON 序列化时需排除，避免与自定义 extra 冲突
|
||
_STD_ATTRS = set(logging.makeLogRecord({}).__dict__.keys()) | {
|
||
"message", "asctime", "msg", "trace_id", "exc_info", "stack_info", "args",
|
||
}
|
||
|
||
|
||
class JsonFormatter(logging.Formatter):
|
||
"""结构化 JSON 日志格式（通过环境变量 AIFLOW_JSON_LOGS=1 启用）。
|
||
|
||
默认不启用，避免破坏现有日志采集/人工查阅习惯。
|
||
"""
|
||
|
||
def format(self, record: logging.LogRecord) -> str:
|
||
data = {
|
||
"ts": self.formatTime(record, "%Y-%m-%dT%H:%M:%S%z"),
|
||
"level": record.levelname,
|
||
"logger": record.name,
|
||
"trace_id": getattr(record, "trace_id", "-"),
|
||
"msg": record.getMessage(),
|
||
}
|
||
for k, v in record.__dict__.items():
|
||
if k not in _STD_ATTRS:
|
||
try:
|
||
json.dumps(v, ensure_ascii=False)
|
||
data[k] = v
|
||
except (TypeError, ValueError):
|
||
data[k] = str(v)
|
||
if record.exc_info:
|
||
data["exc"] = self.formatException(record.exc_info)
|
||
return json.dumps(data, ensure_ascii=False)
|
||
|
||
|
||
def get_logger(name: str | None = None) -> logging.Logger:
|
||
"""获取 logger（向后兼容现有用法）。trace_id 由 root filter 自动注入。"""
|
||
return logging.getLogger(name)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# span：结构化计时上下文管理器（耗时 + trace_id + 自定义字段）
|
||
# ---------------------------------------------------------------------------
|
||
class _Span:
|
||
"""同步/异步通用的计时 span。退出时记录 elapsed_ms 与状态。"""
|
||
|
||
def __init__(self, name: str, **fields):
|
||
self.name = name
|
||
self.fields = fields
|
||
self._logger = logging.getLogger("aiflow.span")
|
||
self._start = None
|
||
self.elapsed_ms = 0.0
|
||
|
||
def _enter(self):
|
||
self._start = time.perf_counter()
|
||
self._logger.info(
|
||
f"SPAN_START {self.name}",
|
||
extra={"span": self.name, "event": "start", "trace_id": get_trace_id(), **self.fields},
|
||
)
|
||
return self
|
||
|
||
def _exit(self, exc_type, exc, tb):
|
||
self.elapsed_ms = (time.perf_counter() - self._start) * 1000
|
||
status = "error" if exc_type else "ok"
|
||
self._logger.info(
|
||
f"SPAN_END {self.name} {self.elapsed_ms:.1f}ms",
|
||
extra={
|
||
"span": self.name,
|
||
"event": "end",
|
||
"elapsed_ms": round(self.elapsed_ms, 1),
|
||
"status": status,
|
||
"trace_id": get_trace_id(),
|
||
**self.fields,
|
||
},
|
||
)
|
||
return False # 不吞掉异常
|
||
|
||
# 同步协议
|
||
def __enter__(self):
|
||
return self._enter()
|
||
|
||
def __exit__(self, exc_type, exc, tb):
|
||
return self._exit(exc_type, exc, tb)
|
||
|
||
# 异步协议
|
||
async def __aenter__(self):
|
||
return self._enter()
|
||
|
||
async def __aexit__(self, exc_type, exc, tb):
|
||
return self._exit(exc_type, exc, tb)
|
||
|
||
|
||
def span(name: str, **fields):
|
||
"""创建一个计时 span。可同步或异步 with。
|
||
|
||
async with span("chat.send_message", conv_id=cid):
|
||
...
|
||
"""
|
||
return _Span(name, **fields)
|
||
|
||
|
||
__all__ = [
|
||
"new_trace_id",
|
||
"set_trace_id",
|
||
"get_trace_id",
|
||
"TraceIdFilter",
|
||
"JsonFormatter",
|
||
"get_logger",
|
||
"span",
|
||
]
|