Delegate save to saviour proc

Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    stackprof (0.2.7)
+    stackprof (0.2.9)
 
 GEM
   remote: https://rubygems.org/
@@ -22,3 +22,6 @@ DEPENDENCIES
   mocha (~> 0.14)
   rake-compiler (~> 0.9)
   stackprof!
+
+BUNDLED WITH
+   1.11.2

lib/stackprof/middleware.rb

@@ -11,26 +11,27 @@ def initialize(app, options = {})
       Middleware.interval = options[:interval] || 1000
       Middleware.raw      = options[:raw] || false
       Middleware.enabled  = options[:enabled]
+      Middleware.saviour  = options[:saviour]
       Middleware.path     = options[:path] || 'tmp'
       at_exit{ Middleware.save } if options[:save_at_exit]
     end
 
     def call(env)
-      enabled = Middleware.enabled?(env)
-      StackProf.start(mode: Middleware.mode, interval: Middleware.interval, raw: Middleware.raw) 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
         StackProf.stop
         if @num_reqs && (@num_reqs-=1) == 0
           @num_reqs = @options[:save_every]
-          Middleware.save
+          Middleware.save(env)
         end
       end
     end
 
     class << self
-      attr_accessor :enabled, :mode, :interval, :raw, :path
+      attr_accessor :enabled, :saviour, :mode, :interval, :raw, :path
 
       def enabled?(env)
         if enabled.respond_to?(:call)
@@ -40,14 +41,18 @@ def enabled?(env)
         end
       end
 
-      def save(filename = nil)
+      def save(env = nil, filename = nil)
         if results = StackProf.results
-          FileUtils.mkdir_p(Middleware.path)
-          filename ||= "stackprof-#{results[:mode]}-#{Process.pid}-#{Time.now.to_i}.dump"
-          File.open(File.join(Middleware.path, filename), 'wb') do |f|
-            f.write Marshal.dump(results)
+          if saviour.respond_to?(:call)
+            saviour.call(env, results)
+          else
+            FileUtils.mkdir_p(Middleware.path)
+            filename ||= "stackprof-#{results[:mode]}-#{Process.pid}-#{Time.now.to_i}.dump"
+            File.open(File.join(Middleware.path, filename), 'wb') do |f|
+              f.write Marshal.dump(results)
+            end
+            filename
           end
-          filename
         end
       end
 

lib/stackprof/report.rb

@@ -12,16 +12,34 @@ def frames(sort_by_total=false)
       @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
+    # output files, (via +() below). In order to simplify combination
+    # of files that are from "different" source revisions, the frames
+    # are normalised by converting them to an MD5 using the "name, file, line"
+    # once frames are converted, every edge that references the original
+    # needs to adjusted also.
     def normalized_frames
       id2hash = {}
       @data[:frames].each do |frame, info|
-        id2hash[frame.to_s] = info[:hash] = Digest::MD5.hexdigest("#{info[:name]}#{info[:file]}#{info[:line]}")
+        # id2 hash contains the original frames converted into a MD5
+        # based on the name, file, line of the original frame
+        # flamegraph likes the frames to be numbers, so md5.to_i(16)
+        id2hash[frame.to_s] = info[:hash] = Digest::MD5.hexdigest("#{info[:name]}#{info[:file]}#{info[:line]}").to_i(16)
       end
-      @data[:frames].inject(Hash.new) do |hash, (frame, info)|
+      # Convert all existing raw frames to use the new mapping via id2hash
+      # the array of raw frames contains a sequence of frame slices. Each
+      # frame slice is preceded by a count of the number of subsequent frames
+      # in the slice. Since the counts need not be normalized and we can
+      # identify frames by looking for a mapping in id2hash, add the
+      # normalized frame value when it is available, and just copy the
+      # existing value otherwise
+      raw_frames = @data[:raw].map {|v| id2hash[v.to_s] || v } if @data[:raw]
+      # return both the normalized frames and the normalized raw frames
+      return @data[:frames].inject(Hash.new) do |hash, (frame, info)|
         info = hash[id2hash[frame.to_s]] = info.dup
         info[:edges] = info[:edges].inject(Hash.new){ |edges, (edge, weight)| edges[id2hash[edge.to_s]] = weight; edges } if info[:edges]
         hash
-      end
+      end, raw_frames || []
     end
 
     def version
@@ -311,7 +329,9 @@ def +(other)
       raise ArgumentError, "cannot combine #{modeline} with #{other.modeline}" unless modeline == other.modeline
       raise ArgumentError, "cannot combine v#{version} with v#{other.version}" unless version == other.version
 
-      f1, f2 = normalized_frames, other.normalized_frames
+      # collect the normalized and the normalized raw frames
+      f1, raw_frames1 = normalized_frames
+      f2, raw_frames2 = other.normalized_frames
       frames = (f1.keys + f2.keys).uniq.inject(Hash.new) do |hash, id|
         if f1[id].nil?
           hash[id] = f2[id]
@@ -339,14 +359,17 @@ def +(other)
       end
 
       d1, d2 = data, other.data
+      # ensure that the new data contains the normalized frames
+      # and the normalized raw frames also
       data = {
         version: version,
         mode: d1[:mode],
         interval: d1[:interval],
         samples: d1[:samples] + d2[:samples],
         gc_samples: d1[:gc_samples] + d2[:gc_samples],
         missed_samples: d1[:missed_samples] + d2[:missed_samples],
-        frames: frames
+        frames: frames,
+        raw: raw_frames1 + raw_frames2
       }
 
       self.class.new(data)

test/test_middleware.rb

@@ -38,6 +38,35 @@ def test_save_custom
     StackProf::Middleware.save
   end
 
+  def test_save_should_use_a_proc_if_passed
+    StackProf.stubs(:results).returns({ mode: 'foo' })
+    FileUtils.expects(:mkdir_p).with('/foo').never
+    File.expects(:open).with(regexp_matches(/^\/foo\/stackprof-foo/), 'wb').never
+
+    proc_called = false
+    StackProf::Middleware.new(Object.new, saviour: Proc.new{ proc_called = true })
+    StackProf::Middleware.save
+    assert proc_called
+  end
+
+  def test_save_proc_should_receive_env_in_proc_if_passed
+    StackProf.stubs(:results).returns({ mode: 'foo' })
+
+    env_set = nil
+    StackProf::Middleware.new(Object.new, saviour: Proc.new{ |env, results| env_set = env['FOO'] })
+    StackProf::Middleware.save({ 'FOO' => 'bar' })
+    assert_equal env_set, 'bar'
+  end
+
+  def test_save_proc_should_receive_results_in_proc_if_passed
+    StackProf.stubs(:results).returns({ mode: 'foo' })
+
+    results_received = nil
+    StackProf::Middleware.new(Object.new, saviour: Proc.new{ |env, results| results_received = results[:mode] })
+    StackProf::Middleware.save({})
+    assert_equal results_received, 'foo'
+  end
+
   def test_enabled_should_use_a_proc_if_passed
     env = {}
 
@@ -64,4 +93,36 @@ def test_raw
     StackProf::Middleware.new(Object.new, raw: true)
     assert StackProf::Middleware.raw
   end
+
+  def test_enabled_should_override_mode_if_a_proc
+    proc_called = false
+    middleware = StackProf::Middleware.new(proc {|env| proc_called = true}, enabled: Proc.new{ [true, 'foo'] })
+    env = Hash.new { true }
+    enabled, mode = StackProf::Middleware.enabled?(env)
+    assert enabled
+    assert_equal 'foo', mode
+
+    StackProf.expects(:start).with({mode: 'foo', interval: StackProf::Middleware.interval, raw: false})
+    StackProf.expects(:stop)
+
+    middleware.call(env)
+    assert proc_called
+  end
+
+  def test_saviour_should_be_called_when_enabled_with_env
+    proc_called = false
+    env_set = nil
+    results_received = nil
+    enable_proc = Proc.new{ [true, 'foo'] }
+    saviour_proc = Proc.new{ |env, results| env_set = env['FOO'] ; results_received = results[:mode] }
+    middleware = StackProf::Middleware.new(proc {|env| proc_called = true}, enabled: enable_proc, saviour: saviour_proc, save_every: 1)
+    StackProf.expects(:start).with({mode: 'foo', interval: StackProf::Middleware.interval, raw: false})
+    StackProf.expects(:stop)
+    StackProf.stubs(:results).returns({ mode: 'foo' })
+
+    middleware.call({ 'FOO' => 'bar' })
+    assert proc_called
+    assert_equal env_set, 'bar'
+    assert_equal results_received, 'foo'
+  end
 end