[ruby] Ohai plugins benchmark script

ohai_bench.rb

#!/usr/bin/env ruby
# encoding: utf-8

# Ohai Plugin Benchmark
# =====================
#
# Usage: ohai_bench.rb [option] [loop_count]
#
# Options:
#     -d PLUGIN_DIR   Ohai plugin path
#     -r              Run rehearsal before executing
#
# License: MIT License (http://opensource.org/licenses/MIT)

require 'ohai'
require 'benchmark'
require 'optparse'

options = {
  :loop_count => 1,
  :plugin_paths => [],
  :rehearsal  => false,
}

opt = OptionParser.new
opt.banner += " [loop_count]\n\nOptions:\n"
opt.on('-d PLUGIN_DIR', 'Ohai plugin path') do |value|
  options[:plugin_paths] << value
end
opt.on('-r', 'Run rehearsal before executing') do |value|
  options[:rehearsal] = true
end
opt.parse!(ARGV)

if ARGV[0]
  options[:loop_count] = ARGV[0].to_i
end

$benchmark = {}

class Ohai::System
  attr_accessor :ohai_bench

  alias_method :orig_from_file, :from_file

  def from_file(filename)
    raise Errno.ENOENT unless File.exist?(filename)

    result = nil
    ohai_bench.benchmark(filename) {
      begin
        result = orig_from_file(filename)
      rescue Exception
      end
    }
    result
  end
end

class OhaiBenchmark
  def initialize(rehearsal = false)
    @rehearsal = rehearsal
    @tms = {}
  end

  def run(loop_count = 1)
    loop_count.times {
      ohai = Ohai::System.new
      ohai.ohai_bench = self
      ohai.all_plugins
    }
  end

  def benchmark(filename, &code_block)
    name = plugin_relpath(filename)
    tms = Benchmark.measure(name, &code_block)
    if @tms[name]
      @tms[name] = Benchmark::Tms.new(
        @tms[name].utime + tms.utime,
        @tms[name].stime + tms.stime,
        @tms[name].cutime + tms.cutime,
        @tms[name].cstime + tms.cstime,
        @tms[name].real + tms.real,
        tms.label
      )
    else
      @tms[name] = tms
    end
  end

  def plugin_relpath(filename)
    Ohai::Config['plugin_path'].each do |path|
      if filename.start_with?(path)
        return filename[path.length+1 .. -1]
      end
    end
    File.basename(filename)
  end

  def total
    @tms.inject(0) do |r, kv|
      r + kv[1].real
    end
  end

  def display
    labelwidth = @tms.keys.max_by {|x| x.length }.length
    line = '-' * (Benchmark::CAPTION.length + labelwidth)

    if @rehearsal
      ins = "Rehearsal "
      puts line.dup.tap {|o| o[0, ins.length] = ins }
    end

    print "%-#{labelwidth}s" % 'plugin'
    puts Benchmark::CAPTION

    @tms.sort {|a, b|
      b[1].real <=> a[1].real
    }.each do |name, tms|
      print "%-#{labelwidth}s" % tms.label
      puts tms
    end

    if @rehearsal
      ins = " total: %.6fsec" % total
      puts line.dup.tap {|o| o[-ins.length, line.length] = ins }
    end
  end
end

options[:plugin_paths].each do |path|
  Ohai::Config[:plugin_path] << File.expand_path(path)
end

if options[:rehearsal]
  bench = OhaiBenchmark.new(true)
  bench.run
  bench.display
  puts
end

bench = OhaiBenchmark.new
bench.run(options[:loop_count])
bench.display