Shopify · shivnagarajan · May 17, 2016 · Feb 12, 2014 · Jun 17, 2014 · Sep 11, 2014
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,4 @@
 /tmp
 /lib/stackprof.bundle
 /lib/stackprof.so
+*.sw?
diff --git a/Gemfile.lock b/Gemfile.lock
@@ -1,22 +1,27 @@
 PATH
   remote: .
   specs:
-    stackprof (0.2.7)
+    stackprof (0.2.9)
 
 GEM
   remote: https://rubygems.org/
   specs:
-    metaclass (0.0.1)
+    metaclass (0.0.4)
+    minitest (5.4.1)
     mocha (0.14.0)
       metaclass (~> 0.0.1)
-    rake (10.1.0)
-    rake-compiler (0.9.1)
+    rake (10.3.2)
+    rake-compiler (0.9.3)
       rake
 
 PLATFORMS
   ruby
 
 DEPENDENCIES
+  minitest (~> 5.0)
   mocha (~> 0.14)
   rake-compiler (~> 0.9)
   stackprof!
+
+BUNDLED WITH
+   1.11.2
diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2013-2015 Aman Gupta
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
@@ -186,7 +186,7 @@ digraph profile {
 }
 ```
 
-#### `StackProf::Report.new(data).print_source(/pow|newobj|math/)`
+#### `StackProf::Report.new(data).print_method(/pow|newobj|math/)`
 
 ```
 A#pow (/Users/tmm1/code/stackprof/sample.rb:11)
@@ -265,11 +265,23 @@ multiple start/stop invocations.
 
 ``` ruby
 StackProf.running?
-StackProf.start
+StackProf.start(mode: :cpu)
 StackProf.stop
-StackProf.results
+StackProf.results('/tmp/some.file')
 ```
 
+### all options
+
+`StackProf.run` accepts an options hash. Currently, the following options are recognized:
+
+Option      | Meaning
+-------     | ---------
+`mode`      | mode of sampling: `:cpu`, `:wall`, `:object`, or `:custom` [c.f.](#sampling)
+`out`       | the target file, which will be overwritten
+`interval`  | mode-relative sample rate [c.f.](#sampling)
+`aggregate` | defaults: `true` - if `false` disables [aggregation](#aggregation)
+`raw`       | defaults `false` - if `true` collects the extra data required by the `--flamegraph` and `--stackcollapse` report types
+
 ### todo
 
 * file/iseq blacklist

diff --git a/Rakefile b/Rakefile
@@ -4,7 +4,7 @@ task :default => :test
 # Packaging
 # ==========================================================
 
-GEMSPEC = eval(File.read('stackprof.gemspec'))
+GEMSPEC = Gem::Specification::load('stackprof.gemspec')
 
 require 'rubygems/package_task'
 Gem::PackageTask.new(GEMSPEC) do |pkg|

diff --git a/bin/stackprof b/bin/stackprof
@@ -2,23 +2,20 @@
 require 'optparse'
 require 'stackprof'
 
-options = {
-  :format => :text,
-  :sort => false,
-  :limit => 30
-}
+options = {}
 
 parser = OptionParser.new(ARGV) do |o|
   o.banner = "Usage: stackprof [file.dump]+ [--text|--method=NAME|--callgrind|--graphviz]"
 
   o.on('--text', 'Text summary per method (default)'){ options[:format] = :text }
   o.on('--files', 'List of files'){ |f| options[:format] = :files }
-  o.on('--limit [num]', Integer, 'Limit --text or --files output to N lines'){ |n| options[:limit] = n }
+  o.on('--limit [num]', Integer, 'Limit --text, --files, or --graphviz output to N entries'){ |n| options[:limit] = n }
   o.on('--sort-total', "Sort --text or --files output on total samples\n\n"){ options[:sort] = true }
   o.on('--method [grep]', 'Zoom into specified method'){ |f| options[:format] = :method; options[:filter] = f }
   o.on('--file [grep]', "Show annotated code for specified file\n\n"){ |f| options[:format] = :file; options[:filter] = f }
   o.on('--callgrind', 'Callgrind output (use with kcachegrind, stackprof-gprof2dot.py)'){ options[:format] = :callgrind }
   o.on('--graphviz', "Graphviz output (use with dot)"){ options[:format] = :graphviz }
+  o.on('--node-fraction [frac]', OptionParser::DecimalNumeric, 'Drop nodes representing less than [frac] fraction of samples'){ |n| options[:node_fraction] = n }
   o.on('--stackcollapse', 'stackcollapse.pl compatible output (use with stackprof-flamegraph.pl)'){ options[:format] = :stackcollapse }
   o.on('--flamegraph', "timeline-flamegraph output (js)"){ options[:format] = :flamegraph }
   o.on('--flamegraph-viewer [f.js]', String, "open html viewer for flamegraph output\n\n"){ |file|
@@ -43,6 +40,20 @@ while ARGV.size > 0
 end
 report = reports.inject(:+)
 
+default_options = {
+  :format => :text,
+  :sort => false,
+  :limit => 30
+}
+
+if options[:format] == :graphviz
+  default_options[:limit] = 120
+  default_options[:node_fraction] = 0.005
+end
+
+options = default_options.merge(options)
+options.delete(:limit) if options[:limit] == 0
+
 case options[:format]
 when :text
   report.print_text(options[:sort], options[:limit])
@@ -53,7 +64,7 @@ when :dump
 when :callgrind
   report.print_callgrind
 when :graphviz
-  report.print_graphviz
+  report.print_graphviz(options)
 when :stackcollapse
   report.print_stackcollapse
 when :flamegraph

diff --git a/ext/stackprof.c b/ext/stackprof.c
@@ -529,7 +529,7 @@ Init_stackprof(void)
     S(aggregate);
 #undef S
 
-    gc_hook = Data_Wrap_Struct(rb_cObject, stackprof_gc_mark, NULL, NULL);
+    gc_hook = Data_Wrap_Struct(rb_cObject, stackprof_gc_mark, NULL, &_stackprof);
     rb_global_variable(&gc_hook);
 
     rb_mStackProf = rb_define_module("StackProf");

diff --git a/lib/stackprof/middleware.rb b/lib/stackprof/middleware.rb
@@ -9,15 +9,15 @@ def initialize(app, options = {})
 
       Middleware.mode     = options[:mode] || :cpu
       Middleware.interval = options[:interval] || 1000
+      Middleware.raw      = options[:raw] || false
       Middleware.enabled  = options[:enabled]
       Middleware.path     = options[:path] || 'tmp'
-      Middleware.raw      = options[:raw] || false
       at_exit{ Middleware.save } if options[:save_at_exit]
     end
 
     def call(env)
-      enabled, mode = Middleware.enabled?
-      StackProf.start(mode: mode || Middleware.mode, raw: Middleware.raw, interval: Middleware.interval) if enabled
+      enabled, mode = Middleware.enabled?(env)
+      StackProf.start(mode: mode || Middleware.mode, interval: Middleware.interval, raw: Middleware.raw) if enabled
       @app.call(env)
     ensure
       if enabled
@@ -30,11 +30,11 @@ def call(env)
     end
 
     class << self
-      attr_accessor :enabled, :mode, :interval, :path, :raw
+      attr_accessor :enabled, :mode, :interval, :raw, :path
 
-      def enabled?
+      def enabled?(env)
         if enabled.respond_to?(:call)
-          enabled.call
+          enabled.call(env)
         else
           enabled
         end

diff --git a/lib/stackprof/report.rb b/lib/stackprof/report.rb
@@ -9,7 +9,7 @@ def initialize(data)
     attr_reader :data
 
     def frames(sort_by_total=false)
-      Hash[ *@data[:frames].sort_by{ |iseq, stats| -stats[sort_by_total ? :total_samples : :samples] }.flatten(1) ]
+      @data[:frames].sort_by{ |iseq, stats| -stats[sort_by_total ? :total_samples : :samples] }.inject({}){|h, (k, v)| h[k] = v; h}
     end
 
     # normalized frames is used when we want to combine multiple
@@ -81,12 +81,12 @@ def print_debug
       pp @data
     end
 
-    def print_dump
-      puts Marshal.dump(@data.reject{|k,v| k == :files })
+    def print_dump(f=STDOUT)
+      f.puts Marshal.dump(@data.reject{|k,v| k == :files })
     end
 
     def print_stackcollapse
-      raise "profile does not include raw samples" unless raw = data[:raw]
+      raise "profile does not include raw samples (add `raw: true` to collecting StackProf.run)" unless raw = data[:raw]
 
       while len = raw.shift
         frames = raw.slice!(0, len)
@@ -98,7 +98,7 @@ def print_stackcollapse
     end
 
     def print_flamegraph(f=STDOUT, skip_common=true)
-      raise "profile does not include raw samples" unless raw = data[:raw]
+      raise "profile does not include raw samples (add `raw: true` to collecting StackProf.run)" unless raw = data[:raw]
 
       stacks = []
       max_x = 0
@@ -165,38 +165,60 @@ def flamegraph_row(f, x, y, weight, addr)
       f.puts %{{"x":#{x},"y":#{y},"width":#{weight},"frame_id":#{addr},"frame":#{frame[:name].dump},"file":#{frame[:file].dump}}}
     end
 
-    def print_graphviz(filter = nil, f = STDOUT)
-      if filter
+    def print_graphviz(options = {}, f = STDOUT)
+      if filter = options[:filter]
         mark_stack = []
-        list = frames
+        list = frames(true)
         list.each{ |addr, frame| mark_stack << addr if frame[:name] =~ filter }
         while addr = mark_stack.pop
           frame = list[addr]
           unless frame[:marked]
-            $stderr.puts frame[:edges].inspect
-            mark_stack += frame[:edges].map{ |addr, weight| addr.to_s if list[addr.to_s][:total_samples] <= weight*1.2 }.compact if frame[:edges]
+            mark_stack += frame[:edges].map{ |addr, weight| addr if list[addr][:total_samples] <= weight*1.2 }.compact if frame[:edges]
             frame[:marked] = true
           end
         end
         list = list.select{ |addr, frame| frame[:marked] }
-        list.each{ |addr, frame| frame[:edges] && frame[:edges].delete_if{ |k,v| list[k.to_s].nil? } }
+        list.each{ |addr, frame| frame[:edges] && frame[:edges].delete_if{ |k,v| list[k].nil? } }
         list
       else
-        list = frames
+        list = frames(true)
       end
 
+
+      limit = options[:limit]
+      fraction = options[:node_fraction]
+
+      included_nodes = {}
+      node_minimum = fraction ? (fraction * overall_samples).ceil : 0
+
       f.puts "digraph profile {"
-      list.each do |frame, info|
+      f.puts "Legend [shape=box,fontsize=24,shape=plaintext,label=\""
+      f.print "Total samples: #{overall_samples}\\l"
+      f.print "Showing top #{limit} nodes\\l" if limit
+      f.print "Dropped nodes with < #{node_minimum} samples\\l" if fraction
+      f.puts "\"];"
+
+      list.each_with_index do |(frame, info), index|
         call, total = info.values_at(:samples, :total_samples)
+        break if total < node_minimum || (limit && index >= limit)
+
         sample = ''
         sample << "#{call} (%2.1f%%)\\rof " % (call*100.0/overall_samples) if call < total
         sample << "#{total} (%2.1f%%)\\r" % (total*100.0/overall_samples)
         fontsize = (1.0 * call / max_samples) * 28 + 10
         size = (1.0 * total / overall_samples) * 2.0 + 0.5
 
         f.puts "  \"#{frame}\" [size=#{size}] [fontsize=#{fontsize}] [penwidth=\"#{size}\"] [shape=box] [label=\"#{info[:name]}\\n#{sample}\"];"
+        included_nodes[frame] = true
+      end
+
+      list.each do |frame, info|
+        next unless included_nodes[frame]
+
         if edges = info[:edges]
           edges.each do |edge, weight|
+            next unless included_nodes[edge]
+
             size = (1.0 * weight / overall_samples) * 2.0 + 0.5
             f.puts "  \"#{frame}\" -> \"#{edge}\" [label=\"#{weight}\"] [weight=\"#{weight}\"] [penwidth=\"#{size}\"];"
           end
@@ -252,7 +274,7 @@ def print_callgrind(f = STDOUT)
     end
 
     def print_method(name, f = STDOUT)
-      name = /#{Regexp.escape name}/ unless Regexp === name
+      name = /#{name}/ unless Regexp === name
       frames.each do |frame, info|
         next unless info[:name] =~ name
         file, line = info.values_at(:file, :line)
@@ -296,8 +318,7 @@ def print_files(sort_by_total=false, limit=nil, f = STDOUT)
 
     def print_file(filter, f = STDOUT)
       filter = /#{Regexp.escape filter}/ unless Regexp === filter
-      list = files
-      list.select!{ |name, lines| name =~ filter }
+      list = files.select{ |name, lines| name =~ filter }
       list.sort_by{ |file, vals| -vals.values.inject(0){ |sum, n| sum + (n.is_a?(Array) ? n[1] : n) } }.each do |file, lines|
         source_display(f, file, lines)
       end

diff --git a/sample.rb b/sample.rb
@@ -34,7 +34,7 @@ def math
 
 result = StackProf::Report.new(profile)
 puts
-result.print_source(/pow|newobj|math/)
+result.print_method(/pow|newobj|math/)
 puts
 result.print_text
 puts

diff --git a/stackprof.gemspec b/stackprof.gemspec
@@ -1,6 +1,6 @@
 Gem::Specification.new do |s|
   s.name = 'stackprof'
-  s.version = '0.2.7'
+  s.version = '0.2.9'
   s.homepage = 'http://github.com/tmm1/stackprof'
 
   s.authors = 'Aman Gupta'
@@ -21,4 +21,5 @@ Gem::Specification.new do |s|
 
   s.add_development_dependency 'rake-compiler', '~> 0.9'
   s.add_development_dependency 'mocha', '~> 0.14'
+  s.add_development_dependency 'minitest', '~> 5.0'
 end