bee_tui/
pprof_bundle.rs

1//! Async pprof + trace bundling for `:diagnose --pprof[=N]`.
2//!
3//! Bee exposes Go's standard `/debug/pprof/profile` (CPU profile) and
4//! `/debug/pprof/trace` (execution trace) endpoints when started with
5//! `--debug-api-enable=true`. Each one blocks for `seconds=N` while
6//! the runtime samples; both are independent so we run them in
7//! parallel.
8//!
9//! The result is a directory next to the existing `.txt` diagnostic
10//! file that the operator can `tar -czf` and ship. We deliberately do
11//! NOT pull in a `tar` dependency — the directory is self-contained
12//! and the operator's shell already knows how to bundle it.
13
14use std::path::PathBuf;
15use std::time::Duration;
16
17use reqwest::StatusCode;
18
19/// Result of a successful pprof+trace fetch. The directory contains
20/// `profile.pprof` + `trace.pprof` (Go's pprof binary format —
21/// readable by `go tool pprof` and `pprof` web).
22#[derive(Debug, Clone)]
23pub struct PprofBundle {
24    pub dir: PathBuf,
25    pub profile_bytes: u64,
26    pub trace_bytes: u64,
27    pub seconds: u32,
28}
29
30impl PprofBundle {
31    /// Operator-friendly summary line.
32    pub fn summary(&self) -> String {
33        format!(
34            "pprof bundle ready · {}s sampling · profile {} · trace {} → {}",
35            self.seconds,
36            human_bytes(self.profile_bytes),
37            human_bytes(self.trace_bytes),
38            self.dir.display(),
39        )
40    }
41}
42
43/// Fetch `/debug/pprof/profile?seconds=N` + `/debug/pprof/trace?seconds=N`
44/// in parallel and write each to `dir/`. The returned bundle's
45/// summary line is the canonical operator-facing message.
46pub async fn fetch_and_write(
47    base_url: &str,
48    auth_token: Option<String>,
49    seconds: u32,
50    dir: PathBuf,
51) -> Result<PprofBundle, String> {
52    std::fs::create_dir_all(&dir).map_err(|e| format!("mkdir {}: {e}", dir.display()))?;
53    // Generous timeout: pprof blocks for `seconds` then streams; add
54    // 30s headroom for trailing read latency.
55    let client = reqwest::Client::builder()
56        .timeout(Duration::from_secs(u64::from(seconds) + 30))
57        .build()
58        .map_err(|e| format!("client build: {e}"))?;
59
60    let base = base_url.trim_end_matches('/').to_string();
61    let profile_url = format!("{base}/debug/pprof/profile?seconds={seconds}");
62    let trace_url = format!("{base}/debug/pprof/trace?seconds={seconds}");
63
64    let (profile, trace) = tokio::join!(
65        fetch_one(&client, &profile_url, auth_token.as_deref()),
66        fetch_one(&client, &trace_url, auth_token.as_deref()),
67    );
68
69    let profile_bytes = profile.map_err(|e| format!("profile: {e}"))?;
70    let trace_bytes = trace.map_err(|e| format!("trace: {e}"))?;
71
72    let profile_path = dir.join("profile.pprof");
73    let trace_path = dir.join("trace.pprof");
74    std::fs::write(&profile_path, &profile_bytes).map_err(|e| format!("write profile: {e}"))?;
75    std::fs::write(&trace_path, &trace_bytes).map_err(|e| format!("write trace: {e}"))?;
76
77    Ok(PprofBundle {
78        dir,
79        profile_bytes: profile_bytes.len() as u64,
80        trace_bytes: trace_bytes.len() as u64,
81        seconds,
82    })
83}
84
85async fn fetch_one(
86    client: &reqwest::Client,
87    url: &str,
88    auth: Option<&str>,
89) -> Result<Vec<u8>, String> {
90    let mut req = client.get(url);
91    if let Some(token) = auth {
92        req = req.bearer_auth(token);
93    }
94    let resp = req.send().await.map_err(|e| format!("GET {url}: {e}"))?;
95    let status = resp.status();
96    if status == StatusCode::NOT_FOUND {
97        return Err("404 — Bee's debug API isn't enabled. Add `--debug-api-enable=true` to your Bee start args (or set `debug-api-enable: true` in bee.yaml) and restart.".to_string());
98    }
99    if !status.is_success() {
100        return Err(format!("{url} returned HTTP {status}"));
101    }
102    let bytes = resp
103        .bytes()
104        .await
105        .map_err(|e| format!("read body of {url}: {e}"))?;
106    Ok(bytes.to_vec())
107}
108
109fn human_bytes(n: u64) -> String {
110    const KIB: u64 = 1 << 10;
111    const MIB: u64 = 1 << 20;
112    if n >= MIB {
113        format!("{:.1} MiB", n as f64 / MIB as f64)
114    } else if n >= KIB {
115        format!("{:.1} KiB", n as f64 / KIB as f64)
116    } else {
117        format!("{n} B")
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn human_bytes_picks_unit() {
127        assert_eq!(human_bytes(0), "0 B");
128        assert_eq!(human_bytes(1023), "1023 B");
129        assert_eq!(human_bytes(2048), "2.0 KiB");
130        assert_eq!(human_bytes(2 * 1024 * 1024), "2.0 MiB");
131    }
132
133    #[test]
134    fn summary_line_includes_dir_and_byte_counts() {
135        let b = PprofBundle {
136            dir: PathBuf::from("/tmp/bee-tui-diagnostic-1234"),
137            profile_bytes: 12_345,
138            trace_bytes: 6_789,
139            seconds: 60,
140        };
141        let s = b.summary();
142        assert!(s.contains("60s sampling"));
143        assert!(s.contains("/tmp/bee-tui-diagnostic-1234"));
144        assert!(s.contains("KiB"));
145    }
146}
bee_tui/pprof_bundle.rs

bee_tui/
pprof_bundle.rs
