Merge pull request #15 from andreasf/flamegraph_support

Add support for FlameGraph output

plop/collector.py

@@ -5,26 +5,28 @@
 import sys
 import thread
 import time
+import argparse
 from plop import platform
 
+
 class Collector(object):
     MODES = {
         'prof': (platform.ITIMER_PROF, signal.SIGPROF),
         'virtual': (platform.ITIMER_VIRTUAL, signal.SIGVTALRM),
         'real': (platform.ITIMER_REAL, signal.SIGALRM),
-        }
+    }
 
     def __init__(self, interval=0.01, mode='virtual'):
         self.interval = interval
         self.mode = mode
         assert mode in Collector.MODES
         timer, sig = Collector.MODES[self.mode]
         signal.signal(sig, self.handler)
+        signal.siginterrupt(sig, False)
         self.reset()
 
     def reset(self):
-        # defaultdict instead of counter for pre-2.7 compatibility
-        self.stack_counts = collections.defaultdict(int)
+        self.stacks = list()
         self.samples_remaining = 0
         self.stopping = False
         self.stopped = False
@@ -45,7 +47,7 @@ def stop(self):
 
     def wait(self):
         while not self.stopped:
-            pass # need busy wait; ITIMER_PROF doesn't proceed while sleeping
+            pass  # need busy wait; ITIMER_PROF doesn't proceed while sleeping
 
     def handler(self, sig, current_frame):
         start = time.time()
@@ -63,56 +65,125 @@ def handler(self, sig, current_frame):
                 code = frame.f_code
                 frames.append((code.co_filename, code.co_firstlineno, code.co_name))
                 frame = frame.f_back
-            self.stack_counts[tuple(frames)] += 1
+            self.stacks.append(frames)
         end = time.time()
         self.samples_taken += 1
         self.sample_time += (end - start)
 
-    def filter(self, max_stacks):
-        self.stack_counts = dict(sorted(self.stack_counts.iteritems(), key=lambda kv: -kv[1])[:max_stacks])
+
+class CollectorFormatter(object):
+    """
+    Abstract class for output formats
+    """
+    def format(self, collector):
+        raise Exception("not implemented")
+
+    def store(self, collector, filename):
+        with open(filename, "w") as f:
+            f.write(self.format(collector))
+
+
+class PlopFormatter(CollectorFormatter):
+    """
+    Formats stack frames for plop.viewer
+    """
+    def __init__(self, max_stacks=50):
+        self.max_stacks = 50
+
+    def format(self, collector):
+        # defaultdict instead of counter for pre-2.7 compatibility
+        stack_counts = collections.defaultdict(int)
+        for frames in collector.stacks:
+            stack_counts[tuple(frames)] += 1
+        stack_counts = dict(sorted(stack_counts.iteritems(),
+                                key=lambda kv: -kv[1])[:self.max_stacks])
+        return repr(dict(stack_counts))
+
+
+class FlamegraphFormatter(CollectorFormatter):
+    """
+    Creates Flamegraph files
+    """
+    def format(self, collector):
+        output = ""
+        previous = None
+        previous_count = 1
+        for stack in collector.stacks:
+            current = self.format_flame(stack)
+            if current == previous:
+                previous_count += 1
+            else:
+                output += "%s %d\n" % (previous, previous_count)
+                previous_count = 1
+                previous = current
+        output += "%s %d\n" % (previous, previous_count)
+        return output
+
+    def format_flame(self, stack):
+        stack.reverse()
+        funcs = map(lambda stack: "%s (%s:%s)" % (stack[2], stack[0], stack[1]), stack)
+        return ";".join(funcs)
 
 
 def main():
     # TODO: more options, refactor this into somewhere shared
     # between tornado.autoreload and auto2to3
-    if len(sys.argv) >= 3 and sys.argv[1] == '-m':
-        mode = 'module'
-        module = filename_base = sys.argv[2]
-        del sys.argv[1:3]
-    elif len(sys.argv) >= 2:
-        mode = "script"
-        script = filename_base = sys.argv[1]
-        sys.argv = sys.argv[1:]
+    parser = argparse.ArgumentParser(description="Plop: Python Low-Overhead Profiler",
+                                     prog="python -m plop.collector",
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument("--format", "-f", help="Output format",
+                        choices=["plop", "flamegraph"], default="plop")
+    parser.add_argument("--module", "-m", help="Execute target as a module",
+                        action="store_const", const=True, default=False)
+    parser.add_argument("--mode", help="Interval timer mode to use, see `man 2 setitimer`",
+                        choices=["prof", "real", "virtual"], default="prof")
+    parser.add_argument("--interval", help="Timer interval in seconds", default=0.01, type=float)
+    parser.add_argument("--duration", help="Profiling duration in seconds", default=3600,
+                        type=int)
+    parser.add_argument("--max-stacks", help=("Number of most frequent stacks to store."
+                        " Ignored for Flamegraph output."), type=int, default=50)
+    parser.add_argument("target", help="Module or script to run")
+    parser.add_argument("arguments", nargs=argparse.REMAINDER,
+                        help="Pass-through arguments for the profiled application")
+    args = parser.parse_args()
+    sys.argv = [args.target] + args.arguments
+
+    if args.format == "flamegraph":
+        extension = "flame"
+        formatter = FlamegraphFormatter()
+    elif args.format == "plop":
+        extension = "plop"
+        formatter = PlopFormatter(max_stacks=args.max_stacks)
     else:
-        print "usage: python -m plop.collector -m module_to_run"
+        sys.stderr.write("Unhandled output format: %s" % args.format)
+        sys.stderr.flush()
         sys.exit(1)
+
     if not os.path.exists('profiles'):
         os.mkdir('profiles')
-    filename = 'profiles/%s-%s.plop' % (filename_base,
-                                        time.strftime('%Y%m%d-%H%M-%S'))
+    filename = 'profiles/%s-%s.%s' % (args.target, time.strftime('%Y%m%d-%H%M-%S'),
+                                      extension)
 
-    collector = Collector()
-    collector.start(duration=3600)
+    collector = Collector(mode=args.mode, interval=args.interval)
+    collector.start(duration=args.duration)
     exit_code = 0
     try:
-        if mode == "module":
+        if args.module:
             import runpy
-            runpy.run_module(module, run_name="__main__", alter_sys=True)
-        elif mode == "script":
-            with open(script) as f:
+            runpy.run_module(args.target, run_name="__main__", alter_sys=True)
+        else:
+            with open(args.target) as f:
                 global __file__
-                __file__ = script
+                __file__ = args.target
                 # Use globals as our "locals" dictionary so that
                 # something that tries to import __main__ (e.g. the unittest
                 # module) will see the right things.
                 exec f.read() in globals(), globals()
     except SystemExit, e:
         exit_code = e.code
     collector.stop()
-    collector.filter(50)
     if collector.samples_taken:
-        with open(filename, 'w') as f:
-            f.write(repr(dict(collector.stack_counts)))
+        formatter.store(collector, filename)
         print "profile output saved to %s" % filename
         overhead = float(collector.sample_time) / collector.samples_taken
         print "overhead was %s per sample (%s%%)" % (
@@ -121,6 +192,6 @@ def main():
         print "no samples collected; program was too fast"
     sys.exit(exit_code)
 
+
 if __name__ == '__main__':
     main()
-