cmd line util for benchmarking

gistfile1.rb

#!/usr/bin/env ruby

require 'rubygems'
require 'net/http'
require 'httpclient'
require 'eventmachine'
require 'em-http'
require 'optparse'

options = {}
opts = OptionParser.new do |opt|
  opt.on('-n', '--num NUM',
    'number of requests to make (default 1000)', Integer) do |val|
    options[:num] = val
  end

  opt.on('-b', '--bytes NUM',
    'byte size of server responses (default 100,000)', Integer) do |val|
    options[:bytes] = val
  end

  opt.on('-l', '--listen SERVER',
    'listen on HOST:PORT or PATH (default localhost:8080') do |val|
    options[:listen] = val
  end

  opt.on('-c', '--config-file FILE',
    'Rainbows! or Unicorn-specific config file') do |val|
    options[:sconf] = File.expand_path val
  end

  opt.on('--unicorn', 'use Unicorn instead of Rainbows!') do
    options[:server] = "unicorn"
  end

  opt.on('--ru FILE', 'use file as config.ru') do |val|
    options[:rack_conf] = File.read File.expand_path val
  end

  opt.on('-R', '--remote', "don't start a local server") do
    options[:remote] = true
  end
end

opts.parse! ARGV

options[:bytes] ||= 100000

TMPDIR = File.expand_path "~/tmp/http_bench.#{(Time.now.to_f * 1000).to_i}"
OUTPUT = []

DEFAULT_RACK_CONF = <<-STR
POOL = ("a".."z").to_a.concat \\
       ("A".."Z").to_a.concat \\
       ("0".."9").to_a.concat ["\\n", "-", "_"]

BODY = ""
#{options[:bytes]}.times{ BODY << POOL[rand(POOL.length)] }

SIZE = BODY.bytesize.to_s

run lambda{|env|
  [200, {"Content-Type" => "text/json", "Content-Length"=> SIZE}, [BODY]] }
STR

num    = options[:num]       || 1000
listen = options[:listen]    || "localhost:8080"
server = options[:server]    || "rainbows"
config = options[:rack_conf] || DEFAULT_RACK_CONF
uri    = URI.parse "http://#{listen}"
pid    = nil


def run_exit pid, status=0
  begin
    if pid && Process.kill(0, pid)
      # Bad bad hack because OSX doesn't seem to forward the sig to child
      Process.kill 'INT', (pid.to_i + 1)
      Process.waitpid2 pid
    end

    print "\n"
    puts OUTPUT.join("\n")
    `rm -r #{TMPDIR}`

  rescue => e
    puts e.message
    puts e.backtrace
    status = 1
  end

  exit status
end

trap 'INT' do
  run_exit pid, 1
end



begin

  unless options[:remote]
    `mkdir -p #{TMPDIR}`
    File.open("#{TMPDIR}/config.ru", "w+"){|file| file.write config }
    puts "starting #{server} at #{listen}"

    cmd = "cd #{TMPDIR} && #{server} -l #{listen}"
    cmd << " -c #{options[:sconf]}" if options[:sconf]

    #pid = fork{ exec "#{cmd} 2> #{TMPDIR}/stderr.log" }
    io = IO.popen "#{cmd} 2> #{TMPDIR}/stderr.log"
    pid = io.pid

    sleep 2

    puts "testing server\n"

    alive = Process.kill(0, pid)
    if alive
      resp  = Net::HTTP.get_response(uri)
      error = resp.code =~ %r{5\d\d}
    end

    unless alive && !error
      puts "server error\n\n"
      puts File.read("#{TMPDIR}/stderr.log")
      run_exit pid, 1
    end

    puts "server returned #{resp.body.bytesize} bytes"

    print "\n"
  end

  def timer
    s = Time.now
    yield
    Time.now - s
  end


  def benchmark num
    avg = 0
    i   = 0

    while i < num
      time = timer do
        yield
      end

      if time.to_i > 5 && i == 0
        puts "not recording first slow connection: #{time.to_f * 1000}"
        next
      end

      avg = ((avg * i) + time) / (i+1)
      i += 1

      $stdout << "."
      $stdout.flush
    end

    puts "\n"
    avg.to_f * 1000
  end



  class EMBench
    class << self
      attr_accessor :avg, :count
    end

    @avg    = 0
    @count  = 0


    def count
      self.class.count
    end


    def count= value
      self.class.count = value
    end


    def avg
      self.class.avg
    end


    def avg= value
      self.class.avg = value
    end


    def request client, head, body
      @start = Time.now
      [head, body]
    end


    def response resp
      time = Time.now - @start

      if time.to_i > 5 && self.count == 0
        puts "not recording first slow connection: #{time.to_f * 1000}"
        return resp
      end

      self.avg = ((self.avg * self.count) + time) / (self.count+1)
      self.count += 1

      $stdout << "."
      #$stdout << "#{time}\n"
      $stdout.flush

      resp
    end
  end



  def em_bench uri, num
    http = EM::HttpRequest.new(uri)
    http.use EMBench

    http.get.callback do
      if EMBench.count >= num
        EM.stop
      else
        em_bench uri, num
      end
    end
  end



  def bench_eventmachine uri, num
    puts "benchmarking EM #{num} times"

    EM::run do
      em_bench uri, num
    end
    puts "\n"

    EMBench.avg = EMBench.avg.to_f * 1000

    "eventmachine avg: #{EMBench.avg}"
  end




  def bench_nethttp uri, num
    puts "benchmarking net/http #{num} times"

    net_clnt = Net::HTTP.new(uri.host, uri.port).start
    get_req  = Net::HTTP::Post.new "/"

    net_res = benchmark num do
      net_clnt.request get_req
    end

    net_clnt.finish

    "net/http avg:     #{net_res}"
  end



  def bench_httpclient uri, num
    puts "benchmarking httpclient #{num} times"

    gem_clnt = HTTPClient.new

    gem_res = benchmark num do
      gem_clnt.post uri
    end

    "httpclient avg:   #{gem_res}"
  end


  OUTPUT << bench_eventmachine(uri, num)
  OUTPUT << bench_nethttp(uri, num)
  OUTPUT << bench_httpclient(uri, num)

  run_exit pid

rescue => e
  puts e.message
  puts e.backtrace

  run_exit pid, 1
end