Neo-ZQYY/apps/backend/app/trace/middleware.py

# -*- coding: utf-8 -*-
"""
TraceMiddleware — ASGI 中间件

拦截 xcx_* 路由前缀（/api/xcx/）的请求，创建 TraceContext 并记录全链路 span。
非 xcx 路由直接跳过，不创建 TraceContext。
DEV_TRACE_ENABLED 关闭时跳过所有采集。

记录的 span 类型：
- HTTP_IN: 请求进入（method, path, query_params, body_preview）
- HTTP_OUT: 请求结束（status_code, duration, body_size）
- MIDDLEWARE: ResponseWrapperMiddleware 执行耗时
- MIDDLEWARE_ERROR: 响应包装失败时记录
"""

from __future__ import annotations

import logging
import time
from datetime import datetime
from urllib.parse import unquote

from starlette.types import ASGIApp, Message, Receive, Scope, Send

from app.trace.config import get_trace_config
from app.trace.context import (
    SpanType,
    TraceContext,
    TraceSpan,
    create_http_trace,
    get_current_trace,
    set_current_trace,
    trace_context_var,
)
from app.trace.writer import get_trace_writer

logger = logging.getLogger(__name__)

# xcx 路由前缀——仅匹配此前缀的请求才采集 trace
XCX_PATH_PREFIX = "/api/xcx/"


def _should_trace(path: str) -> bool:
    """判断请求路径是否属于 xcx_* 路由前缀，需要采集 trace。"""
    return path.startswith(XCX_PATH_PREFIX)


class TraceMiddleware:
    """ASGI 中间件：全链路请求追踪。

    执行顺序（最外层，最先执行）：
    TraceMiddleware → CORSMiddleware → ResponseWrapperMiddleware → 路由处理

    职责：
    1. 检查 DEV_TRACE_ENABLED 开关（运行时检查，支持动态切换）
    2. 仅拦截 /api/xcx/ 前缀的请求
    3. 创建 TraceContext 存入 contextvars
    4. 记录 HTTP_IN / HTTP_OUT / MIDDLEWARE / MIDDLEWARE_ERROR span
    5. 响应头写入 X-Request-ID, X-Process-Time, X-DB-Queries, X-DB-Time
    6. 调用 TraceWriter 写入完整 trace
    """

    def __init__(self, app: ASGIApp) -> None:
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        # 仅处理 HTTP 请求
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path = scope.get("path", "")
        config = get_trace_config()

        # 开关关闭 或 非 xcx 路由 → 直接透传
        if not config.enabled or not _should_trace(path):
            await self.app(scope, receive, send)
            return

        # ── 创建 TraceContext ──
        method = scope.get("method", "UNKNOWN")
        query_string = scope.get("query_string", b"").decode("latin-1", errors="replace")
        ctx = create_http_trace(method, path)
        token = set_current_trace(ctx)

        # 记录 HTTP_IN span
        query_params = {}
        if query_string:
            # 简单解析 query string 为 dict
            for pair in query_string.split("&"):
                if "=" in pair:
                    k, v = pair.split("=", 1)
                    query_params[unquote(k)] = unquote(v)

        ctx.add_span(TraceSpan(
            span_type=SpanType.HTTP_IN,
            module="trace.middleware",
            function="TraceMiddleware.__call__",
            description_zh=f"接收请求 {method} {path}",
            description_en=f"Received request {method} {path}",
            params={"query": query_params},
            result_summary="",
            duration_ms=0.0,
            timestamp=datetime.now().isoformat(),
        ))

        start_time = time.perf_counter()

        # ── 拦截响应，注入 trace 头和采集 HTTP_OUT ──
        status_code = 200
        response_body_parts: list[bytes] = []
        middleware_start = time.perf_counter()
        middleware_error: str | None = None

        async def send_wrapper(message: Message) -> None:
            nonlocal status_code, middleware_error

            if message["type"] == "http.response.start":
                status_code = message.get("status", 200)

                # 计算 trace 统计数据
                elapsed_ms = (time.perf_counter() - start_time) * 1000
                current_ctx = get_current_trace()
                db_queries = 0
                db_time_ms = 0.0
                if current_ctx:
                    for s in current_ctx.spans:
                        if s.span_type == SpanType.DB_QUERY:
                            db_queries += 1
                            db_time_ms += s.duration_ms

                # 注入响应头
                headers = list(message.get("headers", []))
                headers.append([b"x-request-id", ctx.request_id.encode()])
                headers.append([b"x-process-time", f"{elapsed_ms:.1f}ms".encode()])
                headers.append([b"x-db-queries", str(db_queries).encode()])
                headers.append([b"x-db-time", f"{db_time_ms:.1f}ms".encode()])
                message = {**message, "headers": headers}

                await send(message)
                return

            if message["type"] == "http.response.body":
                body = message.get("body", b"")
                response_body_parts.append(body)
                await send(message)
                return

            await send(message)

        try:
            await self.app(scope, receive, send_wrapper)
        except Exception as exc:
            # 即使内层异常，也要记录 trace
            middleware_error = f"{type(exc).__name__}: {exc}"
            raise
        finally:
            elapsed_ms = (time.perf_counter() - start_time) * 1000
            middleware_elapsed_ms = (time.perf_counter() - middleware_start) * 1000
            body_size = sum(len(p) for p in response_body_parts)

            # 记录 MIDDLEWARE span（ResponseWrapperMiddleware 执行耗时）
            if middleware_error:
                ctx.add_span(TraceSpan(
                    span_type=SpanType.MIDDLEWARE_ERROR,
                    module="middleware.response_wrapper",
                    function="ResponseWrapperMiddleware.__call__",
                    description_zh=f"响应包装失败: {middleware_error}",
                    description_en=f"Response wrapping failed: {middleware_error}",
                    params={},
                    result_summary=middleware_error,
                    duration_ms=middleware_elapsed_ms,
                    timestamp=datetime.now().isoformat(),
                ))
            else:
                ctx.add_span(TraceSpan(
                    span_type=SpanType.MIDDLEWARE,
                    module="middleware.response_wrapper",
                    function="ResponseWrapperMiddleware.__call__",
                    description_zh="响应包装中间件执行完成",
                    description_en="Response wrapper middleware completed",
                    params={},
                    result_summary=f"body_size={body_size}",
                    duration_ms=middleware_elapsed_ms,
                    timestamp=datetime.now().isoformat(),
                    extra={"body_size": body_size},
                ))

            # 记录 HTTP_OUT span
            ctx.add_span(TraceSpan(
                span_type=SpanType.HTTP_OUT,
                module="trace.middleware",
                function="TraceMiddleware.__call__",
                description_zh=f"响应返回 {status_code}，耗时 {elapsed_ms:.0f}ms",
                description_en=f"Response sent {status_code}, took {elapsed_ms:.0f}ms",
                params={},
                result_summary=f"{status_code}, {body_size}B body",
                duration_ms=elapsed_ms,
                timestamp=datetime.now().isoformat(),
                extra={"status_code": status_code, "body_size": body_size},
            ))

            # 写入 trace 日志
            try:
                writer = get_trace_writer()
                await writer.write_trace(ctx)
            except Exception:
                logger.warning("Trace 日志写入失败", exc_info=True)

            # 恢复 contextvars
            trace_context_var.reset(token)