benchmarking with memory reporting

mem_bench.rb

require 'benchmark' 

# Example:
#
# bm = MemBench.new
# bm.before_each {require File.dirname(__FILE__) + '/config/environment'}
# bm.benchmark("find") do
#   Product.find(:all, :limit => 1000, :include => [:variants =>
#                                                     :variant_descriptors])}
# end
#
# Output:
#
#           user     system      total        real
# find:  3.270000   0.120000   3.390000 (  3.779030)
# ---- Memory Usage
# before: 62.23Mb
# after:  108.92Mb
# diff:   46.70Mb
#
# Note: does not affect the main process memory size
class MemBench
  include Benchmark
  attr_accessor :before_each

  def benchmark(label, &block)
    fork do
      @before_each.call if @before_each
      memory_usage_before = nil
      memory_usage_after  = nil

      result = benchmark_to_string do |x|
        x.report(label) do 
          memory_usage_before = current_process_size
          yield
          memory_usage_after = current_process_size
        end
      end

      puts result
      puts "---- Memory Usage"
      puts "before: #{format(memory_usage_before)}\n"
      puts "after:  #{format(memory_usage_after)}\n"
      diff = memory_usage_after - memory_usage_before
      puts "diff:   #{format(diff)}\n"
      exit
    end
    Process.wait
  end

  def current_process_size
    `ps -o rss= -p #{Process.pid}`.to_i/1024.0
  end

  private

  def format(float)
    ("%0.02f" % float) + "Mb"
  end

  # benchmark usually goes to stdout. This is ok, but sometimes I like having it return
  # a string before sending it to stdout. Also, this is just a notch prettier.
  def benchmark_to_string
    reports = yield(Report.new(0, Benchmark::FMTSTR))
    reports = [reports].flatten
    largest_label = reports.max {|a,b| a.label.size <=> b.label.size }.label.size

    result = []
    result << "#{" " * largest_label}#{Benchmark::CAPTION}"
    reports.each do |report|
      label_correction = " " * (largest_label - report.label.size)
      result << (report.label + ":" + label_correction + report.to_s)
    end

    return result
  end

end