1#!/usr/bin/env python
2#
3# argdist   Trace a function and display a distribution of its
4#           parameter values as a histogram or frequency count.
5#
6# USAGE: argdist [-h] [-p PID] [-z STRING_SIZE] [-i INTERVAL] [-n COUNT] [-v]
7#                [-c] [-T TOP] [-C specifier] [-H specifier] [-I header]
8#
9# Licensed under the Apache License, Version 2.0 (the "License")
10# Copyright (C) 2016 Sasha Goldshtein.
11
12from bcc import BPF, USDT
13from time import sleep, strftime
14import argparse
15import re
16import traceback
17import os
18import sys
19
20class Probe(object):
21        next_probe_index = 0
22        streq_index = 0
23        aliases = {"$PID": "(bpf_get_current_pid_tgid() >> 32)"}
24
25        def _substitute_aliases(self, expr):
26                if expr is None:
27                        return expr
28                for alias, subst in Probe.aliases.items():
29                        expr = expr.replace(alias, subst)
30                return expr
31
32        def _parse_signature(self):
33                params = map(str.strip, self.signature.split(','))
34                self.param_types = {}
35                for param in params:
36                        # If the type is a pointer, the * can be next to the
37                        # param name. Other complex types like arrays are not
38                        # supported right now.
39                        index = param.rfind('*')
40                        index = index if index != -1 else param.rfind(' ')
41                        param_type = param[0:index + 1].strip()
42                        param_name = param[index + 1:].strip()
43                        self.param_types[param_name] = param_type
44
45        def _generate_entry(self):
46                self.entry_probe_func = self.probe_func_name + "_entry"
47                text = """
48int PROBENAME(struct pt_regs *ctx SIGNATURE)
49{
50        u64 __pid_tgid = bpf_get_current_pid_tgid();
51        u32 __pid      = __pid_tgid;        // lower 32 bits
52        u32 __tgid     = __pid_tgid >> 32;  // upper 32 bits
53        PID_FILTER
54        COLLECT
55        return 0;
56}
57"""
58                text = text.replace("PROBENAME", self.entry_probe_func)
59                text = text.replace("SIGNATURE",
60                     "" if len(self.signature) == 0 else ", " + self.signature)
61                text = text.replace("PID_FILTER", self._generate_pid_filter())
62                collect = ""
63                for pname in self.args_to_probe:
64                        param_hash = self.hashname_prefix + pname
65                        if pname == "__latency":
66                                collect += """
67u64 __time = bpf_ktime_get_ns();
68%s.update(&__pid, &__time);
69                        """ % param_hash
70                        else:
71                                collect += "%s.update(&__pid, &%s);\n" % \
72                                           (param_hash, pname)
73                text = text.replace("COLLECT", collect)
74                return text
75
76        def _generate_entry_probe(self):
77                # Any $entry(name) expressions result in saving that argument
78                # when entering the function.
79                self.args_to_probe = set()
80                regex = r"\$entry\((\w+)\)"
81                for expr in self.exprs:
82                        for arg in re.finditer(regex, expr):
83                                self.args_to_probe.add(arg.group(1))
84                for arg in re.finditer(regex, self.filter):
85                        self.args_to_probe.add(arg.group(1))
86                if any(map(lambda expr: "$latency" in expr, self.exprs)) or \
87                   "$latency" in self.filter:
88                        self.args_to_probe.add("__latency")
89                        self.param_types["__latency"] = "u64"    # nanoseconds
90                for pname in self.args_to_probe:
91                        if pname not in self.param_types:
92                                raise ValueError("$entry(%s): no such param" %
93                                                 arg)
94
95                self.hashname_prefix = "%s_param_" % self.probe_hash_name
96                text = ""
97                for pname in self.args_to_probe:
98                        # Each argument is stored in a separate hash that is
99                        # keyed by pid.
100                        text += "BPF_HASH(%s, u32, %s);\n" % \
101                             (self.hashname_prefix + pname,
102                              self.param_types[pname])
103                text += self._generate_entry()
104                return text
105
106        def _generate_retprobe_prefix(self):
107                # After we're done here, there are __%s_val variables for each
108                # argument we needed to probe using $entry(name), and they all
109                # have values (which isn't necessarily the case if we missed
110                # the method entry probe).
111                text = ""
112                self.param_val_names = {}
113                for pname in self.args_to_probe:
114                        val_name = "__%s_val" % pname
115                        text += "%s *%s = %s.lookup(&__pid);\n" % \
116                                (self.param_types[pname], val_name,
117                                 self.hashname_prefix + pname)
118                        text += "if (%s == 0) { return 0 ; }\n" % val_name
119                        self.param_val_names[pname] = val_name
120                return text
121
122        def _replace_entry_exprs(self):
123                for pname, vname in self.param_val_names.items():
124                        if pname == "__latency":
125                                entry_expr = "$latency"
126                                val_expr = "(bpf_ktime_get_ns() - *%s)" % vname
127                        else:
128                                entry_expr = "$entry(%s)" % pname
129                                val_expr = "(*%s)" % vname
130                        for i in range(0, len(self.exprs)):
131                                self.exprs[i] = self.exprs[i].replace(
132                                                entry_expr, val_expr)
133                        self.filter = self.filter.replace(entry_expr,
134                                                          val_expr)
135
136        def _attach_entry_probe(self):
137                if self.is_user:
138                        self.bpf.attach_uprobe(name=self.library,
139                                               sym=self.function,
140                                               fn_name=self.entry_probe_func,
141                                               pid=self.pid or -1)
142                else:
143                        self.bpf.attach_kprobe(event=self.function,
144                                               fn_name=self.entry_probe_func)
145
146        def _bail(self, error):
147                raise ValueError("error parsing probe '%s': %s" %
148                                 (self.raw_spec, error))
149
150        def _validate_specifier(self):
151                # Everything after '#' is the probe label, ignore it
152                spec = self.raw_spec.split('#')[0]
153                parts = spec.strip().split(':')
154                if len(parts) < 3:
155                        self._bail("at least the probe type, library, and " +
156                                   "function signature must be specified")
157                if len(parts) > 6:
158                        self._bail("extraneous ':'-separated parts detected")
159                if parts[0] not in ["r", "p", "t", "u"]:
160                        self._bail("probe type must be 'p', 'r', 't', or 'u'" +
161                                   " but got '%s'" % parts[0])
162                if re.match(r"\S+\(.*\)", parts[2]) is None:
163                        self._bail(("function signature '%s' has an invalid " +
164                                    "format") % parts[2])
165
166        def _parse_expr_types(self, expr_types):
167                if len(expr_types) == 0:
168                        self._bail("no expr types specified")
169                self.expr_types = expr_types.split(',')
170
171        def _parse_exprs(self, exprs):
172                if len(exprs) == 0:
173                        self._bail("no exprs specified")
174                self.exprs = exprs.split(',')
175
176        def _make_valid_identifier(self, ident):
177                return re.sub(r'[^A-Za-z0-9_]', '_', ident)
178
179        def __init__(self, tool, type, specifier):
180                self.usdt_ctx = None
181                self.streq_functions = ""
182                self.pid = tool.args.pid
183                self.cumulative = tool.args.cumulative or False
184                self.raw_spec = specifier
185                self._validate_specifier()
186
187                spec_and_label = specifier.split('#')
188                self.label = spec_and_label[1] \
189                             if len(spec_and_label) == 2 else None
190
191                parts = spec_and_label[0].strip().split(':')
192                self.type = type    # hist or freq
193                self.probe_type = parts[0]
194                fparts = parts[2].split('(')
195                self.function = fparts[0].strip()
196                if self.probe_type == "t":
197                        self.library = ""       # kernel
198                        self.tp_category = parts[1]
199                        self.tp_event = self.function
200                elif self.probe_type == "u":
201                        self.library = parts[1]
202                        self.probe_func_name = self._make_valid_identifier(
203                                "%s_probe%d" %
204                                (self.function, Probe.next_probe_index))
205                        self._enable_usdt_probe()
206                else:
207                        self.library = parts[1]
208                self.is_user = len(self.library) > 0
209                self.signature = fparts[1].strip()[:-1]
210                self._parse_signature()
211
212                # If the user didn't specify an expression to probe, we probe
213                # the retval in a ret probe, or simply the value "1" otherwise.
214                self.is_default_expr = len(parts) < 5
215                if not self.is_default_expr:
216                        self._parse_expr_types(parts[3])
217                        self._parse_exprs(parts[4])
218                        if len(self.exprs) != len(self.expr_types):
219                                self._bail("mismatched # of exprs and types")
220                        if self.type == "hist" and len(self.expr_types) > 1:
221                                self._bail("histograms can only have 1 expr")
222                else:
223                        if not self.probe_type == "r" and self.type == "hist":
224                                self._bail("histograms must have expr")
225                        self.expr_types = \
226                          ["u64" if not self.probe_type == "r" else "int"]
227                        self.exprs = \
228                          ["1" if not self.probe_type == "r" else "$retval"]
229                self.filter = "" if len(parts) != 6 else parts[5]
230                self._substitute_exprs()
231
232                # Do we need to attach an entry probe so that we can collect an
233                # argument that is required for an exit (return) probe?
234                def check(expr):
235                        keywords = ["$entry", "$latency"]
236                        return any(map(lambda kw: kw in expr, keywords))
237                self.entry_probe_required = self.probe_type == "r" and \
238                        (any(map(check, self.exprs)) or check(self.filter))
239
240                self.probe_func_name = self._make_valid_identifier(
241                        "%s_probe%d" %
242                        (self.function, Probe.next_probe_index))
243                self.probe_hash_name = self._make_valid_identifier(
244                        "%s_hash%d" %
245                        (self.function, Probe.next_probe_index))
246                Probe.next_probe_index += 1
247
248        def _enable_usdt_probe(self):
249                self.usdt_ctx = USDT(path=self.library, pid=self.pid)
250                self.usdt_ctx.enable_probe(
251                        self.function, self.probe_func_name)
252
253        def _generate_streq_function(self, string):
254                fname = "streq_%d" % Probe.streq_index
255                Probe.streq_index += 1
256                self.streq_functions += """
257static inline bool %s(char const *ignored, char const *str) {
258        char needle[] = %s;
259        char haystack[sizeof(needle)];
260        bpf_probe_read(&haystack, sizeof(haystack), (void *)str);
261        for (int i = 0; i < sizeof(needle) - 1; ++i) {
262                if (needle[i] != haystack[i]) {
263                        return false;
264                }
265        }
266        return true;
267}
268                """ % (fname, string)
269                return fname
270
271        def _substitute_exprs(self):
272                def repl(expr):
273                        expr = self._substitute_aliases(expr)
274                        matches = re.finditer('STRCMP\\(("[^"]+\\")', expr)
275                        for match in matches:
276                                string = match.group(1)
277                                fname = self._generate_streq_function(string)
278                                expr = expr.replace("STRCMP", fname, 1)
279                        return expr.replace("$retval", "PT_REGS_RC(ctx)")
280                for i in range(0, len(self.exprs)):
281                        self.exprs[i] = repl(self.exprs[i])
282                self.filter = repl(self.filter)
283
284        def _is_string(self, expr_type):
285                return expr_type == "char*" or expr_type == "char *"
286
287        def _generate_hash_field(self, i):
288                if self._is_string(self.expr_types[i]):
289                        return "struct __string_t v%d;\n" % i
290                else:
291                        return "%s v%d;\n" % (self.expr_types[i], i)
292
293        def _generate_usdt_arg_assignment(self, i):
294                expr = self.exprs[i]
295                if self.probe_type == "u" and expr[0:3] == "arg":
296                        arg_index = int(expr[3])
297                        arg_ctype = self.usdt_ctx.get_probe_arg_ctype(
298                                self.function, arg_index - 1)
299                        return ("        %s %s = 0;\n" +
300                                "        bpf_usdt_readarg(%s, ctx, &%s);\n") \
301                                % (arg_ctype, expr, expr[3], expr)
302                else:
303                        return ""
304
305        def _generate_field_assignment(self, i):
306                text = self._generate_usdt_arg_assignment(i)
307                if self._is_string(self.expr_types[i]):
308                        return (text + "        bpf_probe_read(&__key.v%d.s," +
309                                " sizeof(__key.v%d.s), (void *)%s);\n") % \
310                                (i, i, self.exprs[i])
311                else:
312                        return text + "        __key.v%d = %s;\n" % \
313                               (i, self.exprs[i])
314
315        def _generate_hash_decl(self):
316                if self.type == "hist":
317                        return "BPF_HISTOGRAM(%s, %s);" % \
318                               (self.probe_hash_name, self.expr_types[0])
319                else:
320                        text = "struct %s_key_t {\n" % self.probe_hash_name
321                        for i in range(0, len(self.expr_types)):
322                                text += self._generate_hash_field(i)
323                        text += "};\n"
324                        text += "BPF_HASH(%s, struct %s_key_t, u64);\n" % \
325                                (self.probe_hash_name, self.probe_hash_name)
326                        return text
327
328        def _generate_key_assignment(self):
329                if self.type == "hist":
330                        return self._generate_usdt_arg_assignment(0) + \
331                               ("%s __key = %s;\n" %
332                                (self.expr_types[0], self.exprs[0]))
333                else:
334                        text = "struct %s_key_t __key = {};\n" % \
335                                self.probe_hash_name
336                        for i in range(0, len(self.exprs)):
337                                text += self._generate_field_assignment(i)
338                        return text
339
340        def _generate_hash_update(self):
341                if self.type == "hist":
342                        return "%s.increment(bpf_log2l(__key));" % \
343                                self.probe_hash_name
344                else:
345                        return "%s.increment(__key);" % self.probe_hash_name
346
347        def _generate_pid_filter(self):
348                # Kernel probes need to explicitly filter pid, because the
349                # attach interface doesn't support pid filtering
350                if self.pid is not None and not self.is_user:
351                        return "if (__tgid != %d) { return 0; }" % self.pid
352                else:
353                        return ""
354
355        def generate_text(self):
356                program = ""
357                probe_text = """
358DATA_DECL
359                """ + (
360                    "TRACEPOINT_PROBE(%s, %s)" %
361                    (self.tp_category, self.tp_event)
362                    if self.probe_type == "t"
363                    else "int PROBENAME(struct pt_regs *ctx SIGNATURE)") + """
364{
365        u64 __pid_tgid = bpf_get_current_pid_tgid();
366        u32 __pid      = __pid_tgid;        // lower 32 bits
367        u32 __tgid     = __pid_tgid >> 32;  // upper 32 bits
368        PID_FILTER
369        PREFIX
370        if (!(FILTER)) return 0;
371        KEY_EXPR
372        COLLECT
373        return 0;
374}
375"""
376                prefix = ""
377                signature = ""
378
379                # If any entry arguments are probed in a ret probe, we need
380                # to generate an entry probe to collect them
381                if self.entry_probe_required:
382                        program += self._generate_entry_probe()
383                        prefix += self._generate_retprobe_prefix()
384                        # Replace $entry(paramname) with a reference to the
385                        # value we collected when entering the function:
386                        self._replace_entry_exprs()
387
388                if self.probe_type == "p" and len(self.signature) > 0:
389                        # Only entry uprobes/kprobes can have user-specified
390                        # signatures. Other probes force it to ().
391                        signature = ", " + self.signature
392
393                program += probe_text.replace("PROBENAME",
394                                              self.probe_func_name)
395                program = program.replace("SIGNATURE", signature)
396                program = program.replace("PID_FILTER",
397                                          self._generate_pid_filter())
398
399                decl = self._generate_hash_decl()
400                key_expr = self._generate_key_assignment()
401                collect = self._generate_hash_update()
402                program = program.replace("DATA_DECL", decl)
403                program = program.replace("KEY_EXPR", key_expr)
404                program = program.replace("FILTER",
405                        "1" if len(self.filter) == 0 else self.filter)
406                program = program.replace("COLLECT", collect)
407                program = program.replace("PREFIX", prefix)
408
409                return self.streq_functions + program
410
411        def _attach_u(self):
412                libpath = BPF.find_library(self.library)
413                if libpath is None:
414                        libpath = BPF.find_exe(self.library)
415                if libpath is None or len(libpath) == 0:
416                        self._bail("unable to find library %s" % self.library)
417
418                if self.probe_type == "r":
419                        self.bpf.attach_uretprobe(name=libpath,
420                                                  sym=self.function,
421                                                  fn_name=self.probe_func_name,
422                                                  pid=self.pid or -1)
423                else:
424                        self.bpf.attach_uprobe(name=libpath,
425                                               sym=self.function,
426                                               fn_name=self.probe_func_name,
427                                               pid=self.pid or -1)
428
429        def _attach_k(self):
430                if self.probe_type == "t":
431                        pass    # Nothing to do for tracepoints
432                elif self.probe_type == "r":
433                        self.bpf.attach_kretprobe(event=self.function,
434                                             fn_name=self.probe_func_name)
435                else:
436                        self.bpf.attach_kprobe(event=self.function,
437                                          fn_name=self.probe_func_name)
438
439        def attach(self, bpf):
440                self.bpf = bpf
441                if self.probe_type == "u":
442                        return
443                if self.is_user:
444                        self._attach_u()
445                else:
446                        self._attach_k()
447                if self.entry_probe_required:
448                        self._attach_entry_probe()
449
450        def _v2s(self, v):
451                # Most fields can be converted with plain str(), but strings
452                # are wrapped in a __string_t which has an .s field
453                if "__string_t" in type(v).__name__:
454                        return str(v.s)
455                return str(v)
456
457        def _display_expr(self, i):
458                # Replace ugly latency calculation with $latency
459                expr = self.exprs[i].replace(
460                        "(bpf_ktime_get_ns() - *____latency_val)", "$latency")
461                # Replace alias values back with the alias name
462                for alias, subst in Probe.aliases.items():
463                        expr = expr.replace(subst, alias)
464                # Replace retval expression with $retval
465                expr = expr.replace("PT_REGS_RC(ctx)", "$retval")
466                # Replace ugly (*__param_val) expressions with param name
467                return re.sub(r"\(\*__(\w+)_val\)", r"\1", expr)
468
469        def _display_key(self, key):
470                if self.is_default_expr:
471                        if not self.probe_type == "r":
472                                return "total calls"
473                        else:
474                                return "retval = %s" % str(key.v0)
475                else:
476                        # The key object has v0, ..., vk fields containing
477                        # the values of the expressions from self.exprs
478                        def str_i(i):
479                                key_i = self._v2s(getattr(key, "v%d" % i))
480                                return "%s = %s" % \
481                                        (self._display_expr(i), key_i)
482                        return ", ".join(map(str_i, range(0, len(self.exprs))))
483
484        def display(self, top):
485                data = self.bpf.get_table(self.probe_hash_name)
486                if self.type == "freq":
487                        print(self.label or self.raw_spec)
488                        print("\t%-10s %s" % ("COUNT", "EVENT"))
489                        sdata = sorted(data.items(), key=lambda p: p[1].value)
490                        if top is not None:
491                                sdata = sdata[-top:]
492                        for key, value in sdata:
493                                # Print some nice values if the user didn't
494                                # specify an expression to probe
495                                if self.is_default_expr:
496                                        if not self.probe_type == "r":
497                                                key_str = "total calls"
498                                        else:
499                                                key_str = "retval = %s" % \
500                                                          self._v2s(key.v0)
501                                else:
502                                        key_str = self._display_key(key)
503                                print("\t%-10s %s" %
504                                      (str(value.value), key_str))
505                elif self.type == "hist":
506                        label = self.label or (self._display_expr(0)
507                                if not self.is_default_expr else "retval")
508                        data.print_log2_hist(val_type=label)
509                if not self.cumulative:
510                        data.clear()
511
512        def __str__(self):
513                return self.label or self.raw_spec
514
515class Tool(object):
516        examples = """
517Probe specifier syntax:
518        {p,r,t,u}:{[library],category}:function(signature)[:type[,type...]:expr[,expr...][:filter]][#label]
519Where:
520        p,r,t,u    -- probe at function entry, function exit, kernel
521                      tracepoint, or USDT probe
522                      in exit probes: can use $retval, $entry(param), $latency
523        library    -- the library that contains the function
524                      (leave empty for kernel functions)
525        category   -- the category of the kernel tracepoint (e.g. net, sched)
526        function   -- the function name to trace (or tracepoint name)
527        signature  -- the function's parameters, as in the C header
528        type       -- the type of the expression to collect (supports multiple)
529        expr       -- the expression to collect (supports multiple)
530        filter     -- the filter that is applied to collected values
531        label      -- the label for this probe in the resulting output
532
533EXAMPLES:
534
535argdist -H 'p::__kmalloc(u64 size):u64:size'
536        Print a histogram of allocation sizes passed to kmalloc
537
538argdist -p 1005 -C 'p:c:malloc(size_t size):size_t:size:size==16'
539        Print a frequency count of how many times process 1005 called malloc
540        with an allocation size of 16 bytes
541
542argdist -C 'r:c:gets():char*:(char*)$retval#snooped strings'
543        Snoop on all strings returned by gets()
544
545argdist -H 'r::__kmalloc(size_t size):u64:$latency/$entry(size)#ns per byte'
546        Print a histogram of nanoseconds per byte from kmalloc allocations
547
548argdist -C 'p::__kmalloc(size_t sz, gfp_t flags):size_t:sz:flags&GFP_ATOMIC'
549        Print frequency count of kmalloc allocation sizes that have GFP_ATOMIC
550
551argdist -p 1005 -C 'p:c:write(int fd):int:fd' -T 5
552        Print frequency counts of how many times writes were issued to a
553        particular file descriptor number, in process 1005, but only show
554        the top 5 busiest fds
555
556argdist -p 1005 -H 'r:c:read()'
557        Print a histogram of results (sizes) returned by read() in process 1005
558
559argdist -C 'r::__vfs_read():u32:$PID:$latency > 100000'
560        Print frequency of reads by process where the latency was >0.1ms
561
562argdist -H 'r::__vfs_read(void *file, void *buf, size_t count):size_t:
563            $entry(count):$latency > 1000000'
564        Print a histogram of read sizes that were longer than 1ms
565
566argdist -H \\
567        'p:c:write(int fd, const void *buf, size_t count):size_t:count:fd==1'
568        Print a histogram of buffer sizes passed to write() across all
569        processes, where the file descriptor was 1 (STDOUT)
570
571argdist -C 'p:c:fork()#fork calls'
572        Count fork() calls in libc across all processes
573        Can also use funccount.py, which is easier and more flexible
574
575argdist -H 't:block:block_rq_complete():u32:args->nr_sector'
576        Print histogram of number of sectors in completing block I/O requests
577
578argdist -C 't:irq:irq_handler_entry():int:args->irq'
579        Aggregate interrupts by interrupt request (IRQ)
580
581argdist -C 'u:pthread:pthread_start():u64:arg2' -p 1337
582        Print frequency of function addresses used as a pthread start function,
583        relying on the USDT pthread_start probe in process 1337
584
585argdist -H 'p:c:sleep(u32 seconds):u32:seconds' \\
586        -H 'p:c:nanosleep(struct timespec *req):long:req->tv_nsec'
587        Print histograms of sleep() and nanosleep() parameter values
588
589argdist -p 2780 -z 120 \\
590        -C 'p:c:write(int fd, char* buf, size_t len):char*:buf:fd==1'
591        Spy on writes to STDOUT performed by process 2780, up to a string size
592        of 120 characters
593
594argdist -I 'kernel/sched/sched.h' \\
595        -C 'p::__account_cfs_rq_runtime(struct cfs_rq *cfs_rq):s64:cfs_rq->runtime_remaining'
596        Trace on the cfs scheduling runqueue remaining runtime. The struct cfs_rq is defined
597        in kernel/sched/sched.h which is in kernel source tree and not in kernel-devel
598        package.  So this command needs to run at the kernel source tree root directory
599        so that the added header file can be found by the compiler.
600"""
601
602        def __init__(self):
603                parser = argparse.ArgumentParser(description="Trace a " +
604                  "function and display a summary of its parameter values.",
605                  formatter_class=argparse.RawDescriptionHelpFormatter,
606                  epilog=Tool.examples)
607                parser.add_argument("-p", "--pid", type=int,
608                  help="id of the process to trace (optional)")
609                parser.add_argument("-z", "--string-size", default=80,
610                  type=int,
611                  help="maximum string size to read from char* arguments")
612                parser.add_argument("-i", "--interval", default=1, type=int,
613                  help="output interval, in seconds (default 1 second)")
614                parser.add_argument("-d", "--duration", type=int,
615                  help="total duration of trace, in seconds")
616                parser.add_argument("-n", "--number", type=int, dest="count",
617                  help="number of outputs")
618                parser.add_argument("-v", "--verbose", action="store_true",
619                  help="print resulting BPF program code before executing")
620                parser.add_argument("-c", "--cumulative", action="store_true",
621                  help="do not clear histograms and freq counts at " +
622                       "each interval")
623                parser.add_argument("-T", "--top", type=int,
624                  help="number of top results to show (not applicable to " +
625                  "histograms)")
626                parser.add_argument("-H", "--histogram", action="append",
627                  dest="histspecifier", metavar="specifier",
628                  help="probe specifier to capture histogram of " +
629                  "(see examples below)")
630                parser.add_argument("-C", "--count", action="append",
631                  dest="countspecifier", metavar="specifier",
632                  help="probe specifier to capture count of " +
633                  "(see examples below)")
634                parser.add_argument("-I", "--include", action="append",
635                  metavar="header",
636                  help="additional header files to include in the BPF program "
637                       "as either full path, "
638                       "or relative to relative to current working directory, "
639                       "or relative to default kernel header search path")
640                self.args = parser.parse_args()
641                self.usdt_ctx = None
642
643        def _create_probes(self):
644                self.probes = []
645                for specifier in (self.args.countspecifier or []):
646                        self.probes.append(Probe(self, "freq", specifier))
647                for histspecifier in (self.args.histspecifier or []):
648                        self.probes.append(Probe(self, "hist", histspecifier))
649                if len(self.probes) == 0:
650                        print("at least one specifier is required")
651                        exit(1)
652
653        def _generate_program(self):
654                bpf_source = """
655struct __string_t { char s[%d]; };
656
657#include <uapi/linux/ptrace.h>
658                """ % self.args.string_size
659                for include in (self.args.include or []):
660                        if include.startswith((".", "/")):
661                                include = os.path.abspath(include)
662                                bpf_source += "#include \"%s\"\n" % include
663                        else:
664                                bpf_source += "#include <%s>\n" % include
665
666                bpf_source += BPF.generate_auto_includes(
667                                map(lambda p: p.raw_spec, self.probes))
668                for probe in self.probes:
669                        bpf_source += probe.generate_text()
670                if self.args.verbose:
671                        for text in [probe.usdt_ctx.get_text()
672                                     for probe in self.probes
673                                     if probe.usdt_ctx]:
674                            print(text)
675                        print(bpf_source)
676                usdt_contexts = [probe.usdt_ctx
677                                 for probe in self.probes if probe.usdt_ctx]
678                self.bpf = BPF(text=bpf_source, usdt_contexts=usdt_contexts)
679
680        def _attach(self):
681                for probe in self.probes:
682                        probe.attach(self.bpf)
683                if self.args.verbose:
684                        print("open uprobes: %s" % list(self.bpf.uprobe_fds.keys()))
685                        print("open kprobes: %s" % list(self.bpf.kprobe_fds.keys()))
686
687        def _main_loop(self):
688                count_so_far = 0
689                seconds = 0
690                while True:
691                        try:
692                                sleep(self.args.interval)
693                                seconds += self.args.interval
694                        except KeyboardInterrupt:
695                                exit()
696                        print("[%s]" % strftime("%H:%M:%S"))
697                        for probe in self.probes:
698                                probe.display(self.args.top)
699                        count_so_far += 1
700                        if self.args.count is not None and \
701                           count_so_far >= self.args.count:
702                                exit()
703                        if self.args.duration and \
704                           seconds >= self.args.duration:
705                                exit()
706
707        def run(self):
708                try:
709                        self._create_probes()
710                        self._generate_program()
711                        self._attach()
712                        self._main_loop()
713                except:
714                        exc_info = sys.exc_info()
715                        sys_exit = exc_info[0] is SystemExit
716                        if self.args.verbose:
717                                traceback.print_exc()
718                        elif not sys_exit:
719                                print(exc_info[1])
720                        exit(0 if sys_exit else 1)
721
722if __name__ == "__main__":
723        Tool().run()
724