krabka_telemetry/profiling/
routes.rs

1//! The pprof HTTP routes and the sampling they drive.
2//!
3//! The module holds the `GET /debug/pprof/profile` CPU handler, the optional
4//! `GET /debug/pprof/heap` handler behind the `heap-profiling` feature, the
5//! gzip step that produces the pprof file format, and the router builders that
6//! bind a `ProfilingConfig` to those handlers as axum state.
7
8#[cfg(any(unix, test))]
9use krabka_units::{Time, secs};
10use {
11    super::config::{ProfilingConfig, ProfilingError},
12    axum::{
13        Router,
14        extract::{Query, State},
15        http::StatusCode,
16        response::{IntoResponse, Response},
17        routing::get,
18    },
19    flate2::write::GzEncoder,
20    serde::Deserialize,
21    tokio::{task, time},
22};
23#[cfg(unix)]
24use {krabka_units::convert::TimeExt as _, std::io};
25
26#[derive(Debug, Deserialize)]
27struct CpuQuery {
28    #[cfg_attr(not(unix), allow(dead_code))]
29    seconds: Option<u64>,
30}
31
32#[cfg(all(unix, feature = "heap-profiling"))]
33#[derive(Debug, Deserialize)]
34struct HeapQuery {
35    seconds: Option<u64>,
36}
37
38/// CPU profile in pprof protobuf, sampled for `?seconds=N`.
39///
40/// The default is 30 seconds, and the default configuration clamps the value
41/// to `1..=60` seconds.
42#[cfg(unix)]
43async fn cpu_profile(State(config): State<ProfilingConfig>, Query(q): Query<CpuQuery>) -> Response {
44    let duration = requested_duration(
45        q.seconds,
46        config.profiling_cpu_default_duration,
47        config.profiling_cpu_max_duration,
48    );
49    let guard = match profile_work(move || start_cpu_profile(&config)).await {
50        Ok(guard) => guard,
51        Err(error) => return (StatusCode::INTERNAL_SERVER_ERROR, error).into_response(),
52    };
53    time::sleep(duration.to_std()).await;
54    match profile_work(move || finish_cpu_profile(guard)).await {
55        Ok(body) => (
56            StatusCode::OK,
57            [("content-type", "application/octet-stream")],
58            body,
59        )
60            .into_response(),
61        Err(error) => (StatusCode::INTERNAL_SERVER_ERROR, error).into_response(),
62    }
63}
64
65/// Profiler setup, symbolization, and encoding can take hundreds of milliseconds. Offload that synchronous work so other requests and timers can progress, including on a current-thread runtime.
66#[cfg(unix)]
67async fn profile_work<T: Send + 'static>(
68    work: impl FnOnce() -> Result<T, String> + Send + 'static,
69) -> Result<T, String> {
70    task::spawn_blocking(work)
71        .await
72        .map_err(|error| format!("profiling worker: {error}"))?
73}
74
75#[cfg(unix)]
76fn start_cpu_profile(config: &ProfilingConfig) -> Result<pprof::ProfilerGuard<'static>, String> {
77    let blocklist = config
78        .profiling_native_frame_blocklist
79        .iter()
80        .map(String::as_str)
81        .collect::<Vec<_>>();
82    pprof::ProfilerGuardBuilder::default()
83        .frequency(config.profiling_cpu_sample_frequency.hertz())
84        .blocklist(&blocklist)
85        .build()
86        .map_err(|error| format!("profiler: {error}"))
87}
88
89#[cfg(unix)]
90fn finish_cpu_profile(guard: pprof::ProfilerGuard<'static>) -> Result<Vec<u8>, String> {
91    // This is the prost Message trait used by pprof's generated Profile.
92    use pprof::protos::Message as _;
93
94    let report = guard
95        .report()
96        .build()
97        .map_err(|error| format!("report: {error}"))?;
98    drop(guard);
99    let profile = report.pprof().map_err(|error| format!("pprof: {error}"))?;
100    gzip(&profile.encode_to_vec()).map_err(|error| format!("gzip: {error}"))
101}
102
103/// Gzip a buffer.
104///
105/// The pprof file format is a gzipped `Profile` protobuf.
106#[cfg(unix)]
107fn gzip(raw: &[u8]) -> io::Result<Vec<u8>> {
108    use std::io::Write as _;
109
110    let mut encoder = GzEncoder::new(
111        Vec::with_capacity(raw.len() / 2),
112        flate2::Compression::fast(),
113    );
114    encoder.write_all(raw)?;
115    encoder.finish()
116}
117
118/// Stub for non-Unix targets: CPU profiling is unavailable.
119// cargo-mutants: non-Unix stub is not built or exercised on the default Linux mutation run.
120#[cfg(not(unix))]
121#[cfg_attr(test, mutants::skip)]
122async fn cpu_profile(
123    _config: State<ProfilingConfig>,
124    _q: Query<CpuQuery>,
125) -> axum::response::Response {
126    (
127        StatusCode::SERVICE_UNAVAILABLE,
128        "CPU profiling requires a Unix target",
129    )
130        .into_response()
131}
132
133// cargo-mutants: optional heap-profiling route is feature-gated out of the default mutation run.
134#[cfg(all(unix, feature = "heap-profiling"))]
135#[cfg_attr(test, mutants::skip)]
136async fn heap_profile(
137    State(config): State<ProfilingConfig>,
138    Query(q): Query<HeapQuery>,
139) -> Response {
140    let Some(ctl) = jemalloc_pprof::PROF_CTL.as_ref() else {
141        return (
142            StatusCode::INTERNAL_SERVER_ERROR,
143            "jemalloc profiling not enabled (build with --features heap-profiling and set MALLOC_CONF)",
144        )
145            .into_response();
146    };
147    let mut ctl = ctl.lock().await;
148    let activated_here = !ctl.activated();
149    if activated_here {
150        if let Err(e) = ctl.activate() {
151            return (
152                StatusCode::INTERNAL_SERVER_ERROR,
153                format!("jemalloc prof activate: {e}"),
154            )
155                .into_response();
156        }
157        let duration = requested_duration(
158            q.seconds,
159            config.profiling_heap_default_duration,
160            config.profiling_heap_max_duration,
161        );
162        time::sleep(duration.to_std()).await;
163    }
164    let dump = ctl.dump_pprof();
165    if activated_here && let Err(e) = ctl.deactivate() {
166        tracing::warn!(error = %e, "could not deactivate jemalloc profiling after heap dump");
167    }
168    drop(ctl);
169    match dump {
170        Ok(pprof) => (
171            StatusCode::OK,
172            [("content-type", "application/octet-stream")],
173            pprof,
174        )
175            .into_response(),
176        Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, format!("heap dump: {e}")).into_response(),
177    }
178}
179
180/// The pprof routes with an explicit policy.
181///
182/// The router always has the CPU route, which returns 503 on non-Unix targets.
183/// The router has the heap route only with the `heap-profiling` feature, and
184/// only on Unix.
185///
186/// # Errors
187/// Returns an error when related profiling duration bounds are invalid.
188pub fn pprof_router_with_config(config: ProfilingConfig) -> Result<Router, ProfilingError> {
189    config.validate().map_err(ProfilingError::Config)?;
190    Ok(pprof_router_unchecked(config))
191}
192
193fn pprof_router_unchecked(config: ProfilingConfig) -> Router {
194    let router = Router::new().route("/debug/pprof/profile", get(cpu_profile));
195    #[cfg(all(unix, feature = "heap-profiling"))]
196    let router = router.route("/debug/pprof/heap", get(heap_profile));
197    router.with_state(config)
198}
199
200/// The pprof routes with the compatible default policy.
201pub fn pprof_router() -> Router { pprof_router_unchecked(ProfilingConfig::default()) }
202
203#[cfg(any(unix, test))]
204fn requested_duration(seconds: Option<u64>, default: Time, maximum: Time) -> Time {
205    seconds
206        .map_or(default, |seconds| {
207            secs(u32::try_from(seconds.max(1)).unwrap_or(u32::MAX))
208        })
209        .min(maximum)
210}
211
212#[cfg(test)]
213mod tests {
214    use {
215        super::*,
216        assert2::assert,
217        krabka_units::{millis, minutes},
218        std::{sync::mpsc, time::Duration},
219        tokio::sync::oneshot,
220    };
221
222    /// A blocked profiler worker must leave this single runtime thread free to run the task that releases it.
223    #[cfg(unix)]
224    #[tokio::test]
225    async fn profiling_work_does_not_block_runtime_progress() {
226        let (started, waiting) = oneshot::channel();
227        let (release, released) = mpsc::channel();
228        let work = profile_work(move || {
229            started.send(()).unwrap();
230            released
231                .recv_timeout(Duration::from_secs(5))
232                .map_err(|error| error.to_string())?;
233            Ok(42)
234        });
235        let other_task = async move {
236            waiting.await.unwrap();
237            release.send(()).unwrap();
238        };
239        let (result, ()) = tokio::join!(work, other_task);
240        assert!(result == Ok(42));
241        assert!(
242            profile_work(|| Err::<(), _>("report failed".into())).await
243                == Err("report failed".into())
244        );
245    }
246
247    #[test]
248    fn requested_profile_duration_uses_configured_default_floor_and_cap() {
249        assert!(requested_duration(None, secs(2), secs(5)) == secs(2));
250        assert!(requested_duration(Some(0), secs(2), secs(5)) == secs(1));
251        assert!(requested_duration(Some(3), secs(2), secs(5)) == secs(3));
252        assert!(requested_duration(Some(9), secs(2), secs(5)) == secs(5));
253        assert!(
254            (ProfilingConfig {
255                profiling_cpu_max_duration: millis(500),
256                ..ProfilingConfig::default()
257            }
258            .validate()
259            .is_err())
260        );
261        assert!(
262            (ProfilingConfig {
263                profiling_heap_default_duration: minutes(1),
264                ..ProfilingConfig::default()
265            }
266            .validate()
267            .is_err())
268        );
269    }
270}
krabka_telemetry/profiling/routes.rs

krabka_telemetry/profiling/
routes.rs
