Trace、Span 与关联日志

mkdir l7-trace
cd l7-trace

"""用固定时间输入记录一条退款请求的 span 树与关联日志。"""

from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field

ACTIVE = ContextVar("active_span", default=None)


@dataclass
class Span:
    trace_id: str
    span_id: str
    parent_span_id: str | None
    name: str
    start_ms: int
    end_ms: int
    status: str = "UNSET"
    events: list = field(default_factory=list)


class Recorder:
    def __init__(self, trace_id):
        # OTel 的 trace ID 为 32 位十六进制，span ID 为 16 位。
        if len(trace_id) != 32 or any(c not in "0123456789abcdef" for c in trace_id):
            raise ValueError("invalid trace_id")
        if int(trace_id, 16) == 0:
            raise ValueError("zero trace_id")
        self.trace_id = trace_id
        self.spans = []
        self.logs = []

    @contextmanager
    def span(self, name, start_ms, end_ms):
        # 时间来自预置数据；这里没有测量真实服务。
        if start_ms < 0 or end_ms < start_ms:
            raise ValueError("invalid time interval")
        parent = ACTIVE.get()
        if parent is not None and parent.trace_id != self.trace_id:
            raise ValueError("different trace requires explicit link")
        child = Span(self.trace_id, f"{len(self.spans) + 1:016x}",
                     parent.span_id if parent else None, name, start_ms, end_ms)
        self.spans.append(child)
        token = ACTIVE.set(child)
        try:
            yield child
        except Exception as error:
            child.status = "ERROR"
            child.events.append({"name": "exception", "type": type(error).__name__})
            raise
        finally:
            # 子操作完成后恢复父上下文，避免下一条日志串错 span。
            ACTIVE.reset(token)

    def log(self, message):
        current = ACTIVE.get()
        if current is None or current.trace_id != self.trace_id:
            raise ValueError("log requires active span")
        self.logs.append({"trace_id": self.trace_id, "span_id": current.span_id,
                          "message": message})

    def tree_lines(self):
        lines = []
        def visit(parent_id, depth):
            for span in self.spans:
                if span.parent_span_id == parent_id:
                    lines.append(f"{'  ' * depth}{span.name} id={span.span_id} "
                                 f"status={span.status} duration={span.end_ms - span.start_ms}ms")
                    visit(span.span_id, depth + 1)
        visit(None, 0)
        return lines


def build_trace():
    record = Recorder("000000000000000000000000000000a1")
    with record.span("refund-request", 0, 150) as root:
        with record.span("model-plan", 10, 60) as model:
            record.log("plan ready")
            model.status = "OK"
        try:
            with record.span("policy-attempt-1", 60, 100):
                record.log("policy timeout")
                raise TimeoutError("fixed failure")
        except TimeoutError:
            # 重试是另一个操作，拥有新的 span ID，与首次尝试同属请求。
            with record.span("policy-attempt-2", 110, 140) as retry:
                record.log("policy found")
                retry.status = "OK"
        root.status = "OK"
    return record


def main():
    record = build_trace()
    print(f"trace_id={record.trace_id}")
    print("\n".join(record.tree_lines()))
    for row in record.logs:
        print(f"log span={row['span_id']} message={row['message']}")


if __name__ == "__main__":
    main()

ls
python -X utf8 trace_lab.py

trace_id=000000000000000000000000000000a1
refund-request id=0000000000000001 status=OK duration=150ms
  model-plan id=0000000000000002 status=OK duration=50ms
  policy-attempt-1 id=0000000000000003 status=ERROR duration=40ms
  policy-attempt-2 id=0000000000000004 status=OK duration=30ms
log span=0000000000000002 message=plan ready
log span=0000000000000003 message=policy timeout
log span=0000000000000004 message=policy found

ls
python -X utf8 -m unittest -v test_trace_lab.py