Merge pull request #410 from goldshtn/argdist-silent

Print traceback only if verbose mode was requested

tools/argdist.py

@@ -16,6 +16,8 @@ from bcc import BPF
 from time import sleep, strftime
 import argparse
 import re
+import traceback
+import sys
 
 class Specifier(object):
         probe_text = """
@@ -443,7 +445,8 @@ u64 __time = bpf_ktime_get_ns();
                                 if not self.is_default_expr  else "retval")
                         data.print_log2_hist(val_type=label)
 
-examples = """
+class Tool(object):
+        examples = """
 Probe specifier syntax:
         {p,r}:[library]:function(signature)[:type[,type...]:expr[,expr...][:filter]][#label]
 Where:
@@ -510,69 +513,96 @@ argdist -p 2780 -z 120 \\
         of 120 characters
 """
 
-parser = argparse.ArgumentParser(description=
-        "Trace a function and display a summary of its parameter values.",
+        def __init__(self):
+                parser = argparse.ArgumentParser(description="Trace a " +
+                  "function and display a summary of its parameter values.",
                   formatter_class=argparse.RawDescriptionHelpFormatter,
-        epilog=examples)
-parser.add_argument("-p", "--pid", type=int,
+                  epilog=Tool.examples)
+                parser.add_argument("-p", "--pid", type=int,
                   help="id of the process to trace (optional)")
-parser.add_argument("-z", "--string-size", default=80, type=int,
+                parser.add_argument("-z", "--string-size", default=80,
+                  type=int,
                   help="maximum string size to read from char* arguments")
-parser.add_argument("-i", "--interval", default=1, type=int,
+                parser.add_argument("-i", "--interval", default=1, type=int,
                   help="output interval, in seconds")
-parser.add_argument("-n", "--number", type=int, dest="count",
+                parser.add_argument("-n", "--number", type=int, dest="count",
                   help="number of outputs")
-parser.add_argument("-v", "--verbose", action="store_true",
+                parser.add_argument("-v", "--verbose", action="store_true",
                   help="print resulting BPF program code before executing")
-parser.add_argument("-T", "--top", type=int,
-        help="number of top results to show (not applicable to histograms)")
-parser.add_argument("-H", "--histogram", nargs="*", dest="histspecifier",
-        metavar="specifier",
-        help="probe specifier to capture histogram of (see examples below)")
-parser.add_argument("-C", "--count", nargs="*", dest="countspecifier",
-        metavar="specifier",
-        help="probe specifier to capture count of (see examples below)")
-parser.add_argument("-I", "--include", nargs="*", metavar="header",
+                parser.add_argument("-T", "--top", type=int,
+                  help="number of top results to show (not applicable to " +
+                  "histograms)")
+                parser.add_argument("-H", "--histogram", nargs="*",
+                  dest="histspecifier", metavar="specifier",
+                  help="probe specifier to capture histogram of " +
+                  "(see examples below)")
+                parser.add_argument("-C", "--count", nargs="*",
+                  dest="countspecifier", metavar="specifier",
+                  help="probe specifier to capture count of " +
+                  "(see examples below)")
+                parser.add_argument("-I", "--include", nargs="*",
+                  metavar="header",
                   help="additional header files to include in the BPF program")
-args = parser.parse_args()
-
-specifiers = []
-for specifier in (args.countspecifier or []):
-        specifiers.append(Specifier("freq", specifier, args.pid))
-for histspecifier in (args.histspecifier or []):
-        specifiers.append(Specifier("hist", histspecifier, args.pid))
-if len(specifiers) == 0:
+                self.args = parser.parse_args()
+
+        def _create_specifiers(self):
+                self.specifiers = []
+                for specifier in (self.args.countspecifier or []):
+                        self.specifiers.append(Specifier(
+                                "freq", specifier, self.args.pid))
+                for histspecifier in (self.args.histspecifier or []):
+                        self.specifiers.append(
+                                Specifier("hist", histspecifier, self.args.pid))
+                if len(self.specifiers) == 0:
                         print("at least one specifier is required")
                         exit(1)
 
-bpf_source = """
+        def _generate_program(self):
+                bpf_source = """
 struct __string_t { char s[%d]; };
 
 #include <uapi/linux/ptrace.h>
-""" % args.string_size
-for include in (args.include or []):
+                """ % self.args.string_size
+                for include in (self.args.include or []):
                         bpf_source += "#include <%s>\n" % include
-bpf_source += Specifier.generate_auto_includes(map(lambda s: s.raw_spec, specifiers))
-for specifier in specifiers:
+                bpf_source += Specifier.generate_auto_includes(
+                                map(lambda s: s.raw_spec, self.specifiers))
+                for specifier in self.specifiers:
                         bpf_source += specifier.generate_text()
-
-if args.verbose:
+                if self.args.verbose:
                         print(bpf_source)
+                self.bpf = BPF(text=bpf_source)
 
-bpf = BPF(text=bpf_source)
-
-for specifier in specifiers:
-        specifier.attach(bpf)
+        def _attach(self):
+                for specifier in self.specifiers:
+                        specifier.attach(self.bpf)
 
-count_so_far = 0
-while True:
+        def _main_loop(self):
+                count_so_far = 0
+                while True:
                         try:
-                sleep(args.interval)
+                                sleep(self.args.interval)
                         except KeyboardInterrupt:
                                 exit()
                         print("[%s]" % strftime("%H:%M:%S"))
-        for specifier in specifiers:
-                specifier.display(args.top)
+                        for specifier in self.specifiers:
+                                specifier.display(self.args.top)
                         count_so_far += 1
-        if args.count is not None and count_so_far >= args.count:
+                        if self.args.count is not None and \
+                           count_so_far >= self.args.count:
                                 exit()
+
+        def run(self):
+                try:
+                        self._create_specifiers()
+                        self._generate_program()
+                        self._attach()
+                        self._main_loop()
+                except:
+                        if self.args.verbose:
+                                traceback.print_exc()
+                        else:
+                                print(sys.exc_value)
+
+if __name__ == "__main__":
+        Tool().run()