gistfile1.rb

#!/usr/bin/env ruby

require 'trollop'
require 'em-http'
require 'sha1'
require 'ruby-debug'


# Event driven Dynamo Server Benchmarker
class DynamoBenchmarker

  def initialize(options = {})
    options.each_pair do |k,v|
      self.instance_variable_set("@#{k}".to_sym,v)
    end 
    @successful = 0
    @failed = 0
    @in_progress = 0
    @total_attempted_requests = 0
    @disconnected = 0
  end

  def run
    if @dry_run
      puts "Beginning dynamo-server spam dry run"
    else
      puts "Beginning dynamo-server spam"
    end

    @clients = (0..3).map{random_16_bit_hex}
    @entities = (0..1).map{random_16_bit_hex}
    @attributes = ["phone", "credit_card"]
    @campaigns = (0..7).map{random_16_bit_hex} 
    @medias = (0..7).map{random_16_bit_hex}

    EventMachine.run do
      EventMachine.add_periodic_timer(0.01) { post_attribute }
    end
  end

  private

  def generate_post_params
    [@clients[rand(@clients.size)], @entities[rand(@entities.size)], @attributes[rand(@attributes.size)], @campaigns[rand(@campaigns.size)], @medias[rand(@medias.size)]]
  end
  
  def post_attribute
    if @total_attempted_requests < @num_requests && @in_progress < @concurrency
      client, entity, attribute, campaign, media = *generate_post_params
      uri = "/v1.0/client/%s/entity/%s/attribute/%s?apikey=%s&campaignId=%s&mediaId=%s" % [client, entity, attribute, @api_key, campaign, media]
      path = "http://#{@host}:#{@port}#{uri}"
      puts "POST #{path}"
      http = EventMachine::HttpRequest.new(path).post :body => "phone=iPhone"
      @in_progress += 1
      @total_attempted_requests += 1
      http.callback do
        @in_progress -= 1
        @successful += 1
        report if finished?
      end
      http.errback do
        @in_progress -= 1
        @failed += 1
        report if finished?
      end
    end
  end

  def finished?
    @total_attempted_requests >= @num_requests
  end

  def report
    EventMachine.stop
    puts "
Made a total of #{@total_attempted_requests} out of #{@num_requests} asked for.
#{@successful} requests succeeded.
#{@failed} requests failed.
    "
    exit 0
  end

  def random_16_bit_hex
    srand
    seed = "#{rand(1000000)}#{Time.now}"
    Digest::SHA1.hexdigest(seed)[0..15]
  end
end

if __FILE__ == $0
  opts = Trollop::options do
    opt :host, "Dynamo-server host", :type => :string
    opt :port, "Dynamo-server port", :type => :int, :default => 80
    opt :api_key, "Valid API key for dynamo-server, should be 16-char hex and registered on the server", :type => :string
    opt :num_requests, "Number of requests to send to dynamo-server", :default => 15, :type => :int
    opt :concurrency, "Number of concurrent requests to send to dynamo-server", :default => 20, :type => :int
    opt :dry_run, "Don't make the requests to dynamo-server, just output the URIs that would be hit", :default => false
  end

  #Trollop::die :base_uri, "Must be valid URI" if opts[:base_uri].nil?
  #Trollop::die :api_key, "Must specify API key" if opts[:api_key].nil?

  dynamo = DynamoBenchmarker.new(opts)
  dynamo.run
end