Skip to main content

stoolap_jemalloc/prof/
pprof.rs

1// Copyright 2026 Stoolap Contributors
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//! Encodes the profile as pprof's `profile.proto`, by hand to stay free
16//! of dependencies
17
18use super::StackCounts;
19use super::maps::{self, Mapping};
20use std::collections::HashMap;
21use std::fmt;
22
23/// Why a heap profile could not be dumped
24#[derive(Debug)]
25pub enum DumpError {
26    /// Profiling was never activated, so there is nothing to dump
27    NotActivated,
28    /// Writing the profile file failed
29    Io(std::io::Error),
30}
31
32impl fmt::Display for DumpError {
33    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
34        match self {
35            DumpError::NotActivated => f.write_str("heap profiling was never activated"),
36            DumpError::Io(e) => write!(f, "writing the heap profile: {e}"),
37        }
38    }
39}
40
41impl std::error::Error for DumpError {}
42
43// Wire types
44const VARINT: u64 = 0;
45const LEN: u64 = 2;
46
47fn varint(b: &mut Vec<u8>, mut v: u64) {
48    while v >= 0x80 {
49        b.push(v as u8 | 0x80);
50        v >>= 7;
51    }
52    b.push(v as u8);
53}
54
55fn tag(b: &mut Vec<u8>, field: u64, wire: u64) {
56    varint(b, field << 3 | wire);
57}
58
59fn uint(b: &mut Vec<u8>, field: u64, v: u64) {
60    if v != 0 {
61        tag(b, field, VARINT);
62        varint(b, v);
63    }
64}
65
66fn bytes(b: &mut Vec<u8>, field: u64, data: &[u8]) {
67    tag(b, field, LEN);
68    varint(b, data.len() as u64);
69    b.extend_from_slice(data);
70}
71
72fn packed(b: &mut Vec<u8>, field: u64, values: impl IntoIterator<Item = u64>) {
73    let mut tmp = Vec::new();
74    for v in values {
75        varint(&mut tmp, v);
76    }
77    bytes(b, field, &tmp);
78}
79
80#[derive(Default)]
81struct Strings {
82    ids: HashMap<String, u64>,
83    list: Vec<String>,
84}
85
86impl Strings {
87    fn id(&mut self, s: &str) -> u64 {
88        if self.list.is_empty() {
89            self.list.push(String::new());
90            self.ids.insert(String::new(), 0);
91        }
92        if let Some(&id) = self.ids.get(s) {
93            return id;
94        }
95        let id = self.list.len() as u64;
96        self.list.push(s.to_string());
97        self.ids.insert(s.to_string(), id);
98        id
99    }
100}
101
102/// `ValueType { type, unit }`
103fn value_type(strings: &mut Strings, ty: &str, unit: &str) -> Vec<u8> {
104    let mut b = Vec::new();
105    uint(&mut b, 1, strings.id(ty));
106    uint(&mut b, 2, strings.id(unit));
107    b
108}
109
110/// A source line of a location: the function's name, simplified for
111/// display and in full, its file and the line number
112struct Line {
113    name: String,
114    system_name: String,
115    file: String,
116    number: u64,
117}
118
119#[cfg(feature = "symbolize")]
120/// Rust's demangled name without generic arguments and trait
121/// qualifications, which pprof would otherwise cut into pieces:
122/// `<core::result::Result<T, E> as FnOnce<A>>::call_once` becomes
123/// `core::result::Result::call_once`
124fn simplify(name: &str) -> String {
125    enum Bracket {
126        Generic,
127        Qualified { trait_part: bool },
128    }
129    let mut out = String::with_capacity(name.len());
130    let mut stack: Vec<Bracket> = Vec::new();
131    let mut rest = name;
132    let mut prev = ' ';
133    while let Some(c) = rest.chars().next() {
134        let skipping = stack.iter().any(|b| {
135            matches!(
136                b,
137                Bracket::Generic | Bracket::Qualified { trait_part: true }
138            )
139        });
140        match c {
141            '<' => {
142                // `Vec<T>`, `f::<T>` and `{closure#0}<T>` are generic;
143                // `<impl Trait for T>` is not
144                let generic = skipping
145                    || (!rest.starts_with("<impl ")
146                        && out.ends_with(|p: char| {
147                            p.is_alphanumeric() || p == '_' || p == ':' || p == '}'
148                        }));
149                if generic && !skipping && out.ends_with("::") {
150                    out.truncate(out.len() - 2);
151                }
152                stack.push(if generic {
153                    Bracket::Generic
154                } else {
155                    Bracket::Qualified { trait_part: false }
156                });
157            }
158            // `->` inside generic arguments is not a closing bracket
159            '>' if !stack.is_empty() && prev != '-' => {
160                stack.pop();
161            }
162            _ if skipping => {}
163            ' ' if rest.starts_with(" as ")
164                && matches!(stack.last(), Some(Bracket::Qualified { .. })) =>
165            {
166                stack.pop();
167                stack.push(Bracket::Qualified { trait_part: true });
168                rest = &rest[" as".len()..];
169            }
170            _ => out.push(c),
171        }
172        rest = &rest[c.len_utf8()..];
173        prev = c;
174    }
175    out
176}
177
178/// Whether a line belongs to the allocator itself: by its source file,
179/// or by its name, for the allocator shims that live in the program.
180/// Builds with `debug = "line-tables-only"` name functions without their
181/// paths, so the name alone does not tell.
182fn is_internal(line: &Line) -> bool {
183    let own_source = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src");
184    std::path::Path::new(&line.file).starts_with(own_source)
185        || ["stoolap_jemalloc::", "__rustc::", "__rust_", "__rg_"]
186            .iter()
187            .any(|p| line.name.starts_with(p))
188}
189
190/// A location's lines without the allocator's own, which inlining puts at
191/// the innermost end; `None` when every line is the allocator's
192fn user_lines(addr: usize) -> Option<Vec<Line>> {
193    let mut lines = symbolize(addr);
194    let internal = lines.iter().take_while(|l| is_internal(l)).count();
195    if internal > 0 && internal == lines.len() {
196        return None;
197    }
198    lines.drain(..internal);
199    Some(lines)
200}
201
202#[cfg(feature = "symbolize")]
203fn symbolize(addr: usize) -> Vec<Line> {
204    let mut lines = Vec::new();
205    // Innermost first, as pprof wants inlined frames ordered
206    backtrace::resolve(addr as *mut core::ffi::c_void, |sym| {
207        let Some(name) = sym.name() else { return };
208        let system_name = format!("{name:#}");
209        lines.push(Line {
210            name: simplify(&system_name),
211            system_name,
212            file: sym
213                .filename()
214                .map(|f| f.display().to_string())
215                .unwrap_or_default(),
216            number: u64::from(sym.lineno().unwrap_or(0)),
217        });
218    });
219    lines
220}
221
222#[cfg(not(feature = "symbolize"))]
223fn symbolize(_addr: usize) -> Vec<Line> {
224    Vec::new()
225}
226
227pub(super) fn encode(stacks: &[StackCounts], interval: usize) -> Vec<u8> {
228    let mut strings = Strings::default();
229    strings.id("");
230    let mut out = Vec::new();
231
232    for (ty, unit) in [
233        ("alloc_objects", "count"),
234        ("alloc_space", "bytes"),
235        ("inuse_objects", "count"),
236        ("inuse_space", "bytes"),
237    ] {
238        let vt = value_type(&mut strings, ty, unit);
239        bytes(&mut out, 1, &vt);
240    }
241
242    let mut mappings = maps::mappings();
243    mappings.sort_by_key(|m| m.start);
244    let mapping_of = |addr: u64| -> u64 {
245        let i = mappings.partition_point(|m| m.start <= addr);
246        if i > 0 && addr < mappings[i - 1].limit {
247            i as u64
248        } else {
249            0
250        }
251    };
252
253    let mut symbols: HashMap<usize, Option<Vec<Line>>> = HashMap::new();
254    let mut locations = Locations::default();
255    // Whether every location in each mapping got functions, file names and
256    // line numbers; the symbolizer misses libraries loaded after its first
257    // use, and those are left for pprof to symbolize. Inline frames come
258    // with the debug information that gives file names.
259    let mut complete = vec![[cfg!(feature = "symbolize"); 3]; mappings.len() + 1];
260
261    for s in stacks {
262        if s.alloc_objects == 0 {
263            continue;
264        }
265        let mut ids = Vec::with_capacity(s.frames.len());
266        for &addr in s.frames {
267            let Some(lines) = symbols.entry(addr).or_insert_with(|| user_lines(addr)) else {
268                // The allocator's frames are only at the top of a stack
269                continue;
270            };
271            let mapping = mapping_of(addr as u64);
272            let c = &mut complete[mapping as usize];
273            c[0] &= !lines.is_empty();
274            c[1] &= !lines.is_empty() && lines.iter().all(|l| !l.file.is_empty());
275            c[2] &= !lines.is_empty() && lines.iter().all(|l| l.number > 0);
276            ids.push(locations.id(addr, lines, mapping, &mut strings));
277        }
278        let mut sample = Vec::new();
279        packed(&mut sample, 1, ids);
280        packed(
281            &mut sample,
282            2,
283            [
284                s.alloc_objects,
285                s.alloc_bytes,
286                s.inuse_objects,
287                s.inuse_bytes,
288            ],
289        );
290        bytes(&mut out, 2, &sample);
291    }
292
293    for (i, m) in mappings.iter().enumerate() {
294        out.extend(encode_mapping(
295            &mut strings,
296            i as u64 + 1,
297            m,
298            complete[i + 1],
299        ));
300    }
301    out.extend(locations.encoded);
302    out.extend(locations.functions_encoded);
303
304    let now = std::time::SystemTime::now()
305        .duration_since(std::time::UNIX_EPOCH)
306        .map_or(0, |d| d.as_nanos() as u64);
307    uint(&mut out, 9, now);
308    let period_type = value_type(&mut strings, "space", "bytes");
309    bytes(&mut out, 11, &period_type);
310    uint(&mut out, 12, interval as u64);
311    let default = strings.id("inuse_space");
312    uint(&mut out, 14, default);
313
314    for s in &strings.list {
315        bytes(&mut out, 6, s.as_bytes());
316    }
317    out
318}
319
320/// Locations and the functions their lines name, each encoded once
321#[derive(Default)]
322struct Locations {
323    ids: HashMap<usize, u64>,
324    /// By full name and file: simplified names of different functions,
325    /// such as two traits' methods on one type, can be the same
326    functions: HashMap<(String, String), u64>,
327    encoded: Vec<u8>,
328    functions_encoded: Vec<u8>,
329}
330
331impl Locations {
332    /// The id of the location at `addr`, encoding it on first use
333    fn id(&mut self, addr: usize, lines: &[Line], mapping: u64, strings: &mut Strings) -> u64 {
334        if let Some(&id) = self.ids.get(&addr) {
335            return id;
336        }
337        let id = self.ids.len() as u64 + 1;
338        self.ids.insert(addr, id);
339        let mut loc = Vec::new();
340        uint(&mut loc, 1, id);
341        uint(&mut loc, 2, mapping);
342        uint(&mut loc, 3, addr as u64);
343        for line in lines {
344            let mut l = Vec::new();
345            uint(&mut l, 1, self.function(line, strings));
346            uint(&mut l, 2, line.number);
347            bytes(&mut loc, 4, &l);
348        }
349        bytes(&mut self.encoded, 4, &loc);
350        id
351    }
352
353    fn function(&mut self, line: &Line, strings: &mut Strings) -> u64 {
354        let key = (line.system_name.clone(), line.file.clone());
355        if let Some(&id) = self.functions.get(&key) {
356            return id;
357        }
358        let id = self.functions.len() as u64 + 1;
359        let mut f = Vec::new();
360        uint(&mut f, 1, id);
361        uint(&mut f, 2, strings.id(&line.name));
362        uint(&mut f, 3, strings.id(&line.system_name));
363        uint(&mut f, 4, strings.id(&line.file));
364        bytes(&mut self.functions_encoded, 5, &f);
365        self.functions.insert(key, id);
366        id
367    }
368}
369
370/// A mapping; `symbolized` says whether its locations carry their
371/// functions and lines already
372/// `complete`: whether every location in the mapping has functions, file
373/// names and line numbers
374fn encode_mapping(strings: &mut Strings, id: u64, m: &Mapping, complete: [bool; 3]) -> Vec<u8> {
375    let [functions, filenames, lines] = complete;
376    let mut b = Vec::new();
377    uint(&mut b, 1, id);
378    uint(&mut b, 2, m.start);
379    uint(&mut b, 3, m.limit);
380    uint(&mut b, 4, m.offset);
381    uint(&mut b, 5, strings.id(&m.path));
382    uint(&mut b, 7, u64::from(functions));
383    uint(&mut b, 8, u64::from(filenames));
384    uint(&mut b, 9, u64::from(lines));
385    uint(&mut b, 10, u64::from(filenames));
386    let mut out = Vec::new();
387    bytes(&mut out, 3, &b);
388    out
389}
390
391#[cfg(all(test, feature = "symbolize"))]
392mod tests {
393    use super::{Line, Strings, encode_mapping, is_internal, simplify};
394    use crate::prof::maps::Mapping;
395
396    /// Functions found say nothing of file names or lines: a library
397    /// with a symbol table and no debug information has only functions
398    #[test]
399    fn mapping_flags_are_separate() {
400        let m = Mapping {
401            start: 0x1000,
402            limit: 0x2000,
403            offset: 0,
404            path: "/lib/x.so".into(),
405        };
406        let b = encode_mapping(&mut Strings::default(), 1, &m, [true, false, true]);
407        // Fields 7 to 10 are varints; a false one is left out
408        let has = |field: u8| b.windows(2).any(|w| w == [field << 3, 1]);
409        assert!(has(7), "has_functions");
410        assert!(!has(8), "has_filenames");
411        assert!(has(9), "has_line_numbers");
412        assert!(!has(10), "has_inline_frames");
413    }
414
415    fn line(name: &str, file: &str) -> Line {
416        Line {
417            name: name.to_string(),
418            system_name: name.to_string(),
419            file: file.to_string(),
420            number: 1,
421        }
422    }
423
424    /// With line tables only, names come without paths: the allocator's own
425    /// frames are still told by their files
426    #[test]
427    fn internal_frames_without_paths() {
428        let own = format!("{}/src/malloc.rs", env!("CARGO_MANIFEST_DIR"));
429        assert!(is_internal(&line("alloc", &own)));
430        assert!(is_internal(&line("stoolap_jemalloc::malloc::alloc", "")));
431        assert!(is_internal(&line("__rust_alloc", "/app/src/lib.rs")));
432        assert!(!is_internal(&line(
433            "alloc",
434            "/rustlib/src/rust/library/alloc/src/alloc.rs"
435        )));
436        assert!(!is_internal(&line(
437            "add_versions_batch",
438            "/app/src/storage/mvcc.rs"
439        )));
440    }
441
442    #[test]
443    fn simplifies_rust_names() {
444        let cases = [
445            ("alloc::raw_vec::finish_grow", "alloc::raw_vec::finish_grow"),
446            (
447                "<alloc::raw_vec::RawVec<T,A>>::finish_grow",
448                "alloc::raw_vec::RawVec::finish_grow",
449            ),
450            (
451                "<core::result::Result<T, E> as core::ops::function::FnOnce<A>>::call_once",
452                "core::result::Result::call_once",
453            ),
454            ("<&T as core::fmt::Display>::fmt", "&T::fmt"),
455            (
456                "std::thread::Builder::spawn::{{closure}}",
457                "std::thread::Builder::spawn::{{closure}}",
458            ),
459            ("foo::<impl Fn() -> u8>::bar", "foo::impl Fn() -> u8::bar"),
460            (
461                "RawVecInner::reserve::do_reserve_and_handle::<alloc::alloc::Global>",
462                "RawVecInner::reserve::do_reserve_and_handle",
463            ),
464            (
465                "core::ptr::drop_in_place::<alloc::vec::Vec<u8>>",
466                "core::ptr::drop_in_place",
467            ),
468            // An inlined closure's name, as debug information has it
469            (
470                "{closure#0}<i32, alloc::boxed::Box<[u8; 64], alloc::alloc::Global>, ()>",
471                "{closure#0}",
472            ),
473        ];
474        for (name, want) in cases {
475            assert_eq!(simplify(name), want, "{name}");
476        }
477    }
478}