xref: /aosp_15_r20/external/bcc/tools/funccount.py (revision 387f9dfdfa2baef462e92476d413c7bc2470293e)
1#!/usr/bin/env python
2# @lint-avoid-python-3-compatibility-imports
3#
4# funccount Count functions, tracepoints, and USDT probes.
5#           For Linux, uses BCC, eBPF.
6#
7# USAGE: funccount [-h] [-p PID] [-i INTERVAL] [-d DURATION] [-T] [-r]
8#                  [-c CPU] pattern
9#
10# The pattern is a string with optional '*' wildcards, similar to file
11# globbing. If you'd prefer to use regular expressions, use the -r option.
12#
13# Copyright (c) 2015 Brendan Gregg.
14# Licensed under the Apache License, Version 2.0 (the "License")
15#
16# 09-Sep-2015   Brendan Gregg       Created this.
17# 18-Oct-2016   Sasha Goldshtein    Generalized for uprobes, tracepoints, USDT.
18
19from __future__ import print_function
20from bcc import ArgString, BPF, USDT
21from time import sleep, strftime
22import argparse
23import re
24import signal
25import sys
26import traceback
27
28debug = False
29
30def verify_limit(num):
31    probe_limit = BPF.get_probe_limit()
32    if num > probe_limit:
33        raise Exception("maximum of %d probes allowed, attempted %d" %
34                        (probe_limit, num))
35
36class Probe(object):
37    def __init__(self, pattern, use_regex=False, pid=None, cpu=None):
38        """Init a new probe.
39
40        Init the probe from the pattern provided by the user. The supported
41        patterns mimic the 'trace' and 'argdist' tools, but are simpler because
42        we don't have to distinguish between probes and retprobes.
43
44            func            -- probe a kernel function
45            lib:func        -- probe a user-space function in the library 'lib'
46            /path:func      -- probe a user-space function in binary '/path'
47            p::func         -- same thing as 'func'
48            p:lib:func      -- same thing as 'lib:func'
49            t:cat:event     -- probe a kernel tracepoint
50            u:lib:probe     -- probe a USDT tracepoint
51        """
52        parts = bytes(pattern).split(b':')
53        if len(parts) == 1:
54            parts = [b"p", b"", parts[0]]
55        elif len(parts) == 2:
56            parts = [b"p", parts[0], parts[1]]
57        elif len(parts) == 3:
58            if parts[0] == b"t":
59                parts = [b"t", b"", b"%s:%s" % tuple(parts[1:])]
60            if parts[0] not in [b"p", b"t", b"u"]:
61                raise Exception("Type must be 'p', 't', or 'u', but got %s" %
62                                parts[0])
63        else:
64            raise Exception("Too many ':'-separated components in pattern %s" %
65                            pattern)
66
67        (self.type, self.library, self.pattern) = parts
68        if not use_regex:
69            self.pattern = self.pattern.replace(b'*', b'.*')
70            self.pattern = b'^' + self.pattern + b'$'
71
72        if (self.type == b"p" and self.library) or self.type == b"u":
73            libpath = BPF.find_library(self.library)
74            if libpath is None:
75                # This might be an executable (e.g. 'bash')
76                libpath = BPF.find_exe(str(self.library))
77            if libpath is None or len(libpath) == 0:
78                raise Exception("unable to find library %s" % self.library)
79            self.library = libpath
80
81        self.pid = pid
82        self.cpu = cpu
83        self.matched = 0
84        self.trace_functions = {}   # map location number to function name
85
86    def is_kernel_probe(self):
87        return self.type == b"t" or (self.type == b"p" and self.library == b"")
88
89    def attach(self):
90        if self.type == b"p" and not self.library:
91            for index, function in self.trace_functions.items():
92                self.bpf.attach_kprobe(
93                        event=function,
94                        fn_name="trace_count_%d" % index)
95        elif self.type == b"p" and self.library:
96            for index, function in self.trace_functions.items():
97                self.bpf.attach_uprobe(
98                        name=self.library,
99                        sym=function,
100                        fn_name="trace_count_%d" % index,
101                        pid=self.pid or -1)
102        elif self.type == b"t":
103            for index, function in self.trace_functions.items():
104                self.bpf.attach_tracepoint(
105                        tp=function,
106                        fn_name="trace_count_%d" % index)
107        elif self.type == b"u":
108            pass    # Nothing to do -- attach already happened in `load`
109
110    def _add_function(self, template, probe_name):
111        new_func = b"trace_count_%d" % self.matched
112        text = template.replace(b"PROBE_FUNCTION", new_func)
113        text = text.replace(b"LOCATION", b"%d" % self.matched)
114        self.trace_functions[self.matched] = probe_name
115        self.matched += 1
116        return text
117
118    def _generate_functions(self, template):
119        self.usdt = None
120        text = b""
121        if self.type == b"p" and not self.library:
122            functions = BPF.get_kprobe_functions(self.pattern)
123            verify_limit(len(functions))
124            for function in functions:
125                text += self._add_function(template, function)
126        elif self.type == b"p" and self.library:
127            # uprobes are tricky because the same function may have multiple
128            # addresses, and the same address may be mapped to multiple
129            # functions. We aren't allowed to create more than one uprobe
130            # per address, so track unique addresses and ignore functions that
131            # map to an address that we've already seen. Also ignore functions
132            # that may repeat multiple times with different addresses.
133            addresses, functions = (set(), set())
134            functions_and_addresses = BPF.get_user_functions_and_addresses(
135                                        self.library, self.pattern)
136            verify_limit(len(functions_and_addresses))
137            for function, address in functions_and_addresses:
138                if address in addresses or function in functions:
139                    continue
140                addresses.add(address)
141                functions.add(function)
142                text += self._add_function(template, function)
143        elif self.type == b"t":
144            tracepoints = BPF.get_tracepoints(self.pattern)
145            verify_limit(len(tracepoints))
146            for tracepoint in tracepoints:
147                text += self._add_function(template, tracepoint)
148        elif self.type == b"u":
149            self.usdt = USDT(path=str(self.library), pid=self.pid)
150            matches = []
151            for probe in self.usdt.enumerate_probes():
152                if not self.pid and (probe.bin_path != self.library):
153                    continue
154                if re.match(self.pattern, probe.name):
155                    matches.append(probe.name)
156            verify_limit(len(matches))
157            for match in matches:
158                new_func = b"trace_count_%d" % self.matched
159                text += self._add_function(template, match)
160                self.usdt.enable_probe(match, new_func)
161            if debug:
162                print(self.usdt.get_text())
163        return text
164
165    def load(self):
166        trace_count_text = b"""
167int PROBE_FUNCTION(void *ctx) {
168    FILTERPID
169    FILTERCPU
170    int loc = LOCATION;
171    counts.atomic_increment(loc);
172    return 0;
173}
174        """
175        bpf_text = b"""#include <uapi/linux/ptrace.h>
176
177BPF_ARRAY(counts, u64, NUMLOCATIONS);
178        """
179
180        # We really mean the tgid from the kernel's perspective, which is in
181        # the top 32 bits of bpf_get_current_pid_tgid().
182        if self.pid:
183            trace_count_text = trace_count_text.replace(b'FILTERPID',
184                b"""u32 pid = bpf_get_current_pid_tgid() >> 32;
185                   if (pid != %d) { return 0; }""" % self.pid)
186        else:
187            trace_count_text = trace_count_text.replace(b'FILTERPID', b'')
188
189        if self.cpu:
190            trace_count_text = trace_count_text.replace(b'FILTERCPU',
191                b"""u32 cpu = bpf_get_smp_processor_id();
192                   if (cpu != %d) { return 0; }""" % int(self.cpu))
193        else:
194            trace_count_text = trace_count_text.replace(b'FILTERCPU', b'')
195
196        bpf_text += self._generate_functions(trace_count_text)
197        bpf_text = bpf_text.replace(b"NUMLOCATIONS",
198                                    b"%d" % len(self.trace_functions))
199        if debug:
200            print(bpf_text)
201
202        if self.matched == 0:
203            raise Exception("No functions matched by pattern %s" %
204                            self.pattern)
205
206        self.bpf = BPF(text=bpf_text,
207                       usdt_contexts=[self.usdt] if self.usdt else [])
208        self.clear()    # Initialize all array items to zero
209
210    def counts(self):
211        return self.bpf["counts"]
212
213    def clear(self):
214        counts = self.bpf["counts"]
215        for location, _ in list(self.trace_functions.items()):
216            counts[counts.Key(location)] = counts.Leaf()
217
218class Tool(object):
219    def __init__(self):
220        examples = """examples:
221    ./funccount 'vfs_*'             # count kernel fns starting with "vfs"
222    ./funccount -r '^vfs.*'         # same as above, using regular expressions
223    ./funccount -Ti 5 'vfs_*'       # output every 5 seconds, with timestamps
224    ./funccount -d 10 'vfs_*'       # trace for 10 seconds only
225    ./funccount -p 185 'vfs_*'      # count vfs calls for PID 181 only
226    ./funccount t:sched:sched_fork  # count calls to the sched_fork tracepoint
227    ./funccount -p 185 u:node:gc*   # count all GC USDT probes in node, PID 185
228    ./funccount c:malloc            # count all malloc() calls in libc
229    ./funccount go:os.*             # count all "os.*" calls in libgo
230    ./funccount -p 185 go:os.*      # count all "os.*" calls in libgo, PID 185
231    ./funccount ./test:read*        # count "read*" calls in the ./test binary
232    ./funccount -c 1 'vfs_*'        # count vfs calls on CPU 1 only
233    """
234        parser = argparse.ArgumentParser(
235            description="Count functions, tracepoints, and USDT probes",
236            formatter_class=argparse.RawDescriptionHelpFormatter,
237            epilog=examples)
238        parser.add_argument("-p", "--pid", type=int,
239            help="trace this PID only")
240        parser.add_argument("-i", "--interval",
241            help="summary interval, seconds")
242        parser.add_argument("-d", "--duration",
243            help="total duration of trace, seconds")
244        parser.add_argument("-T", "--timestamp", action="store_true",
245            help="include timestamp on output")
246        parser.add_argument("-r", "--regexp", action="store_true",
247            help="use regular expressions. Default is \"*\" wildcards only.")
248        parser.add_argument("-D", "--debug", action="store_true",
249            help="print BPF program before starting (for debugging purposes)")
250        parser.add_argument("-c", "--cpu",
251            help="trace this CPU only")
252        parser.add_argument("pattern",
253            type=ArgString,
254            help="search expression for events")
255        self.args = parser.parse_args()
256        global debug
257        debug = self.args.debug
258        self.probe = Probe(self.args.pattern, self.args.regexp, self.args.pid,
259                           self.args.cpu)
260        if self.args.duration and not self.args.interval:
261            self.args.interval = self.args.duration
262        if not self.args.interval:
263            self.args.interval = 99999999
264
265    @staticmethod
266    def _signal_ignore(signal, frame):
267        print()
268
269    def run(self):
270        self.probe.load()
271        self.probe.attach()
272        print("Tracing %d functions for \"%s\"... Hit Ctrl-C to end." %
273              (self.probe.matched, bytes(self.args.pattern)))
274        exiting = 0 if self.args.interval else 1
275        seconds = 0
276        while True:
277            try:
278                sleep(int(self.args.interval))
279                seconds += int(self.args.interval)
280            except KeyboardInterrupt:
281                exiting = 1
282                # as cleanup can take many seconds, trap Ctrl-C:
283                signal.signal(signal.SIGINT, Tool._signal_ignore)
284            if self.args.duration and seconds >= int(self.args.duration):
285                exiting = 1
286
287            print()
288            if self.args.timestamp:
289                print("%-8s\n" % strftime("%H:%M:%S"), end="")
290
291            print("%-36s %8s" % ("FUNC", "COUNT"))
292            counts = self.probe.counts()
293            for k, v in sorted(counts.items(),
294                               key=lambda counts: counts[1].value):
295                if v.value == 0:
296                    continue
297                print("%-36s %8d" %
298                      (self.probe.trace_functions[k.value].decode('utf-8', 'replace'), v.value))
299
300            if exiting:
301                print("Detaching...")
302                exit()
303            else:
304                self.probe.clear()
305
306if __name__ == "__main__":
307    try:
308        Tool().run()
309    except Exception:
310        if debug:
311            traceback.print_exc()
312        elif sys.exc_info()[0] is not SystemExit:
313            print(sys.exc_info()[1])
314