allocative/
flamegraph.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 *
4 * This source code is licensed under both the MIT license found in the
5 * LICENSE-MIT file in the root directory of this source tree and the Apache
6 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
7 * of this source tree.
8 */
9
10use std::collections::hash_map;
11use std::collections::HashMap;
12use std::collections::HashSet;
13use std::fmt::Write as _;
14use std::mem;
15use std::ops::Index;
16use std::ops::IndexMut;
17
18use crate::global_root::roots;
19use crate::key::Key;
20use crate::visitor::NodeKind;
21use crate::visitor::Visitor;
22use crate::visitor::VisitorImpl;
23use crate::Allocative;
24
25/// Node in flamegraph tree.
26///
27/// Can be written to flamegraph format with [`write`](FlameGraph::write).
28#[derive(Debug, Default, Clone)]
29pub struct FlameGraph {
30    children: HashMap<Key, FlameGraph>,
31    /// Total size of all children, cached.
32    children_size: usize,
33    /// Node size excluding children.
34    node_size: usize,
35}
36
37impl FlameGraph {
38    pub fn total_size(&self) -> usize {
39        self.node_size + self.children_size
40    }
41
42    /// Add another flamegraph to this one.
43    pub fn add(&mut self, other: FlameGraph) {
44        self.node_size += other.node_size;
45        for (key, child) in other.children {
46            self.add_child(key, child);
47        }
48    }
49
50    /// Add a child node to the flamegraph, merging if it already exists.
51    pub fn add_child(&mut self, key: Key, child: FlameGraph) {
52        self.children_size += child.total_size();
53        match self.children.entry(key) {
54            hash_map::Entry::Occupied(mut entry) => {
55                entry.get_mut().add(child);
56            }
57            hash_map::Entry::Vacant(entry) => {
58                entry.insert(child);
59            }
60        }
61    }
62
63    /// Add size to this node.
64    pub fn add_self(&mut self, size: usize) {
65        self.node_size += size;
66    }
67
68    fn write_flame_graph_impl(&self, stack: &[&str], w: &mut String) {
69        if self.node_size != 0 {
70            if !stack.is_empty() {
71                writeln!(w, "{} {}", stack.join(";"), self.node_size).unwrap();
72            } else {
73                // Don't care.
74            }
75        }
76        let mut stack = stack.to_vec();
77        let mut children = Vec::from_iter(self.children.iter());
78        children.sort_by_key(|(key, _)| *key);
79        for (key, child) in children {
80            stack.push(key);
81            child.write_flame_graph_impl(&stack, w);
82            stack.pop().unwrap();
83        }
84    }
85
86    /// Write flamegraph in format suitable for [`flamegraph.pl`] or [inferno].
87    ///
88    /// [flamegraph.pl]: https://github.com/brendangregg/FlameGraph
89    /// [inferno]: https://github.com/jonhoo/inferno
90    pub fn write(&self) -> String {
91        let mut r = String::new();
92        self.write_flame_graph_impl(&[], &mut r);
93        r
94    }
95}
96
97#[derive(Debug)]
98pub struct FlameGraphOutput {
99    flamegraph: FlameGraph,
100    warnings: String,
101}
102
103impl FlameGraphOutput {
104    /// Flamegraph source, can be fed to `flamegraph.pl` or `inferno`.
105    pub fn flamegraph(&self) -> &FlameGraph {
106        &self.flamegraph
107    }
108
109    /// Warnings. Text file in unspecified format.
110    pub fn warnings(&self) -> String {
111        self.warnings.clone()
112    }
113}
114
115#[derive(Default, Eq, PartialEq, Clone, Debug)]
116struct TreeData {
117    /// Size of this node including children but excluding unique/shared children.
118    /// For example for `String` this would be `size_of::<String>()`.
119    size: usize,
120    /// Size excluding children. This value is output to flamegraph for given stack.
121    /// Can be negative if nodes provides sizes incorrectly.
122    rem_size: isize,
123    /// Whether this node is `Box` something.
124    unique: bool,
125    /// Child nodes.
126    children: HashMap<Key, TreeId>,
127}
128
129impl TreeData {
130    fn inline_children_size(&self) -> isize {
131        self.size as isize - self.rem_size
132    }
133}
134
135struct TreeRef<'a> {
136    trees: &'a Trees,
137    tree_id: TreeId,
138}
139
140impl<'a> TreeRef<'a> {
141    fn write_flame_graph(&self, stack: &[&str], warnings: &mut String) -> FlameGraph {
142        let mut flame_graph = FlameGraph::default();
143        let tree = &self.trees[self.tree_id];
144        if tree.rem_size > 0 {
145            if stack.is_empty() {
146                // don't care.
147            } else {
148                flame_graph.node_size = tree.rem_size as usize;
149            }
150        } else if tree.rem_size < 0 && !stack.is_empty() {
151            writeln!(
152                warnings,
153                "Incorrect size declaration for node `{}`, size of self: {}, size of inline children: {}",
154                stack.join(";"),
155                tree.size,
156                tree.inline_children_size()
157            )
158                .unwrap();
159        }
160        let mut children: Vec<(&Key, &TreeId)> = Vec::from_iter(&tree.children);
161        let mut stack = stack.to_vec();
162        children.sort_by_key(|(k, _)| *k);
163        for (key, child) in children {
164            stack.push(key);
165            let child = TreeRef {
166                trees: self.trees,
167                tree_id: *child,
168            };
169            let child_framegraph = child.write_flame_graph(&stack, warnings);
170            flame_graph.add_child(key.clone(), child_framegraph);
171            stack.pop().unwrap();
172        }
173        flame_graph
174    }
175
176    fn to_flame_graph(&self) -> (FlameGraph, String) {
177        let mut warnings = String::new();
178        let flame_graph = self.write_flame_graph(&[], &mut warnings);
179        (flame_graph, warnings)
180    }
181}
182
183#[derive(Debug, Eq, PartialEq)]
184struct Tree {
185    trees: Trees,
186    tree_id: TreeId,
187}
188
189impl Tree {
190    fn as_ref(&self) -> TreeRef {
191        TreeRef {
192            trees: &self.trees,
193            tree_id: self.tree_id,
194        }
195    }
196
197    fn to_flame_graph(&self) -> (FlameGraph, String) {
198        self.as_ref().to_flame_graph()
199    }
200}
201
202#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
203struct TreeId(usize);
204
205#[derive(Default, Clone, Debug, Eq, PartialEq)]
206struct Trees {
207    trees: Vec<TreeData>,
208}
209
210impl Trees {
211    fn new_tree(&mut self) -> TreeId {
212        let id = TreeId(self.trees.len());
213        self.trees.push(TreeData::default());
214        id
215    }
216}
217
218impl Index<TreeId> for Trees {
219    type Output = TreeData;
220
221    fn index(&self, index: TreeId) -> &Self::Output {
222        &self.trees[index.0]
223    }
224}
225
226impl IndexMut<TreeId> for Trees {
227    fn index_mut(&mut self, index: TreeId) -> &mut Self::Output {
228        &mut self.trees[index.0]
229    }
230}
231
232#[derive(Clone, Debug)]
233struct TreeStack {
234    stack: Vec<TreeId>,
235    tree: TreeId,
236}
237
238struct TreeStackRef<'t, 's> {
239    trees: &'t mut Trees,
240    stack: &'s mut TreeStack,
241}
242
243impl<'t, 's> TreeStackRef<'t, 's> {
244    fn current_data(&'t mut self) -> &'t mut TreeData {
245        &mut self.trees[self.stack.tree]
246    }
247
248    fn down(&mut self, key: Key) {
249        self.stack.stack.push(self.stack.tree);
250        let next_tree_id = TreeId(self.trees.trees.len());
251        let child = match self.trees[self.stack.tree].children.entry(key) {
252            hash_map::Entry::Occupied(e) => *e.get(),
253            hash_map::Entry::Vacant(e) => {
254                e.insert(next_tree_id);
255                let child = self.trees.new_tree();
256                assert_eq!(child, next_tree_id);
257                child
258            }
259        };
260        self.stack.tree = child;
261    }
262
263    #[must_use]
264    fn up(&mut self) -> bool {
265        if let Some(pop) = self.stack.stack.pop() {
266            self.stack.tree = pop;
267            true
268        } else {
269            false
270        }
271    }
272}
273
274#[derive(Eq, PartialEq, Hash, Clone, Copy, Debug)]
275struct VisitedSharedPointer(*const ());
276unsafe impl Send for VisitedSharedPointer {}
277
278/// Build a flamegraph from given root objects.
279///
280/// # Example
281///
282/// ```
283/// use allocative::Allocative;
284/// use allocative::FlameGraphBuilder;
285///
286/// #[derive(Allocative)]
287/// struct Foo {
288///     data: String,
289/// };
290///
291/// let foo1 = Foo {
292///     data: "Hello, world!".to_owned(),
293/// };
294/// let foo2 = Foo {
295///     data: "Another message!".to_owned(),
296/// };
297///
298/// let mut flamegraph = FlameGraphBuilder::default();
299/// flamegraph.visit_root(&foo1);
300/// flamegraph.visit_root(&foo2);
301/// let flamegraph_src = flamegraph.finish().flamegraph();
302/// ```
303///
304/// And now `flamegraph_src` can be fed to either [flamegraph.pl] or [inferno].
305///
306/// [flamegraph.pl]: https://github.com/brendangregg/FlameGraph
307/// [inferno]: https://github.com/jonhoo/inferno
308#[derive(Debug)]
309pub struct FlameGraphBuilder {
310    /// Visited shared pointers.
311    visited_shared: HashSet<VisitedSharedPointer>,
312    /// Tree data storage.
313    trees: Trees,
314    /// Current node we are processing in `Visitor`.
315    current: TreeStack,
316    /// Previous stack when entering shared pointer.
317    shared: Vec<TreeStack>,
318    /// Data root.
319    root: TreeId,
320    /// Is root visitor created?
321    entered_root_visitor: bool,
322}
323
324fn _assert_flame_graph_builder_is_send() {
325    fn assert_send<T: Send>() {}
326    assert_send::<FlameGraphBuilder>();
327}
328
329impl Default for FlameGraphBuilder {
330    fn default() -> FlameGraphBuilder {
331        let mut trees = Trees::default();
332        let root = trees.new_tree();
333        FlameGraphBuilder {
334            trees,
335            visited_shared: HashSet::new(),
336            current: TreeStack {
337                stack: Vec::new(),
338                tree: root,
339            },
340            shared: Vec::new(),
341            root,
342            entered_root_visitor: false,
343        }
344    }
345}
346
347impl FlameGraphBuilder {
348    pub fn root_visitor(&mut self) -> Visitor {
349        assert!(!self.entered_root_visitor);
350        self.entered_root_visitor = true;
351        Visitor {
352            visitor: self,
353            node_kind: NodeKind::Root,
354        }
355    }
356
357    /// Collect tree sizes starting from given root.
358    pub fn visit_root(&mut self, root: &dyn Allocative) {
359        let mut visitor = self.root_visitor();
360        root.visit(&mut visitor);
361        visitor.exit();
362    }
363
364    /// Collect data from global roots registered with
365    /// [`register_root`](crate::register_root).
366    pub fn visit_global_roots(&mut self) {
367        for root in roots() {
368            self.visit_root(root);
369        }
370    }
371
372    fn finish_impl(mut self) -> Tree {
373        assert!(self.shared.is_empty());
374        assert!(self.current.stack.is_empty());
375        assert!(!self.entered_root_visitor);
376        Self::update_sizes(self.root, &mut self.trees);
377        Tree {
378            trees: self.trees,
379            tree_id: self.root,
380        }
381    }
382
383    /// Finish building the flamegraph.
384    pub fn finish(self) -> FlameGraphOutput {
385        let tree = self.finish_impl();
386        let (flamegraph, warnings) = tree.to_flame_graph();
387        FlameGraphOutput {
388            flamegraph,
389            warnings,
390        }
391    }
392
393    /// Finish building the flamegraph and return the flamegraph output.
394    pub fn finish_and_write_flame_graph(self) -> String {
395        self.finish().flamegraph.write()
396    }
397
398    fn update_sizes(tree_id: TreeId, trees: &mut Trees) {
399        let tree = &mut trees[tree_id];
400        for child in tree.children.values().copied().collect::<Vec<_>>() {
401            Self::update_sizes(child, trees);
402        }
403        let tree = &mut trees[tree_id];
404        let children_size = if tree.unique {
405            0
406        } else {
407            let tree = &trees[tree_id];
408            tree.children
409                .values()
410                .map(|child| trees[*child].size)
411                .sum::<usize>()
412        };
413        let tree = &mut trees[tree_id];
414        let size = tree.size;
415        tree.rem_size = (size as isize).saturating_sub(children_size as isize);
416    }
417
418    fn current(&mut self) -> TreeStackRef {
419        TreeStackRef {
420            trees: &mut self.trees,
421            stack: &mut self.current,
422        }
423    }
424
425    fn exit_impl(&mut self) {
426        assert!(self.entered_root_visitor);
427
428        let up = self.current().up();
429        if !up {
430            if let Some(shared) = self.shared.pop() {
431                self.current = shared;
432                assert!(self.current().up());
433            } else {
434                self.entered_root_visitor = false;
435            }
436        }
437    }
438}
439
440impl VisitorImpl for FlameGraphBuilder {
441    fn enter_inline_impl(&mut self, name: Key, size: usize, _parent: NodeKind) {
442        self.current().down(name);
443        self.current().current_data().size += size;
444    }
445
446    fn enter_unique_impl(&mut self, name: Key, size: usize, _parent: NodeKind) {
447        self.current().down(name);
448        self.current().current_data().size += size;
449        // TODO: deal with potential issue when node is both unique and not.
450        // TODO: record some malloc overhead.
451        self.current().current_data().unique = true;
452    }
453
454    #[must_use]
455    fn enter_shared_impl(
456        &mut self,
457        name: Key,
458        size: usize,
459        ptr: *const (),
460        _parent: NodeKind,
461    ) -> bool {
462        self.current().down(name);
463        self.current().current_data().size += size;
464
465        if !self.visited_shared.insert(VisitedSharedPointer(ptr)) {
466            self.exit_impl();
467            return false;
468        }
469
470        self.shared.push(mem::replace(
471            &mut self.current,
472            TreeStack {
473                stack: Vec::new(),
474                tree: self.root,
475            },
476        ));
477        true
478    }
479
480    fn exit_inline_impl(&mut self) {
481        self.exit_impl();
482    }
483
484    fn exit_unique_impl(&mut self) {
485        self.exit_impl();
486    }
487
488    fn exit_shared_impl(&mut self) {
489        self.exit_impl();
490    }
491
492    fn exit_root_impl(&mut self) {
493        self.exit_impl();
494    }
495}
496
497#[cfg(test)]
498mod tests {
499    use crate::flamegraph::FlameGraphBuilder;
500    use crate::flamegraph::Tree;
501    use crate::flamegraph::Trees;
502    use crate::key::Key;
503    use crate::FlameGraph;
504
505    #[test]
506    fn test_empty() {
507        let mut fg = FlameGraphBuilder::default();
508        fg.root_visitor().exit();
509        let tree = fg.finish_impl();
510
511        let mut expected_trees = Trees::default();
512        let expected_id = expected_trees.new_tree();
513        let expected = Tree {
514            trees: expected_trees,
515            tree_id: expected_id,
516        };
517
518        assert_eq!(expected, tree);
519        assert_eq!("", tree.to_flame_graph().0.write());
520    }
521
522    #[test]
523    fn test_simple() {
524        let mut fg = FlameGraphBuilder::default();
525        fg.root_visitor().visit_simple(Key::new("a"), 10);
526        let tree = fg.finish_impl();
527
528        let mut expected = Trees::default();
529        let expected_root = expected.new_tree();
530        let expected_child = expected.new_tree();
531        expected[expected_root].size = 0;
532        expected[expected_root].rem_size = -10;
533        expected[expected_root]
534            .children
535            .insert(Key::new("a"), expected_child);
536        expected[expected_child].size = 10;
537        expected[expected_child].rem_size = 10;
538        let expected = Tree {
539            trees: expected,
540            tree_id: expected_root,
541        };
542        assert_eq!(expected, tree);
543        assert_eq!("a 10\n", tree.to_flame_graph().0.write());
544    }
545
546    #[test]
547    fn test_unique() {
548        let mut fg = FlameGraphBuilder::default();
549        let mut visitor = fg.root_visitor();
550        let mut s = visitor.enter(Key::new("Struct"), 10);
551        s.visit_simple(Key::new("a"), 3);
552        let mut un = s.enter_unique(Key::new("p"), 6);
553        un.visit_simple(Key::new("x"), 13);
554        un.exit();
555        s.exit();
556        visitor.exit();
557
558        let tree = fg.finish_impl();
559
560        assert_eq!(
561            "\
562                Struct 1\n\
563                Struct;a 3\n\
564                Struct;p 6\n\
565                Struct;p;x 13\n\
566            ",
567            tree.to_flame_graph().0.write(),
568            "{:#?}",
569            tree,
570        );
571    }
572
573    #[test]
574    fn test_shared() {
575        let p = 10;
576
577        let mut fg = FlameGraphBuilder::default();
578        let mut visitor = fg.root_visitor();
579
580        for _ in 0..2 {
581            let mut s = visitor.enter(Key::new("Struct"), 10);
582            s.visit_simple(Key::new("a"), 3);
583            {
584                let sh = s.enter_shared(Key::new("p"), 6, &p as *const i32 as *const ());
585                if let Some(mut sh) = sh {
586                    sh.visit_simple(Key::new("Shared"), 13);
587                    sh.exit();
588                }
589            }
590            s.exit();
591        }
592
593        visitor.exit();
594
595        let tree = fg.finish_impl();
596
597        assert_eq!(
598            "\
599            Shared 13\n\
600            Struct 2\n\
601            Struct;a 6\n\
602            Struct;p 12\n\
603        ",
604            tree.to_flame_graph().0.write(),
605            "{:#?}",
606            tree,
607        );
608    }
609
610    #[test]
611    fn test_inline_children_too_large() {
612        let mut fg = FlameGraphBuilder::default();
613        let mut visitor = fg.root_visitor();
614        let mut child_visitor = visitor.enter(Key::new("a"), 10);
615        child_visitor.visit_simple(Key::new("b"), 13);
616        child_visitor.exit();
617        visitor.exit();
618        let output = fg.finish();
619        assert_eq!("a;b 13\n", output.flamegraph().write());
620        assert_eq!(
621            "Incorrect size declaration for node `a`, size of self: 10, size of inline children: 13\n",
622            output.warnings()
623        );
624    }
625
626    #[test]
627    fn test_flamegraph_add() {
628        let mut a = FlameGraph::default();
629
630        let mut b_1 = FlameGraph::default();
631        b_1.add_self(10);
632
633        let mut b = FlameGraph::default();
634        b.add_child(Key::new("x"), b_1);
635
636        a.add(b);
637        assert_eq!(10, a.total_size());
638    }
639}
allocative/flamegraph.rs

allocative/
flamegraph.rs
