162306a36Sopenharmony_ci# flamegraph.py - create flame graphs from perf samples
262306a36Sopenharmony_ci# SPDX-License-Identifier: GPL-2.0
362306a36Sopenharmony_ci#
462306a36Sopenharmony_ci# Usage:
562306a36Sopenharmony_ci#
662306a36Sopenharmony_ci#     perf record -a -g -F 99 sleep 60
762306a36Sopenharmony_ci#     perf script report flamegraph
862306a36Sopenharmony_ci#
962306a36Sopenharmony_ci# Combined:
1062306a36Sopenharmony_ci#
1162306a36Sopenharmony_ci#     perf script flamegraph -a -F 99 sleep 60
1262306a36Sopenharmony_ci#
1362306a36Sopenharmony_ci# Written by Andreas Gerstmayr <agerstmayr@redhat.com>
1462306a36Sopenharmony_ci# Flame Graphs invented by Brendan Gregg <bgregg@netflix.com>
1562306a36Sopenharmony_ci# Works in tandem with d3-flame-graph by Martin Spier <mspier@netflix.com>
1662306a36Sopenharmony_ci#
1762306a36Sopenharmony_ci# pylint: disable=missing-module-docstring
1862306a36Sopenharmony_ci# pylint: disable=missing-class-docstring
1962306a36Sopenharmony_ci# pylint: disable=missing-function-docstring
2062306a36Sopenharmony_ci
2162306a36Sopenharmony_cifrom __future__ import print_function
2262306a36Sopenharmony_ciimport argparse
2362306a36Sopenharmony_ciimport hashlib
2462306a36Sopenharmony_ciimport io
2562306a36Sopenharmony_ciimport json
2662306a36Sopenharmony_ciimport os
2762306a36Sopenharmony_ciimport subprocess
2862306a36Sopenharmony_ciimport sys
2962306a36Sopenharmony_ciimport urllib.request
3062306a36Sopenharmony_ci
3162306a36Sopenharmony_ciminimal_html = """<head>
3262306a36Sopenharmony_ci  <link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/d3-flamegraph.css">
3362306a36Sopenharmony_ci</head>
3462306a36Sopenharmony_ci<body>
3562306a36Sopenharmony_ci  <div id="chart"></div>
3662306a36Sopenharmony_ci  <script type="text/javascript" src="https://d3js.org/d3.v7.js"></script>
3762306a36Sopenharmony_ci  <script type="text/javascript" src="https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/d3-flamegraph.min.js"></script>
3862306a36Sopenharmony_ci  <script type="text/javascript">
3962306a36Sopenharmony_ci  const stacks = [/** @flamegraph_json **/];
4062306a36Sopenharmony_ci  // Note, options is unused.
4162306a36Sopenharmony_ci  const options = [/** @options_json **/];
4262306a36Sopenharmony_ci
4362306a36Sopenharmony_ci  var chart = flamegraph();
4462306a36Sopenharmony_ci  d3.select("#chart")
4562306a36Sopenharmony_ci        .datum(stacks[0])
4662306a36Sopenharmony_ci        .call(chart);
4762306a36Sopenharmony_ci  </script>
4862306a36Sopenharmony_ci</body>
4962306a36Sopenharmony_ci"""
5062306a36Sopenharmony_ci
5162306a36Sopenharmony_ci# pylint: disable=too-few-public-methods
5262306a36Sopenharmony_ciclass Node:
5362306a36Sopenharmony_ci    def __init__(self, name, libtype):
5462306a36Sopenharmony_ci        self.name = name
5562306a36Sopenharmony_ci        # "root" | "kernel" | ""
5662306a36Sopenharmony_ci        # "" indicates user space
5762306a36Sopenharmony_ci        self.libtype = libtype
5862306a36Sopenharmony_ci        self.value = 0
5962306a36Sopenharmony_ci        self.children = []
6062306a36Sopenharmony_ci
6162306a36Sopenharmony_ci    def to_json(self):
6262306a36Sopenharmony_ci        return {
6362306a36Sopenharmony_ci            "n": self.name,
6462306a36Sopenharmony_ci            "l": self.libtype,
6562306a36Sopenharmony_ci            "v": self.value,
6662306a36Sopenharmony_ci            "c": self.children
6762306a36Sopenharmony_ci        }
6862306a36Sopenharmony_ci
6962306a36Sopenharmony_ci
7062306a36Sopenharmony_ciclass FlameGraphCLI:
7162306a36Sopenharmony_ci    def __init__(self, args):
7262306a36Sopenharmony_ci        self.args = args
7362306a36Sopenharmony_ci        self.stack = Node("all", "root")
7462306a36Sopenharmony_ci
7562306a36Sopenharmony_ci    @staticmethod
7662306a36Sopenharmony_ci    def get_libtype_from_dso(dso):
7762306a36Sopenharmony_ci        """
7862306a36Sopenharmony_ci        when kernel-debuginfo is installed,
7962306a36Sopenharmony_ci        dso points to /usr/lib/debug/lib/modules/*/vmlinux
8062306a36Sopenharmony_ci        """
8162306a36Sopenharmony_ci        if dso and (dso == "[kernel.kallsyms]" or dso.endswith("/vmlinux")):
8262306a36Sopenharmony_ci            return "kernel"
8362306a36Sopenharmony_ci
8462306a36Sopenharmony_ci        return ""
8562306a36Sopenharmony_ci
8662306a36Sopenharmony_ci    @staticmethod
8762306a36Sopenharmony_ci    def find_or_create_node(node, name, libtype):
8862306a36Sopenharmony_ci        for child in node.children:
8962306a36Sopenharmony_ci            if child.name == name:
9062306a36Sopenharmony_ci                return child
9162306a36Sopenharmony_ci
9262306a36Sopenharmony_ci        child = Node(name, libtype)
9362306a36Sopenharmony_ci        node.children.append(child)
9462306a36Sopenharmony_ci        return child
9562306a36Sopenharmony_ci
9662306a36Sopenharmony_ci    def process_event(self, event):
9762306a36Sopenharmony_ci        pid = event.get("sample", {}).get("pid", 0)
9862306a36Sopenharmony_ci        # event["dso"] sometimes contains /usr/lib/debug/lib/modules/*/vmlinux
9962306a36Sopenharmony_ci        # for user-space processes; let's use pid for kernel or user-space distinction
10062306a36Sopenharmony_ci        if pid == 0:
10162306a36Sopenharmony_ci            comm = event["comm"]
10262306a36Sopenharmony_ci            libtype = "kernel"
10362306a36Sopenharmony_ci        else:
10462306a36Sopenharmony_ci            comm = "{} ({})".format(event["comm"], pid)
10562306a36Sopenharmony_ci            libtype = ""
10662306a36Sopenharmony_ci        node = self.find_or_create_node(self.stack, comm, libtype)
10762306a36Sopenharmony_ci
10862306a36Sopenharmony_ci        if "callchain" in event:
10962306a36Sopenharmony_ci            for entry in reversed(event["callchain"]):
11062306a36Sopenharmony_ci                name = entry.get("sym", {}).get("name", "[unknown]")
11162306a36Sopenharmony_ci                libtype = self.get_libtype_from_dso(entry.get("dso"))
11262306a36Sopenharmony_ci                node = self.find_or_create_node(node, name, libtype)
11362306a36Sopenharmony_ci        else:
11462306a36Sopenharmony_ci            name = event.get("symbol", "[unknown]")
11562306a36Sopenharmony_ci            libtype = self.get_libtype_from_dso(event.get("dso"))
11662306a36Sopenharmony_ci            node = self.find_or_create_node(node, name, libtype)
11762306a36Sopenharmony_ci        node.value += 1
11862306a36Sopenharmony_ci
11962306a36Sopenharmony_ci    def get_report_header(self):
12062306a36Sopenharmony_ci        if self.args.input == "-":
12162306a36Sopenharmony_ci            # when this script is invoked with "perf script flamegraph",
12262306a36Sopenharmony_ci            # no perf.data is created and we cannot read the header of it
12362306a36Sopenharmony_ci            return ""
12462306a36Sopenharmony_ci
12562306a36Sopenharmony_ci        try:
12662306a36Sopenharmony_ci            output = subprocess.check_output(["perf", "report", "--header-only"])
12762306a36Sopenharmony_ci            return output.decode("utf-8")
12862306a36Sopenharmony_ci        except Exception as err:  # pylint: disable=broad-except
12962306a36Sopenharmony_ci            print("Error reading report header: {}".format(err), file=sys.stderr)
13062306a36Sopenharmony_ci            return ""
13162306a36Sopenharmony_ci
13262306a36Sopenharmony_ci    def trace_end(self):
13362306a36Sopenharmony_ci        stacks_json = json.dumps(self.stack, default=lambda x: x.to_json())
13462306a36Sopenharmony_ci
13562306a36Sopenharmony_ci        if self.args.format == "html":
13662306a36Sopenharmony_ci            report_header = self.get_report_header()
13762306a36Sopenharmony_ci            options = {
13862306a36Sopenharmony_ci                "colorscheme": self.args.colorscheme,
13962306a36Sopenharmony_ci                "context": report_header
14062306a36Sopenharmony_ci            }
14162306a36Sopenharmony_ci            options_json = json.dumps(options)
14262306a36Sopenharmony_ci
14362306a36Sopenharmony_ci            template_md5sum = None
14462306a36Sopenharmony_ci            if self.args.format == "html":
14562306a36Sopenharmony_ci                if os.path.isfile(self.args.template):
14662306a36Sopenharmony_ci                    template = f"file://{self.args.template}"
14762306a36Sopenharmony_ci                else:
14862306a36Sopenharmony_ci                    if not self.args.allow_download:
14962306a36Sopenharmony_ci                        print(f"""Warning: Flame Graph template '{self.args.template}'
15062306a36Sopenharmony_cidoes not exist. To avoid this please install a package such as the
15162306a36Sopenharmony_cijs-d3-flame-graph or libjs-d3-flame-graph, specify an existing flame
15262306a36Sopenharmony_cigraph template (--template PATH) or use another output format (--format
15362306a36Sopenharmony_ciFORMAT).""",
15462306a36Sopenharmony_ci                              file=sys.stderr)
15562306a36Sopenharmony_ci                        if self.args.input == "-":
15662306a36Sopenharmony_ci                            print("""Not attempting to download Flame Graph template as script command line
15762306a36Sopenharmony_ciinput is disabled due to using live mode. If you want to download the
15862306a36Sopenharmony_citemplate retry without live mode. For example, use 'perf record -a -g
15962306a36Sopenharmony_ci-F 99 sleep 60' and 'perf script report flamegraph'. Alternatively,
16062306a36Sopenharmony_cidownload the template from:
16162306a36Sopenharmony_cihttps://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/templates/d3-flamegraph-base.html
16262306a36Sopenharmony_ciand place it at:
16362306a36Sopenharmony_ci/usr/share/d3-flame-graph/d3-flamegraph-base.html""",
16462306a36Sopenharmony_ci                                  file=sys.stderr)
16562306a36Sopenharmony_ci                            quit()
16662306a36Sopenharmony_ci                        s = None
16762306a36Sopenharmony_ci                        while s != "y" and s != "n":
16862306a36Sopenharmony_ci                            s = input("Do you wish to download a template from cdn.jsdelivr.net? (this warning can be suppressed with --allow-download) [yn] ").lower()
16962306a36Sopenharmony_ci                        if s == "n":
17062306a36Sopenharmony_ci                            quit()
17162306a36Sopenharmony_ci                    template = "https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/templates/d3-flamegraph-base.html"
17262306a36Sopenharmony_ci                    template_md5sum = "143e0d06ba69b8370b9848dcd6ae3f36"
17362306a36Sopenharmony_ci
17462306a36Sopenharmony_ci            try:
17562306a36Sopenharmony_ci                with urllib.request.urlopen(template) as template:
17662306a36Sopenharmony_ci                    output_str = "".join([
17762306a36Sopenharmony_ci                        l.decode("utf-8") for l in template.readlines()
17862306a36Sopenharmony_ci                    ])
17962306a36Sopenharmony_ci            except Exception as err:
18062306a36Sopenharmony_ci                print(f"Error reading template {template}: {err}\n"
18162306a36Sopenharmony_ci                      "a minimal flame graph will be generated", file=sys.stderr)
18262306a36Sopenharmony_ci                output_str = minimal_html
18362306a36Sopenharmony_ci                template_md5sum = None
18462306a36Sopenharmony_ci
18562306a36Sopenharmony_ci            if template_md5sum:
18662306a36Sopenharmony_ci                download_md5sum = hashlib.md5(output_str.encode("utf-8")).hexdigest()
18762306a36Sopenharmony_ci                if download_md5sum != template_md5sum:
18862306a36Sopenharmony_ci                    s = None
18962306a36Sopenharmony_ci                    while s != "y" and s != "n":
19062306a36Sopenharmony_ci                        s = input(f"""Unexpected template md5sum.
19162306a36Sopenharmony_ci{download_md5sum} != {template_md5sum}, for:
19262306a36Sopenharmony_ci{output_str}
19362306a36Sopenharmony_cicontinue?[yn] """).lower()
19462306a36Sopenharmony_ci                    if s == "n":
19562306a36Sopenharmony_ci                        quit()
19662306a36Sopenharmony_ci
19762306a36Sopenharmony_ci            output_str = output_str.replace("/** @options_json **/", options_json)
19862306a36Sopenharmony_ci            output_str = output_str.replace("/** @flamegraph_json **/", stacks_json)
19962306a36Sopenharmony_ci
20062306a36Sopenharmony_ci            output_fn = self.args.output or "flamegraph.html"
20162306a36Sopenharmony_ci        else:
20262306a36Sopenharmony_ci            output_str = stacks_json
20362306a36Sopenharmony_ci            output_fn = self.args.output or "stacks.json"
20462306a36Sopenharmony_ci
20562306a36Sopenharmony_ci        if output_fn == "-":
20662306a36Sopenharmony_ci            with io.open(sys.stdout.fileno(), "w", encoding="utf-8", closefd=False) as out:
20762306a36Sopenharmony_ci                out.write(output_str)
20862306a36Sopenharmony_ci        else:
20962306a36Sopenharmony_ci            print("dumping data to {}".format(output_fn))
21062306a36Sopenharmony_ci            try:
21162306a36Sopenharmony_ci                with io.open(output_fn, "w", encoding="utf-8") as out:
21262306a36Sopenharmony_ci                    out.write(output_str)
21362306a36Sopenharmony_ci            except IOError as err:
21462306a36Sopenharmony_ci                print("Error writing output file: {}".format(err), file=sys.stderr)
21562306a36Sopenharmony_ci                sys.exit(1)
21662306a36Sopenharmony_ci
21762306a36Sopenharmony_ci
21862306a36Sopenharmony_ciif __name__ == "__main__":
21962306a36Sopenharmony_ci    parser = argparse.ArgumentParser(description="Create flame graphs.")
22062306a36Sopenharmony_ci    parser.add_argument("-f", "--format",
22162306a36Sopenharmony_ci                        default="html", choices=["json", "html"],
22262306a36Sopenharmony_ci                        help="output file format")
22362306a36Sopenharmony_ci    parser.add_argument("-o", "--output",
22462306a36Sopenharmony_ci                        help="output file name")
22562306a36Sopenharmony_ci    parser.add_argument("--template",
22662306a36Sopenharmony_ci                        default="/usr/share/d3-flame-graph/d3-flamegraph-base.html",
22762306a36Sopenharmony_ci                        help="path to flame graph HTML template")
22862306a36Sopenharmony_ci    parser.add_argument("--colorscheme",
22962306a36Sopenharmony_ci                        default="blue-green",
23062306a36Sopenharmony_ci                        help="flame graph color scheme",
23162306a36Sopenharmony_ci                        choices=["blue-green", "orange"])
23262306a36Sopenharmony_ci    parser.add_argument("-i", "--input",
23362306a36Sopenharmony_ci                        help=argparse.SUPPRESS)
23462306a36Sopenharmony_ci    parser.add_argument("--allow-download",
23562306a36Sopenharmony_ci                        default=False,
23662306a36Sopenharmony_ci                        action="store_true",
23762306a36Sopenharmony_ci                        help="allow unprompted downloading of HTML template")
23862306a36Sopenharmony_ci
23962306a36Sopenharmony_ci    cli_args = parser.parse_args()
24062306a36Sopenharmony_ci    cli = FlameGraphCLI(cli_args)
24162306a36Sopenharmony_ci
24262306a36Sopenharmony_ci    process_event = cli.process_event
24362306a36Sopenharmony_ci    trace_end = cli.trace_end
244