1use super::StackCounts;
19use super::maps::{self, Mapping};
20use std::collections::HashMap;
21use std::fmt;
22
23#[derive(Debug)]
25pub enum DumpError {
26 NotActivated,
28 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
43const 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
102fn 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
110struct Line {
113 name: String,
114 system_name: String,
115 file: String,
116 number: u64,
117}
118
119#[cfg(feature = "symbolize")]
120fn 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 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 '>' 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
178fn 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
190fn 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 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 let mut symbolized = vec![cfg!(feature = "symbolize"); mappings.len() + 1];
259
260 for s in stacks {
261 if s.alloc_objects == 0 {
262 continue;
263 }
264 let mut ids = Vec::with_capacity(s.frames.len());
265 for &addr in s.frames {
266 let Some(lines) = symbols.entry(addr).or_insert_with(|| user_lines(addr)) else {
267 continue;
269 };
270 let mapping = mapping_of(addr as u64);
271 if lines.is_empty() {
272 symbolized[mapping as usize] = false;
273 }
274 ids.push(locations.id(addr, lines, mapping, &mut strings));
275 }
276 let mut sample = Vec::new();
277 packed(&mut sample, 1, ids);
278 packed(
279 &mut sample,
280 2,
281 [
282 s.alloc_objects,
283 s.alloc_bytes,
284 s.inuse_objects,
285 s.inuse_bytes,
286 ],
287 );
288 bytes(&mut out, 2, &sample);
289 }
290
291 for (i, m) in mappings.iter().enumerate() {
292 let complete = u64::from(symbolized[i + 1]);
293 out.extend(encode_mapping(&mut strings, i as u64 + 1, m, complete));
294 }
295 out.extend(locations.encoded);
296 out.extend(locations.functions_encoded);
297
298 let now = std::time::SystemTime::now()
299 .duration_since(std::time::UNIX_EPOCH)
300 .map_or(0, |d| d.as_nanos() as u64);
301 uint(&mut out, 9, now);
302 let period_type = value_type(&mut strings, "space", "bytes");
303 bytes(&mut out, 11, &period_type);
304 uint(&mut out, 12, interval as u64);
305 let default = strings.id("inuse_space");
306 uint(&mut out, 14, default);
307
308 for s in &strings.list {
309 bytes(&mut out, 6, s.as_bytes());
310 }
311 out
312}
313
314#[derive(Default)]
316struct Locations {
317 ids: HashMap<usize, u64>,
318 functions: HashMap<(String, String), u64>,
321 encoded: Vec<u8>,
322 functions_encoded: Vec<u8>,
323}
324
325impl Locations {
326 fn id(&mut self, addr: usize, lines: &[Line], mapping: u64, strings: &mut Strings) -> u64 {
328 if let Some(&id) = self.ids.get(&addr) {
329 return id;
330 }
331 let id = self.ids.len() as u64 + 1;
332 self.ids.insert(addr, id);
333 let mut loc = Vec::new();
334 uint(&mut loc, 1, id);
335 uint(&mut loc, 2, mapping);
336 uint(&mut loc, 3, addr as u64);
337 for line in lines {
338 let mut l = Vec::new();
339 uint(&mut l, 1, self.function(line, strings));
340 uint(&mut l, 2, line.number);
341 bytes(&mut loc, 4, &l);
342 }
343 bytes(&mut self.encoded, 4, &loc);
344 id
345 }
346
347 fn function(&mut self, line: &Line, strings: &mut Strings) -> u64 {
348 let key = (line.system_name.clone(), line.file.clone());
349 if let Some(&id) = self.functions.get(&key) {
350 return id;
351 }
352 let id = self.functions.len() as u64 + 1;
353 let mut f = Vec::new();
354 uint(&mut f, 1, id);
355 uint(&mut f, 2, strings.id(&line.name));
356 uint(&mut f, 3, strings.id(&line.system_name));
357 uint(&mut f, 4, strings.id(&line.file));
358 bytes(&mut self.functions_encoded, 5, &f);
359 self.functions.insert(key, id);
360 id
361 }
362}
363
364fn encode_mapping(strings: &mut Strings, id: u64, m: &Mapping, symbolized: u64) -> Vec<u8> {
367 let mut b = Vec::new();
368 uint(&mut b, 1, id);
369 uint(&mut b, 2, m.start);
370 uint(&mut b, 3, m.limit);
371 uint(&mut b, 4, m.offset);
372 uint(&mut b, 5, strings.id(&m.path));
373 for field in 7..=10 {
374 uint(&mut b, field, symbolized);
375 }
376 let mut out = Vec::new();
377 bytes(&mut out, 3, &b);
378 out
379}
380
381#[cfg(all(test, feature = "symbolize"))]
382mod tests {
383 use super::{Line, is_internal, simplify};
384
385 fn line(name: &str, file: &str) -> Line {
386 Line {
387 name: name.to_string(),
388 system_name: name.to_string(),
389 file: file.to_string(),
390 number: 1,
391 }
392 }
393
394 #[test]
397 fn internal_frames_without_paths() {
398 let own = format!("{}/src/malloc.rs", env!("CARGO_MANIFEST_DIR"));
399 assert!(is_internal(&line("alloc", &own)));
400 assert!(is_internal(&line("stoolap_jemalloc::malloc::alloc", "")));
401 assert!(is_internal(&line("__rust_alloc", "/app/src/lib.rs")));
402 assert!(!is_internal(&line(
403 "alloc",
404 "/rustlib/src/rust/library/alloc/src/alloc.rs"
405 )));
406 assert!(!is_internal(&line(
407 "add_versions_batch",
408 "/app/src/storage/mvcc.rs"
409 )));
410 }
411
412 #[test]
413 fn simplifies_rust_names() {
414 let cases = [
415 ("alloc::raw_vec::finish_grow", "alloc::raw_vec::finish_grow"),
416 (
417 "<alloc::raw_vec::RawVec<T,A>>::finish_grow",
418 "alloc::raw_vec::RawVec::finish_grow",
419 ),
420 (
421 "<core::result::Result<T, E> as core::ops::function::FnOnce<A>>::call_once",
422 "core::result::Result::call_once",
423 ),
424 ("<&T as core::fmt::Display>::fmt", "&T::fmt"),
425 (
426 "std::thread::Builder::spawn::{{closure}}",
427 "std::thread::Builder::spawn::{{closure}}",
428 ),
429 ("foo::<impl Fn() -> u8>::bar", "foo::impl Fn() -> u8::bar"),
430 (
431 "RawVecInner::reserve::do_reserve_and_handle::<alloc::alloc::Global>",
432 "RawVecInner::reserve::do_reserve_and_handle",
433 ),
434 (
435 "core::ptr::drop_in_place::<alloc::vec::Vec<u8>>",
436 "core::ptr::drop_in_place",
437 ),
438 (
440 "{closure#0}<i32, alloc::boxed::Box<[u8; 64], alloc::alloc::Global>, ()>",
441 "{closure#0}",
442 ),
443 ];
444 for (name, want) in cases {
445 assert_eq!(simplify(name), want, "{name}");
446 }
447 }
448}