15287f926SAndreas Gerstmayr# flamegraph.py - create flame graphs from perf samples
25287f926SAndreas Gerstmayr# SPDX-License-Identifier: GPL-2.0
35287f926SAndreas Gerstmayr#
45287f926SAndreas Gerstmayr# Usage:
55287f926SAndreas Gerstmayr#
65287f926SAndreas Gerstmayr#     perf record -a -g -F 99 sleep 60
75287f926SAndreas Gerstmayr#     perf script report flamegraph
85287f926SAndreas Gerstmayr#
95287f926SAndreas Gerstmayr# Combined:
105287f926SAndreas Gerstmayr#
115287f926SAndreas Gerstmayr#     perf script flamegraph -a -F 99 sleep 60
125287f926SAndreas Gerstmayr#
135287f926SAndreas Gerstmayr# Written by Andreas Gerstmayr <[email protected]>
145287f926SAndreas Gerstmayr# Flame Graphs invented by Brendan Gregg <[email protected]>
155287f926SAndreas Gerstmayr# Works in tandem with d3-flame-graph by Martin Spier <[email protected]>
165287f926SAndreas Gerstmayr
175287f926SAndreas Gerstmayrfrom __future__ import print_function
185287f926SAndreas Gerstmayrimport sys
195287f926SAndreas Gerstmayrimport os
20*c42ad5d4SAndreas Gerstmayrimport io
215287f926SAndreas Gerstmayrimport argparse
225287f926SAndreas Gerstmayrimport json
235287f926SAndreas Gerstmayr
245287f926SAndreas Gerstmayr
255287f926SAndreas Gerstmayrclass Node:
265287f926SAndreas Gerstmayr    def __init__(self, name, libtype=""):
275287f926SAndreas Gerstmayr        self.name = name
285287f926SAndreas Gerstmayr        self.libtype = libtype
295287f926SAndreas Gerstmayr        self.value = 0
305287f926SAndreas Gerstmayr        self.children = []
315287f926SAndreas Gerstmayr
325287f926SAndreas Gerstmayr    def toJSON(self):
335287f926SAndreas Gerstmayr        return {
345287f926SAndreas Gerstmayr            "n": self.name,
355287f926SAndreas Gerstmayr            "l": self.libtype,
365287f926SAndreas Gerstmayr            "v": self.value,
375287f926SAndreas Gerstmayr            "c": self.children
385287f926SAndreas Gerstmayr        }
395287f926SAndreas Gerstmayr
405287f926SAndreas Gerstmayr
415287f926SAndreas Gerstmayrclass FlameGraphCLI:
425287f926SAndreas Gerstmayr    def __init__(self, args):
435287f926SAndreas Gerstmayr        self.args = args
445287f926SAndreas Gerstmayr        self.stack = Node("root")
455287f926SAndreas Gerstmayr
465287f926SAndreas Gerstmayr        if self.args.format == "html" and \
475287f926SAndreas Gerstmayr                not os.path.isfile(self.args.template):
485287f926SAndreas Gerstmayr            print("Flame Graph template {} does not exist. Please install "
495287f926SAndreas Gerstmayr                  "the js-d3-flame-graph (RPM) or libjs-d3-flame-graph (deb) "
505287f926SAndreas Gerstmayr                  "package, specify an existing flame graph template "
515287f926SAndreas Gerstmayr                  "(--template PATH) or another output format "
525287f926SAndreas Gerstmayr                  "(--format FORMAT).".format(self.args.template),
535287f926SAndreas Gerstmayr                  file=sys.stderr)
545287f926SAndreas Gerstmayr            sys.exit(1)
555287f926SAndreas Gerstmayr
565287f926SAndreas Gerstmayr    def find_or_create_node(self, node, name, dso):
575287f926SAndreas Gerstmayr        libtype = "kernel" if dso == "[kernel.kallsyms]" else ""
585287f926SAndreas Gerstmayr        if name is None:
595287f926SAndreas Gerstmayr            name = "[unknown]"
605287f926SAndreas Gerstmayr
615287f926SAndreas Gerstmayr        for child in node.children:
625287f926SAndreas Gerstmayr            if child.name == name and child.libtype == libtype:
635287f926SAndreas Gerstmayr                return child
645287f926SAndreas Gerstmayr
655287f926SAndreas Gerstmayr        child = Node(name, libtype)
665287f926SAndreas Gerstmayr        node.children.append(child)
675287f926SAndreas Gerstmayr        return child
685287f926SAndreas Gerstmayr
695287f926SAndreas Gerstmayr    def process_event(self, event):
705287f926SAndreas Gerstmayr        node = self.find_or_create_node(self.stack, event["comm"], None)
715287f926SAndreas Gerstmayr        if "callchain" in event:
725287f926SAndreas Gerstmayr            for entry in reversed(event['callchain']):
735287f926SAndreas Gerstmayr                node = self.find_or_create_node(
745287f926SAndreas Gerstmayr                    node, entry.get("sym", {}).get("name"), event.get("dso"))
755287f926SAndreas Gerstmayr        else:
765287f926SAndreas Gerstmayr            node = self.find_or_create_node(
775287f926SAndreas Gerstmayr                node, entry.get("symbol"), event.get("dso"))
785287f926SAndreas Gerstmayr        node.value += 1
795287f926SAndreas Gerstmayr
805287f926SAndreas Gerstmayr    def trace_end(self):
815287f926SAndreas Gerstmayr        json_str = json.dumps(self.stack, default=lambda x: x.toJSON())
825287f926SAndreas Gerstmayr
835287f926SAndreas Gerstmayr        if self.args.format == "html":
845287f926SAndreas Gerstmayr            try:
85*c42ad5d4SAndreas Gerstmayr                with io.open(self.args.template, encoding="utf-8") as f:
865287f926SAndreas Gerstmayr                    output_str = f.read().replace("/** @flamegraph_json **/",
875287f926SAndreas Gerstmayr                                                  json_str)
885287f926SAndreas Gerstmayr            except IOError as e:
895287f926SAndreas Gerstmayr                print("Error reading template file: {}".format(e), file=sys.stderr)
905287f926SAndreas Gerstmayr                sys.exit(1)
915287f926SAndreas Gerstmayr            output_fn = self.args.output or "flamegraph.html"
925287f926SAndreas Gerstmayr        else:
935287f926SAndreas Gerstmayr            output_str = json_str
945287f926SAndreas Gerstmayr            output_fn = self.args.output or "stacks.json"
955287f926SAndreas Gerstmayr
965287f926SAndreas Gerstmayr        if output_fn == "-":
97*c42ad5d4SAndreas Gerstmayr            with io.open(sys.stdout.fileno(), "w", encoding="utf-8", closefd=False) as out:
98*c42ad5d4SAndreas Gerstmayr                out.write(output_str)
995287f926SAndreas Gerstmayr        else:
1005287f926SAndreas Gerstmayr            print("dumping data to {}".format(output_fn))
1015287f926SAndreas Gerstmayr            try:
102*c42ad5d4SAndreas Gerstmayr                with io.open(output_fn, "w", encoding="utf-8") as out:
1035287f926SAndreas Gerstmayr                    out.write(output_str)
1045287f926SAndreas Gerstmayr            except IOError as e:
1055287f926SAndreas Gerstmayr                print("Error writing output file: {}".format(e), file=sys.stderr)
1065287f926SAndreas Gerstmayr                sys.exit(1)
1075287f926SAndreas Gerstmayr
1085287f926SAndreas Gerstmayr
1095287f926SAndreas Gerstmayrif __name__ == "__main__":
1105287f926SAndreas Gerstmayr    parser = argparse.ArgumentParser(description="Create flame graphs.")
1115287f926SAndreas Gerstmayr    parser.add_argument("-f", "--format",
1125287f926SAndreas Gerstmayr                        default="html", choices=["json", "html"],
1135287f926SAndreas Gerstmayr                        help="output file format")
1145287f926SAndreas Gerstmayr    parser.add_argument("-o", "--output",
1155287f926SAndreas Gerstmayr                        help="output file name")
1165287f926SAndreas Gerstmayr    parser.add_argument("--template",
1175287f926SAndreas Gerstmayr                        default="/usr/share/d3-flame-graph/d3-flamegraph-base.html",
1185287f926SAndreas Gerstmayr                        help="path to flamegraph HTML template")
1195287f926SAndreas Gerstmayr    parser.add_argument("-i", "--input",
1205287f926SAndreas Gerstmayr                        help=argparse.SUPPRESS)
1215287f926SAndreas Gerstmayr
1225287f926SAndreas Gerstmayr    args = parser.parse_args()
1235287f926SAndreas Gerstmayr    cli = FlameGraphCLI(args)
1245287f926SAndreas Gerstmayr
1255287f926SAndreas Gerstmayr    process_event = cli.process_event
1265287f926SAndreas Gerstmayr    trace_end = cli.trace_end
127