microchassis/profiling/
jeprof.rs

1// Copyright 2023 Folke Behrens
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Contains HTTP handler for jeprof support (/pprof/heap).
16//! Based on <https://gperftools.github.io/gperftools/pprof_remote_servers.html>,
17//! <https://jemalloc.net/jemalloc.3.html#mallctl_namespace>,
18//! <https://github.com/jemalloc/jemalloc/blob/master/bin/jeprof.in>.
19
20use crate::profiling::mallctl;
21use http::{header, Method, Request, Response, StatusCode};
22use std::{collections::HashMap, env, fmt};
23
24#[inline]
25pub fn router(req: Request<Vec<u8>>) -> http::Result<Response<Vec<u8>>> {
26    match (req.method(), req.uri().path()) {
27        (&Method::GET, "/pprof/conf") => JeprofHandler(get_pprof_conf_handler).call(req),
28        (&Method::POST, "/pprof/conf") => JeprofHandler(post_pprof_conf_handler).call(req),
29        (&Method::GET, "/pprof/heap") => JeprofHandler(get_pprof_heap_handler).call(req),
30        (&Method::GET, "/pprof/cmdline") => JeprofHandler(get_pprof_cmdline_handler).call(req),
31        (&Method::GET, "/pprof/symbol") => JeprofHandler(get_pprof_symbol_handler).call(req),
32        (&Method::POST, "/pprof/symbol") => JeprofHandler(post_pprof_symbol_handler).call(req),
33        (&Method::GET, "/pprof/stats") => JeprofHandler(get_pprof_stats_handler).call(req),
34        _ => {
35            let body = b"Bad Request\r\n";
36            Response::builder()
37                .status(StatusCode::BAD_REQUEST)
38                .header(header::CONTENT_TYPE, "application/octet-stream")
39                .header(header::CONTENT_LENGTH, body.len())
40                .body(body.to_vec())
41        }
42    }
43}
44
45#[cfg(feature = "actix-handlers")]
46#[inline]
47pub fn actix_routes(cfg: &mut actix_web::web::ServiceConfig) {
48    cfg.service(
49        actix_web::web::scope("/pprof")
50            .route("/conf", actix_web::web::get().to(JeprofHandler(get_pprof_conf_handler)))
51            .route("/conf", actix_web::web::post().to(JeprofHandler(post_pprof_conf_handler)))
52            .route("/heap", actix_web::web::get().to(JeprofHandler(get_pprof_heap_handler)))
53            .route("/cmdline", actix_web::web::get().to(JeprofHandler(get_pprof_cmdline_handler)))
54            .route("/symbol", actix_web::web::get().to(JeprofHandler(get_pprof_symbol_handler)))
55            .route("/symbol", actix_web::web::post().to(JeprofHandler(post_pprof_symbol_handler)))
56            .route("/stats", actix_web::web::get().to(JeprofHandler(get_pprof_stats_handler))),
57    );
58}
59
60#[derive(Debug)]
61pub struct ErrorResponse(String);
62
63impl fmt::Display for ErrorResponse {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        write!(f, "ERROR: {}", self.0)
66    }
67}
68
69#[cfg(feature = "actix-handlers")]
70impl actix_web::ResponseError for ErrorResponse {}
71
72#[derive(Clone, Debug)]
73struct JeprofHandler<F>(F)
74where
75    F: Fn(&[u8], &HashMap<String, String>) -> Result<(Vec<u8>, Option<String>), ErrorResponse>
76        + Clone
77        + 'static;
78
79impl<F> JeprofHandler<F>
80where
81    F: Fn(&[u8], &HashMap<String, String>) -> Result<(Vec<u8>, Option<String>), ErrorResponse>
82        + Clone
83        + 'static,
84{
85    fn call(&self, req: Request<Vec<u8>>) -> http::Result<Response<Vec<u8>>> {
86        let params: HashMap<String, String> = parse_malloc_conf_query(req.uri().query())
87            .iter()
88            .map(|(k, v)| ((*k).to_string(), v.unwrap_or_default().to_string()))
89            .collect();
90        match self.0(req.body(), &params) {
91            Ok((body, Some(content_disposition))) => response_ok_binary(body, &content_disposition),
92            Ok((body, None)) => response_ok(body),
93            Err(err) => response_err(&err.0),
94        }
95    }
96}
97
98#[cfg(feature = "actix-handlers")]
99impl<F>
100    actix_web::Handler<(actix_web::web::Payload, actix_web::web::Query<HashMap<String, String>>)>
101    for JeprofHandler<F>
102where
103    F: Fn(&[u8], &HashMap<String, String>) -> Result<(Vec<u8>, Option<String>), ErrorResponse>
104        + Clone
105        + 'static,
106{
107    type Output = Result<actix_web::HttpResponse, ErrorResponse>;
108    type Future = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output>>>;
109
110    fn call(
111        &self,
112        (mut body, query): (
113            actix_web::web::Payload,
114            actix_web::web::Query<HashMap<String, String>>,
115        ),
116    ) -> Self::Future {
117        use futures_util::StreamExt as _;
118
119        let f = self.0.clone();
120        Box::pin(async move {
121            let mut data = Vec::<u8>::new();
122            while let Some(item) = body.next().await {
123                data.extend_from_slice(&item.map_err(|e| ErrorResponse(e.to_string()))?);
124            }
125            f(&data, &query.0).map(|(body, content_disposition)| {
126                let mut resp = actix_web::HttpResponse::Ok();
127                if let Some(filename) = content_disposition {
128                    resp.insert_header(actix_web::http::header::ContentDisposition::attachment(
129                        filename,
130                    ));
131                }
132                resp.body(actix_web::web::Bytes::from(body))
133            })
134        })
135    }
136}
137
138#[inline]
139pub fn get_pprof_conf_handler(
140    _body: &[u8],
141    _params: &HashMap<String, String>,
142) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
143    match mallctl::enabled() {
144        Ok(true) => (),
145        _ => return Err(ErrorResponse("jemalloc profiling not enabled".to_owned())),
146    };
147
148    let Ok(state) = mallctl::active() else {
149        return Err(ErrorResponse("failed to read prof.active\r\n".to_owned()));
150    };
151    let Ok(sample) = mallctl::sample_interval() else {
152        return Err(ErrorResponse("failed to read prof.lg_sample\r\n".to_owned()));
153    };
154    let body = format!("prof.active:{state},prof.lg_sample:{sample}\r\n");
155    Ok((body.into_bytes(), None))
156}
157
158#[inline]
159pub fn post_pprof_conf_handler(
160    _body: &[u8],
161    params: &HashMap<String, String>,
162) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
163    match mallctl::enabled() {
164        Ok(true) => (),
165        _ => return Err(ErrorResponse("jemalloc profiling not enabled\r\n".to_owned())),
166    };
167
168    for (name, value) in params {
169        if let Err(e) = match name.as_str() {
170            "prof.reset" => {
171                let sample = value.parse().map_err(|_| {
172                    ErrorResponse(format!("invalid prof.reset value: {value:?}\r\n"))
173                })?;
174                mallctl::reset(Some(sample))
175            }
176            "prof.active" => {
177                let Some(state) = value.parse().ok() else {
178                    return Err(ErrorResponse(format!("invalid prof.active value: {value:?}\r\n")));
179                };
180                mallctl::set_active(state)
181            }
182            _ => {
183                return Err(ErrorResponse(format!("{name}={value:?} unknown\r\n")));
184            }
185        } {
186            return Err(ErrorResponse(format!("{name}={value:?} failed: {e}\r\n")));
187        }
188    }
189
190    Ok((b"OK\r\n".to_vec(), None))
191}
192
193#[inline]
194pub fn get_pprof_heap_handler(
195    _body: &[u8],
196    _params: &HashMap<String, String>,
197) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
198    match mallctl::enabled() {
199        Ok(true) => (),
200        _ => return Err(ErrorResponse("jemalloc profiling not enabled\r\n".to_owned())),
201    };
202
203    let Ok(f) = tempfile::Builder::new().prefix("jemalloc.").suffix(".prof").tempfile() else {
204        return Err(ErrorResponse("cannot create temporary file for profile dump\r\n".to_owned()));
205    };
206
207    let Ok(profile) = mallctl::dump(f.path().to_str()) else {
208        return Err(ErrorResponse("failed to dump profile\r\n".to_owned()));
209    };
210
211    let filename = f.path().file_name().expect("proper filename from tempfile");
212    Ok((profile.expect("profile not None"), Some(filename.to_string_lossy().to_string())))
213}
214
215/// HTTP handler for GET /pprof/cmdline.
216#[inline]
217pub fn get_pprof_cmdline_handler(
218    _body: &[u8],
219    _params: &HashMap<String, String>,
220) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
221    let mut body = String::new();
222    for arg in env::args() {
223        body.push_str(arg.as_str());
224        body.push_str("\r\n");
225    }
226    Ok((body.into_bytes(), None))
227}
228
229/// HTTP handler for GET /pprof/symbol.
230#[inline]
231pub fn get_pprof_symbol_handler(
232    _body: &[u8],
233    _params: &HashMap<String, String>,
234) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
235    // TODO: any quick way to check if binary is stripped?
236    let body = b"num_symbols: 1\r\n";
237    Ok((body.to_vec(), None))
238}
239
240/// HTTP handler for POST /pprof/symbol.
241#[inline]
242pub fn post_pprof_symbol_handler(
243    body: &[u8],
244    _params: &HashMap<String, String>,
245) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
246    fn lookup_symbol(addr: u64) -> Option<String> {
247        let mut s: Option<String> = None;
248        backtrace::resolve(addr as *mut _, |symbol| {
249            s = symbol.name().map(|n| n.to_string());
250        });
251        s
252    }
253
254    let body = String::from_utf8_lossy(body);
255    let addrs = body
256        .split('+')
257        .filter_map(|addr| u64::from_str_radix(addr.trim_start_matches("0x"), 16).ok())
258        .map(|addr| (addr, lookup_symbol(addr)))
259        .filter_map(|(addr, sym)| sym.map(|sym| (addr, sym)));
260
261    let mut body = String::new();
262    for (addr, sym) in addrs {
263        body.push_str(format!("{addr:#x}\t{sym}\r\n").as_str());
264    }
265
266    Ok((body.into_bytes(), None))
267}
268
269/// HTTP handler for GET /pprof/stats.
270#[inline]
271pub fn get_pprof_stats_handler(
272    _body: &[u8],
273    _params: &HashMap<String, String>,
274) -> Result<(Vec<u8>, Option<String>), ErrorResponse> {
275    let body = match mallctl::stats() {
276        Ok(body) => body,
277        Err(e) => return Err(ErrorResponse(format!("failed to print stats: {e}\r\n"))),
278    };
279    Ok((body, None))
280}
281
282fn parse_malloc_conf_query(query: Option<&str>) -> Vec<(&str, Option<&str>)> {
283    query
284        .map(|q| {
285            q.split(',')
286                .map(|kv| kv.splitn(2, ':').collect::<Vec<_>>())
287                .map(|v| match v.len() {
288                    1 => (v[0], None),
289                    2 => (v[0], Some(v[1])),
290                    _ => unreachable!(),
291                })
292                .collect()
293        })
294        .unwrap_or_default()
295}
296
297fn response_ok(body: Vec<u8>) -> http::Result<Response<Vec<u8>>> {
298    Response::builder()
299        .status(StatusCode::OK)
300        .header(header::CONTENT_TYPE, "text/plain; charset=UTF-8")
301        .header(header::CONTENT_LENGTH, body.len())
302        .body(body)
303}
304
305fn response_ok_binary(body: Vec<u8>, filename: &str) -> http::Result<Response<Vec<u8>>> {
306    Response::builder()
307        .status(StatusCode::OK)
308        .header(header::CONTENT_TYPE, "application/octet-stream")
309        .header(header::CONTENT_DISPOSITION, format!("attachment; filename=\"{filename}\""))
310        .header(header::CONTENT_LENGTH, body.len())
311        .body(body)
312}
313
314fn response_err(msg: &str) -> http::Result<Response<Vec<u8>>> {
315    Response::builder()
316        .status(StatusCode::BAD_REQUEST)
317        .header(header::CONTENT_TYPE, "text/plain; charset=UTF-8")
318        .header(header::CONTENT_LENGTH, msg.len())
319        .body(msg.as_bytes().to_owned())
320}
microchassis/profiling/jeprof.rs

microchassis/profiling/
jeprof.rs
