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