1# flamegraph.py - create flame graphs from perf samples
2# SPDX-License-Identifier: GPL-2.0
3#
4# Usage:
5#
6#     perf record -a -g -F 99 sleep 60
7#     perf script report flamegraph
8#
9# Combined:
10#
11#     perf script flamegraph -a -F 99 sleep 60
12#
13# Written by Andreas Gerstmayr <agerstmayr@redhat.com>
14# Flame Graphs invented by Brendan Gregg <bgregg@netflix.com>
15# Works in tandem with d3-flame-graph by Martin Spier <mspier@netflix.com>
16#
17# pylint: disable=missing-module-docstring
18# pylint: disable=missing-class-docstring
19# pylint: disable=missing-function-docstring
20
21from __future__ import print_function
22import argparse
23import hashlib
24import io
25import json
26import os
27import subprocess
28import sys
29import urllib.request
30
31minimal_html = """<head>
32  <link rel="stylesheet" type="text/css" href="https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/d3-flamegraph.css">
33</head>
34<body>
35  <div id="chart"></div>
36  <script type="text/javascript" src="https://d3js.org/d3.v7.js"></script>
37  <script type="text/javascript" src="https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/d3-flamegraph.min.js"></script>
38  <script type="text/javascript">
39  const stacks = [/** @flamegraph_json **/];
40  // Note, options is unused.
41  const options = [/** @options_json **/];
42
43  var chart = flamegraph();
44  d3.select("#chart")
45        .datum(stacks[0])
46        .call(chart);
47  </script>
48</body>
49"""
50
51# pylint: disable=too-few-public-methods
52class Node:
53    def __init__(self, name, libtype):
54        self.name = name
55        # "root" | "kernel" | ""
56        # "" indicates user space
57        self.libtype = libtype
58        self.value = 0
59        self.children = []
60
61    def to_json(self):
62        return {
63            "n": self.name,
64            "l": self.libtype,
65            "v": self.value,
66            "c": self.children
67        }
68
69
70class FlameGraphCLI:
71    def __init__(self, args):
72        self.args = args
73        self.stack = Node("all", "root")
74
75    @staticmethod
76    def get_libtype_from_dso(dso):
77        """
78        when kernel-debuginfo is installed,
79        dso points to /usr/lib/debug/lib/modules/*/vmlinux
80        """
81        if dso and (dso == "[kernel.kallsyms]" or dso.endswith("/vmlinux")):
82            return "kernel"
83
84        return ""
85
86    @staticmethod
87    def find_or_create_node(node, name, libtype):
88        for child in node.children:
89            if child.name == name:
90                return child
91
92        child = Node(name, libtype)
93        node.children.append(child)
94        return child
95
96    def process_event(self, event):
97        pid = event.get("sample", {}).get("pid", 0)
98        # event["dso"] sometimes contains /usr/lib/debug/lib/modules/*/vmlinux
99        # for user-space processes; let's use pid for kernel or user-space distinction
100        if pid == 0:
101            comm = event["comm"]
102            libtype = "kernel"
103        else:
104            comm = "{} ({})".format(event["comm"], pid)
105            libtype = ""
106        node = self.find_or_create_node(self.stack, comm, libtype)
107
108        if "callchain" in event:
109            for entry in reversed(event["callchain"]):
110                name = entry.get("sym", {}).get("name", "[unknown]")
111                libtype = self.get_libtype_from_dso(entry.get("dso"))
112                node = self.find_or_create_node(node, name, libtype)
113        else:
114            name = event.get("symbol", "[unknown]")
115            libtype = self.get_libtype_from_dso(event.get("dso"))
116            node = self.find_or_create_node(node, name, libtype)
117        node.value += 1
118
119    def get_report_header(self):
120        if self.args.input == "-":
121            # when this script is invoked with "perf script flamegraph",
122            # no perf.data is created and we cannot read the header of it
123            return ""
124
125        try:
126            output = subprocess.check_output(["perf", "report", "--header-only"])
127            return output.decode("utf-8")
128        except Exception as err:  # pylint: disable=broad-except
129            print("Error reading report header: {}".format(err), file=sys.stderr)
130            return ""
131
132    def trace_end(self):
133        stacks_json = json.dumps(self.stack, default=lambda x: x.to_json())
134
135        if self.args.format == "html":
136            report_header = self.get_report_header()
137            options = {
138                "colorscheme": self.args.colorscheme,
139                "context": report_header
140            }
141            options_json = json.dumps(options)
142
143            template_md5sum = None
144            if self.args.format == "html":
145                if os.path.isfile(self.args.template):
146                    template = f"file://{self.args.template}"
147                else:
148                    if not self.args.allow_download:
149                        print(f"""Warning: Flame Graph template '{self.args.template}'
150does not exist. To avoid this please install a package such as the
151js-d3-flame-graph or libjs-d3-flame-graph, specify an existing flame
152graph template (--template PATH) or use another output format (--format
153FORMAT).""",
154                              file=sys.stderr)
155                        if self.args.input == "-":
156                            print("""Not attempting to download Flame Graph template as script command line
157input is disabled due to using live mode. If you want to download the
158template retry without live mode. For example, use 'perf record -a -g
159-F 99 sleep 60' and 'perf script report flamegraph'. Alternatively,
160download the template from:
161https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/templates/d3-flamegraph-base.html
162and place it at:
163/usr/share/d3-flame-graph/d3-flamegraph-base.html""",
164                                  file=sys.stderr)
165                            quit()
166                        s = None
167                        while s != "y" and s != "n":
168                            s = input("Do you wish to download a template from cdn.jsdelivr.net? (this warning can be suppressed with --allow-download) [yn] ").lower()
169                        if s == "n":
170                            quit()
171                    template = "https://cdn.jsdelivr.net/npm/d3-flame-graph@4.1.3/dist/templates/d3-flamegraph-base.html"
172                    template_md5sum = "143e0d06ba69b8370b9848dcd6ae3f36"
173
174            try:
175                with urllib.request.urlopen(template) as template:
176                    output_str = "".join([
177                        l.decode("utf-8") for l in template.readlines()
178                    ])
179            except Exception as err:
180                print(f"Error reading template {template}: {err}\n"
181                      "a minimal flame graph will be generated", file=sys.stderr)
182                output_str = minimal_html
183                template_md5sum = None
184
185            if template_md5sum:
186                download_md5sum = hashlib.md5(output_str.encode("utf-8")).hexdigest()
187                if download_md5sum != template_md5sum:
188                    s = None
189                    while s != "y" and s != "n":
190                        s = input(f"""Unexpected template md5sum.
191{download_md5sum} != {template_md5sum}, for:
192{output_str}
193continue?[yn] """).lower()
194                    if s == "n":
195                        quit()
196
197            output_str = output_str.replace("/** @options_json **/", options_json)
198            output_str = output_str.replace("/** @flamegraph_json **/", stacks_json)
199
200            output_fn = self.args.output or "flamegraph.html"
201        else:
202            output_str = stacks_json
203            output_fn = self.args.output or "stacks.json"
204
205        if output_fn == "-":
206            with io.open(sys.stdout.fileno(), "w", encoding="utf-8", closefd=False) as out:
207                out.write(output_str)
208        else:
209            print("dumping data to {}".format(output_fn))
210            try:
211                with io.open(output_fn, "w", encoding="utf-8") as out:
212                    out.write(output_str)
213            except IOError as err:
214                print("Error writing output file: {}".format(err), file=sys.stderr)
215                sys.exit(1)
216
217
218if __name__ == "__main__":
219    parser = argparse.ArgumentParser(description="Create flame graphs.")
220    parser.add_argument("-f", "--format",
221                        default="html", choices=["json", "html"],
222                        help="output file format")
223    parser.add_argument("-o", "--output",
224                        help="output file name")
225    parser.add_argument("--template",
226                        default="/usr/share/d3-flame-graph/d3-flamegraph-base.html",
227                        help="path to flame graph HTML template")
228    parser.add_argument("--colorscheme",
229                        default="blue-green",
230                        help="flame graph color scheme",
231                        choices=["blue-green", "orange"])
232    parser.add_argument("-i", "--input",
233                        help=argparse.SUPPRESS)
234    parser.add_argument("--allow-download",
235                        default=False,
236                        action="store_true",
237                        help="allow unprompted downloading of HTML template")
238
239    cli_args = parser.parse_args()
240    cli = FlameGraphCLI(cli_args)
241
242    process_event = cli.process_event
243    trace_end = cli.trace_end
244