72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""trace_id 全链路贯通（contextvars · 架构 §5.1）。
|
||
|
||
约束（编码前必读）：
|
||
- 中间件必须在 call_next 之前 set()，禁止在 endpoint/service 内重新 set()；
|
||
- 后台任务（create_task）须显式 contextvars.copy_context()；
|
||
- 同步 def 路由跑线程池时 context 经 anyio 传播，由 B8 trace 一致性断言实测兜底。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
from contextvars import ContextVar, Token
|
||
from uuid import uuid4
|
||
|
||
# 透传外部 X-Trace-Id / X-Request-Id 的格式白名单（防响应头注入；不合规一律新生成）
|
||
_HEADER_ID_PATTERN = re.compile(r"^[A-Za-z0-9._-]{1,64}$")
|
||
|
||
_trace_id: ContextVar[str] = ContextVar("trace_id", default="")
|
||
# 独立请求级 ID（T-02 · B7 复审 P3-4）：trace_id 贯通链路，request_id 标识单次
|
||
# HTTP 请求（幂等/重试对账，手册 §4.6），两者不再互用。
|
||
_request_id: ContextVar[str] = ContextVar("request_id", default="")
|
||
|
||
|
||
def new_trace(trace_id: str | None = None) -> str:
|
||
"""生成并绑定新 trace_id；透传外部 X-Trace-Id 时传入该值。"""
|
||
tid = trace_id or f"trc-{uuid4().hex[:16]}"
|
||
_trace_id.set(tid)
|
||
return tid
|
||
|
||
|
||
def set_trace(trace_id: str) -> Token:
|
||
"""绑定外部 trace_id（中间件入口），返回 token 供请求结束复位。"""
|
||
return _trace_id.set(trace_id)
|
||
|
||
|
||
def reset_trace(token: Token) -> None:
|
||
"""复位到中间件之前的上下文（防异步任务间串号）。"""
|
||
_trace_id.reset(token)
|
||
|
||
|
||
def current_trace() -> str:
|
||
"""读取当前 trace_id；未初始化时返回空串（调用方应兜底生成）。"""
|
||
return _trace_id.get()
|
||
|
||
|
||
def ensure_trace() -> None:
|
||
"""无上下文时兜底归因（脚本/引擎入口），有值时保留（中间件场景不重新 set）。"""
|
||
if not current_trace():
|
||
new_trace()
|
||
|
||
|
||
def new_request_id() -> str:
|
||
"""生成并绑定 request_id（中间件入口每请求一次）。"""
|
||
rid = f"req-{uuid4().hex[:16]}"
|
||
_request_id.set(rid)
|
||
return rid
|
||
|
||
|
||
def bind_request_id(request_id: str | None) -> tuple[str, Token]:
|
||
"""绑定请求 request_id（合法透传否则生成）；返回 (id, token) 供中间件复位。"""
|
||
rid = request_id if _HEADER_ID_PATTERN.fullmatch(request_id or "") else f"req-{uuid4().hex[:16]}"
|
||
return rid, _request_id.set(rid)
|
||
|
||
|
||
def reset_request_id(token: Token) -> None:
|
||
_request_id.reset(token)
|
||
|
||
|
||
def current_request_id() -> str:
|
||
"""读取当前请求 request_id；未初始化返回空串（错误体兜底生成）。"""
|
||
return _request_id.get()
|