tdm_server_rust/profile/

mod.rs

//! 开发模式请求级调用栈耗时采集 (Request Profile)
//!
//! 在 dev 环境下自动采集每个 `#[tracing::instrument]` 的 self/inclusive 耗时，
//! 支持终端树形输出、热点高亮、JetBrains/Cursor 源码跳转链接。
//!
//! ## 功能特性
//!
//! - **Self/Inclusive 耗时**: 区分自身耗时和包含子调用的总耗时
//! - **热点高亮**: 自动标记 self 耗时最大的 span
//! - **源码跳转**: JetBrains 过滤器格式 / Cursor/VS Code OSC 8 链接
//! - **Server-Timing**: 生成 `Server-Timing` 响应头供前端联调
//! - **JSONL 导出**: `DEV_PROFILE_JSONL` 环境变量控制 JSONL 文件输出
//!
//! ## 环境变量
//!
//! | 变量 | 默认值 | 说明 |
//! |------|--------|------|
//! | `DEV_PROFILE_MIN_MS` | 0.1 | 最小展示阈值 (ms) |
//! | `DEV_PROFILE_LINK_SCHEME` | 自动检测 | 源码跳转协议 |
//! | `DEV_PROFILE_NO_LINKS` | - | 禁用源码链接 |
//! | `DEV_PROFILE_JSONL` | - | JSONL 输出文件路径 |
//! | `DEV_SERVER_TIMING` | 1 | 是否输出 Server-Timing 头 |

mod layer;

/// dev 模式 tracing Layer，自动采集 span 耗时并生成调用栈树
pub use layer::DevProfileLayer;

use crate::middleware::HttpLogStyle;
use axum::http::{Method, StatusCode, Uri};
use serde::Serialize;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tracing::Span;

/// 耗时列宽
const TIME_WIDTH: usize = 8;

/// Server-Timing 除 total/hot 外最多附加 span 数
const SERVER_TIMING_EXTRA: usize = 4;

/// 单次 span 记录
#[derive(Debug, Clone)]
struct SpanRecord {
    /// 展示名（module::name）
    label: String,
    /// 开始时刻
    started: Instant,
    /// 结束时刻
    ended: Option<Instant>,
    /// 父 span id
    parent: Option<tracing::Id>,
    /// 源码文件（instrument 位置）
    source_file: Option<String>,
    /// 源码行号
    source_line: Option<u32>,
}

/// 热点 span 摘要（JSONL 用）
#[derive(Debug, Clone, Serialize)]
pub struct ProfileHotspot {
    /// 展示名
    pub label: String,
    /// self 毫秒
    pub self_ms: f64,
    /// self 占墙钟百分比
    pub self_pct: f64,
    /// 源码文件
    pub file: Option<String>,
    /// 源码行号
    pub line: Option<u32>,
}

/// 单个 span 节点（JSONL 用）
#[derive(Debug, Clone, Serialize)]
pub struct ProfileSpanNode {
    /// 展示名
    pub label: String,
    /// self 毫秒
    pub self_ms: f64,
    /// inclusive 毫秒
    pub inclusive_ms: f64,
    /// self 占墙钟百分比
    pub self_pct: f64,
    /// 源码文件
    pub file: Option<String>,
    /// 源码行号
    pub line: Option<u32>,
}

/// 单次 HTTP 请求 profile 报告（JSONL 一行）
#[derive(Debug, Clone, Serialize)]
pub struct ProfileReport {
    /// HTTP 方法
    pub method: String,
    /// 完整 URI（含 query）
    pub uri: String,
    /// 墙钟毫秒
    pub wall_ms: f64,
    /// HTTP 状态码
    pub status: u16,
    /// self 最大的 span
    pub hotspot: Option<ProfileHotspot>,
    /// 可见 span 列表
    pub spans: Vec<ProfileSpanNode>,
}

/// 待打印节点（扁平收集后统一排版）
struct PrintNode {
    /// 树深度
    depth: usize,
    /// 展示名
    label: String,
    /// self 毫秒
    self_ms: f64,
    /// inclusive 毫秒
    inclusive_ms: f64,
    /// self 占墙钟时间百分比
    self_pct: f64,
    /// 源码文件
    source_file: Option<String>,
    /// 源码行号
    source_line: Option<u32>,
}

/// 单个 HTTP 请求内的 span 聚合
#[derive(Debug, Default)]
pub struct RequestProfile {
    /// span id → 记录
    spans: HashMap<tracing::Id, SpanRecord>,
    /// 创建顺序
    order: Vec<tracing::Id>,
}

impl RequestProfile {
    /// 创建空采集器
    pub fn new() -> Self {
        Self::default()
    }

    /// span 进入时登记
    pub fn on_enter(
        &mut self,
        id: tracing::Id,
        label: String,
        parent: Option<tracing::Id>,
        source_file: Option<String>,
        source_line: Option<u32>,
    ) {
        if self.spans.contains_key(&id) {
            return;
        }
        self.order.push(id.clone());
        self.spans.insert(
            id,
            SpanRecord {
                label,
                started: Instant::now(),
                ended: None,
                parent,
                source_file,
                source_line,
            },
        );
    }

    /// span 关闭时记录结束时刻
    pub fn on_close(&mut self, id: tracing::Id) {
        if let Some(record) = self.spans.get_mut(&id) {
            record.ended = Some(Instant::now());
        }
    }

    /// 收集可见 span 节点
    fn collect_print_nodes(&self, wall: Duration, min_show: Duration) -> Vec<PrintNode> {
        if self.spans.is_empty() {
            return Vec::new();
        }
        let wall_ms = wall.as_secs_f64() * 1000.0;
        let roots = self.root_ids();
        let mut nodes = Vec::new();
        for root in &roots {
            self.collect_nodes(root, 0, wall_ms, min_show, &mut nodes);
        }
        nodes
    }

    /// 构建结构化 profile 报告
    pub fn build_report(
        &self,
        method: &Method,
        uri: &Uri,
        status: StatusCode,
        wall: Duration,
        min_show: Duration,
    ) -> Option<ProfileReport> {
        let nodes = self.collect_print_nodes(wall, min_show);
        if nodes.is_empty() {
            return None;
        }
        let wall_ms = wall.as_secs_f64() * 1000.0;
        let endpoint = format_uri(uri);
        let hotspot = nodes
            .iter()
            .max_by(|a, b| {
                a.self_ms
                    .partial_cmp(&b.self_ms)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .filter(|n| n.self_ms > 0.0)
            .map(|n| ProfileHotspot {
                label: n.label.clone(),
                self_ms: n.self_ms,
                self_pct: n.self_pct,
                file: n.source_file.clone(),
                line: n.source_line,
            });
        Some(ProfileReport {
            method: method.to_string(),
            uri: endpoint,
            wall_ms,
            status: status.as_u16(),
            hotspot,
            spans: nodes
                .iter()
                .map(|n| ProfileSpanNode {
                    label: n.label.clone(),
                    self_ms: n.self_ms,
                    inclusive_ms: n.inclusive_ms,
                    self_pct: n.self_pct,
                    file: n.source_file.clone(),
                    line: n.source_line,
                })
                .collect(),
        })
    }

    /// 格式化调用栈耗时树（不直接写 stderr，供整段原子输出）
    pub fn format_report(
        &self,
        method: &Method,
        uri: &Uri,
        wall: Duration,
        min_show: Duration,
    ) -> Vec<String> {
        let nodes = self.collect_print_nodes(wall, min_show);
        if nodes.is_empty() {
            return Vec::new();
        }
        HttpLogStyle::enable_ansi_support();
        let style = HttpLogStyle::detect();

        let hotspot_idx = nodes
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| {
                a.self_ms
                    .partial_cmp(&b.self_ms)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .filter(|(_, n)| n.self_ms > 0.0)
            .map(|(i, _)| i);

        let mut lines = Vec::new();
        lines.push(style.divider());
        lines.push(format!(
            "{} {} {}  {}",
            style.label("prof │"),
            style.method(method),
            style.endpoint(uri),
            style.elapsed(wall),
        ));

        for (i, node) in nodes.iter().enumerate() {
            let is_hot = hotspot_idx == Some(i);
            lines.extend(Self::format_node(node, is_hot, &style));
        }

        if let Some(idx) = hotspot_idx {
            if let Some(node) = nodes.get(idx) {
                let hint = format!(
                    "hotspot {}  self {}  ({:.1}%)",
                    node.label,
                    format_ms(node.self_ms),
                    node.self_pct,
                );
                lines.push(format!(
                    "{} {}",
                    style.label("prof │"),
                    style.hotspot(&hint),
                ));
                if let Some(jb) = jetbrains_filter_line(node.source_file.as_deref(), node.source_line)
                {
                    lines.push(jb);
                }
            }
        }

        lines.push(style.divider());
        lines
    }

    /// 输出调用栈耗时树到 stderr
    pub fn print_report(
        &self,
        method: &Method,
        uri: &Uri,
        wall: Duration,
        min_show: Duration,
    ) {
        HttpLogStyle::print_lines(&self.format_report(method, uri, wall, min_show));
    }

    /// 生成 `Server-Timing` 响应头（dev 前后端联调）
    pub fn format_server_timing(&self, wall: Duration, min_show: Duration) -> Option<String> {
        if !server_timing_enabled() {
            return None;
        }
        let nodes = self.collect_print_nodes(wall, min_show);
        if nodes.is_empty() {
            return None;
        }
        let wall_ms = wall.as_secs_f64() * 1000.0;
        let mut parts = vec![format!("total;dur={wall_ms:.2}")];

        let hotspot_idx = nodes
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| {
                a.self_ms
                    .partial_cmp(&b.self_ms)
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .filter(|(_, n)| n.self_ms > 0.0)
            .map(|(i, _)| i);

        if let Some(idx) = hotspot_idx {
            let node = &nodes[idx];
            parts.push(server_timing_metric("hot", node.self_ms, &node.label));
        }

        let mut ranked: Vec<_> = nodes
            .iter()
            .enumerate()
            .filter(|(i, n)| Some(*i) != hotspot_idx && n.self_ms > 0.0)
            .collect();
        ranked.sort_by(|(_, a), (_, b)| {
            b.self_ms
                .partial_cmp(&a.self_ms)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        for (seq, (_, node)) in ranked
            .into_iter()
            .take(SERVER_TIMING_EXTRA)
            .enumerate()
        {
            parts.push(server_timing_metric(
                &format!("s{}", seq + 1),
                node.self_ms,
                &node.label,
            ));
        }

        Some(parts.join(", "))
    }

    /// 取根 span（父不在表内或无父）
    fn root_ids(&self) -> Vec<tracing::Id> {
        self.order
            .iter()
            .filter(|id| {
                self.spans
                    .get(*id)
                    .and_then(|s| s.parent.as_ref())
                    .map(|p| !self.spans.contains_key(p))
                    .unwrap_or(true)
            })
            .cloned()
            .collect()
    }

    /// 递归收集可见节点（跳过 http.request 容器）
    fn collect_nodes(
        &self,
        id: &tracing::Id,
        depth: usize,
        wall_ms: f64,
        min_show: Duration,
        out: &mut Vec<PrintNode>,
    ) {
        let Some(record) = self.spans.get(id) else {
            return;
        };

        if record.label.ends_with("::http.request") || record.label == "http.request" {
            for child in self.direct_children(id) {
                self.collect_nodes(&child, depth, wall_ms, min_show, out);
            }
            return;
        }

        let inclusive = record.duration();
        let inclusive_ms = inclusive.as_secs_f64() * 1000.0;
        let children_ms: f64 = self
            .direct_children(id)
            .iter()
            .filter_map(|cid| self.spans.get(cid))
            .map(|c| c.duration().as_secs_f64() * 1000.0)
            .sum();
        let self_ms = (inclusive_ms - children_ms).max(0.0);

        if Duration::from_secs_f64(self_ms / 1000.0) < min_show {
            for child in self.direct_children(id) {
                self.collect_nodes(&child, depth + 1, wall_ms, min_show, out);
            }
            return;
        }

        let self_pct = if wall_ms > 0.0 {
            self_ms / wall_ms * 100.0
        } else {
            0.0
        };
        out.push(PrintNode {
            depth,
            label: record.label.clone(),
            self_ms,
            inclusive_ms,
            self_pct,
            source_file: record.source_file.clone(),
            source_line: record.source_line,
        });

        for child in self.direct_children(id) {
            self.collect_nodes(&child, depth + 1, wall_ms, min_show, out);
        }
    }

    /// 格式化单行节点
    fn format_node(node: &PrintNode, is_hot: bool, style: &HttpLogStyle) -> Vec<String> {
        let branch = tree_branch(node.depth);
        let label_col = format_method_label(
            style,
            &node.label,
            node.source_file.as_deref(),
            node.source_line,
        );

        let hot_mark = if is_hot { " ◀ hot" } else { "" };
        let line = format!(
            "{branch}{label_col}  self {:>tw$}  incl {:>tw$}  {:>5.1}%{hot_mark}",
            format_ms(node.self_ms),
            format_ms(node.inclusive_ms),
            node.self_pct,
            tw = TIME_WIDTH,
        );

        let mut lines = vec![format!(
            "{} {}",
            style.label("prof │"),
            style.prof_text(&line, is_hot),
        )];
        if let Some(jb) = jetbrains_filter_line(node.source_file.as_deref(), node.source_line) {
            lines.push(jb);
        }
        lines
    }

    /// 直接子 span
    fn direct_children(&self, id: &tracing::Id) -> Vec<tracing::Id> {
        self.order
            .iter()
            .filter(|child| {
                self.spans
                    .get(*child)
                    .and_then(|s| s.parent.as_ref())
                    == Some(id)
            })
            .cloned()
            .collect()
    }
}

impl SpanRecord {
    /// span inclusive 耗时
    fn duration(&self) -> Duration {
        self.ended
            .map(|end| end.duration_since(self.started))
            .unwrap_or_else(|| self.started.elapsed())
    }
}

tokio::task_local! {
    /// 当前请求的耗时采集器
    static REQUEST_PROFILE: Arc<Mutex<RequestProfile>>;
}

/// 是否处于请求 profiling 上下文
pub fn is_active() -> bool {
    REQUEST_PROFILE.try_with(|_| ()).is_ok()
}

/// 在请求 profiling 作用域内执行 future
pub async fn scope<F>(profile: Arc<Mutex<RequestProfile>>, fut: F) -> F::Output
where
    F: std::future::Future,
{
    REQUEST_PROFILE.scope(profile, fut).await
}

/// 写入 span 进入事件（Layer 调用）
pub fn record_enter(
    id: tracing::Id,
    label: String,
    parent: Option<tracing::Id>,
    source_file: Option<&str>,
    source_line: Option<u32>,
) {
    let _ = REQUEST_PROFILE.try_with(|cell| {
        if let Ok(mut guard) = cell.lock() {
            guard.on_enter(
                id,
                label,
                parent,
                source_file.map(str::to_string),
                source_line,
            );
        }
    });
}

/// 写入 span 关闭事件（Layer 调用）
pub fn record_close(id: tracing::Id) {
    let _ = REQUEST_PROFILE.try_with(|cell| {
        if let Ok(mut guard) = cell.lock() {
            guard.on_close(id);
        }
    });
}

/// 从 tracing metadata 组装展示名
pub fn span_label(meta: &tracing::Metadata<'_>) -> String {
    let target = meta.target();
    let short = target.rsplit("::").take(2).collect::<Vec<_>>();
    let module = short.into_iter().rev().collect::<Vec<_>>().join("::");
    format!("{}::{}", module, meta.name())
}

/// 是否输出 Server-Timing 响应头（`DEV_SERVER_TIMING=0` 关闭）
fn server_timing_enabled() -> bool {
    std::env::var("DEV_SERVER_TIMING").ok().as_deref() != Some("0")
}

/// 组装单条 Server-Timing metric
fn server_timing_metric(name: &str, dur_ms: f64, label: &str) -> String {
    let desc = escape_server_timing_desc(&server_timing_desc(label));
    format!("{name};dur={dur_ms:.2};desc=\"{desc}\"")
}

/// 缩短 span 展示名（取末尾两段 `::`）
fn server_timing_desc(label: &str) -> String {
    let parts: Vec<&str> = label.split("::").collect();
    if parts.len() <= 2 {
        label.to_string()
    } else {
        parts[parts.len() - 2..].join("::")
    }
}

/// 转义 Server-Timing desc 中的引号与反斜杠
fn escape_server_timing_desc(s: &str) -> String {
    s.replace('\\', "\\\\").replace('"', "\\\"")
}

/// 读取最小展示阈值（毫秒），默认 0.1ms
pub fn min_show_duration() -> Duration {
    std::env::var("DEV_PROFILE_MIN_MS")
        .ok()
        .and_then(|v| v.parse::<f64>().ok())
        .map(|ms| Duration::from_secs_f64(ms / 1000.0))
        .unwrap_or_else(|| Duration::from_secs_f64(0.0001))
}

/// 将 profile 报告追加到 JSONL 文件（`DEV_PROFILE_JSONL` 环境变量）
pub fn append_jsonl_report(report: &ProfileReport) {
    let Ok(path) = std::env::var("DEV_PROFILE_JSONL") else {
        return;
    };
    let Ok(line) = serde_json::to_string(report) else {
        return;
    };
    HttpLogStyle::append_jsonl(&path, &line);
}

/// 格式化 URI（含 query）
fn format_uri(uri: &Uri) -> String {
    match uri.query() {
        Some(q) => format!("{}?{q}", uri.path()),
        None => uri.path().to_string(),
    }
}

/// 创建 HTTP 根 span
pub fn http_request_span(method: &Method, uri: &Uri) -> Span {
    tracing::debug_span!(
        "http.request",
        method = %method,
        uri = %uri,
    )
}

/// 树形前缀
fn tree_branch(depth: usize) -> String {
    if depth == 0 {
        "├─ ".to_string()
    } else {
        format!("{}└─ ", "│  ".repeat(depth.saturating_sub(1)))
    }
}

/// 格式化毫秒
fn format_ms(ms: f64) -> String {
    if ms >= 100.0 {
        format!("{ms:.0}ms")
    } else if ms >= 10.0 {
        format!("{ms:.1}ms")
    } else {
        format!("{ms:.2}ms")
    }
}

/// 链接协议
enum LinkScheme {
    /// JetBrains / RustRover（独立一行 ` at path:line:col`）
    JetBrains,
    /// Cursor
    Cursor,
    /// VS Code
    Vscode,
    /// VS Code Insiders
    VscodeInsiders,
    /// file://
    File,
    /// 禁用
    None,
}

impl LinkScheme {
    /// 解析链接协议（环境变量 > 终端自动检测）
    fn detect() -> Self {
        if std::env::var("DEV_PROFILE_NO_LINKS").is_ok() {
            return Self::None;
        }
        match std::env::var("DEV_PROFILE_LINK_SCHEME")
            .unwrap_or_default()
            .to_ascii_lowercase()
            .as_str()
        {
            "none" => Self::None,
            "jetbrains" | "rustrover" | "idea" | "goland" | "webstorm" => Self::JetBrains,
            "vscode" => Self::Vscode,
            "vscode-insiders" => Self::VscodeInsiders,
            "file" => Self::File,
            "cursor" => Self::Cursor,
            "" => {
                let emulator = std::env::var("TERMINAL_EMULATOR").unwrap_or_default();
                if emulator.contains("JetBrains") {
                    Self::JetBrains
                } else if std::env::var("TERM_PROGRAM").unwrap_or_default() == "vscode" {
                    Self::Cursor
                } else {
                    Self::JetBrains
                }
            }
            _ => Self::Cursor,
        }
    }
}

/// 方法名（Cursor 等 IDE 可 OSC 8 跳转）
fn format_method_label(
    style: &HttpLogStyle,
    method: &str,
    file: Option<&str>,
    line: Option<u32>,
) -> String {
    if matches!(LinkScheme::detect(), LinkScheme::JetBrains) {
        return method.to_string();
    }
    let Some(uri) = build_osc8_uri(file, line) else {
        return method.to_string();
    };
    style.link(&uri, method)
}

/// 生成 JetBrains 过滤器可识别的整行文本
fn jetbrains_filter_line(file: Option<&str>, line: Option<u32>) -> Option<String> {
    if !matches!(LinkScheme::detect(), LinkScheme::JetBrains) {
        return None;
    }
    let file = file?;
    let line = line?;
    let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let path = resolve_source_path(&manifest, file);
    let rel = project_relative_path(&path, &manifest);
    let style = std::env::var("DEV_PROFILE_JB_STYLE").unwrap_or_else(|_| "rustc".into());
    Some(match style.as_str() {
        // RsBacktraceFilter: ^\s+at FILE:LINE(:COL)?$
        "stack" | "at" => format!(" at {rel}:{line}:1"),
        // RsConsoleFilter: ^(?:\s+--> )?FILE:LINE(:COL)?.*$
        _ => format!("   --> {rel}:{line}:1"),
    })
}

/// crate 内相对路径（正斜杠）
fn project_relative_path(path: &Path, manifest: &Path) -> String {
    path.strip_prefix(manifest)
        .map(path_to_link)
        .unwrap_or_else(|_| path_to_link(path))
}

/// 构建 OSC 8 URI（Cursor / VS Code）
fn build_osc8_uri(file: Option<&str>, line: Option<u32>) -> Option<String> {
    let file = file?;
    let line = line?;
    let scheme = LinkScheme::detect();
    if matches!(scheme, LinkScheme::None | LinkScheme::JetBrains) {
        return None;
    }
    let manifest = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
    let path = resolve_source_path(&manifest, file);
    let abs = path_to_link(&path);
    Some(match scheme {
        LinkScheme::File => format!("file://{abs}:{line}:1"),
        LinkScheme::Vscode => format!("vscode://file/{abs}:{line}:1"),
        LinkScheme::VscodeInsiders => format!("vscode-insiders://file/{abs}:{line}:1"),
        LinkScheme::Cursor => format!("cursor://file/{abs}:{line}:1"),
        LinkScheme::JetBrains | LinkScheme::None => return None,
    })
}

/// 路径转链接用字符串（去掉 Windows \\?\ 扩展前缀）
fn path_to_link(path: &Path) -> String {
    let raw = path.to_string_lossy();
    let stripped = raw
        .strip_prefix(r"\\?\")
        .or_else(|| raw.strip_prefix("//?/"))
        .unwrap_or(&raw);
    stripped.replace('\\', "/")
}

/// 解析 instrument 返回的 file 路径
fn resolve_source_path(manifest: &Path, file: &str) -> PathBuf {
    let p = PathBuf::from(file);
    if p.is_absolute() {
        p
    } else {
        manifest.join(p)
    }
}