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 18from __future__ import print_function 19from bcc import ArgString, BPF, USDT 20from time import sleep, strftime 21import argparse 22import os 23import re 24import signal 25import sys 26import traceback 27 28debug = False 29 30def 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 36class 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""" 166int 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 179BPF_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 213class 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 297if __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