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