rampup-client

#!/usr/bin/env ruby

require 'rubygems'
require 'digest/md5'
require 'net/http'
require 'uri'

require './rampup-store'

$cfg_defaults = {
  :prefix => "",            # Prefix for every item key.
  :max_ops => 0,            # Max number of requests (sets and gets) before exiting.
  :max_items => -1,         # Max number of items, if we ever create any; default 10000.
  :max_creates => -1,       # Max number of creates; defaults to max_items.
  :min_value_size => 10,    # Minimal value size in bytes during SET's.
  :ratio_sets => 0.1,       # Fraction of requests that should be SET's.
  :ratio_creates => 0.1,    # Fraction of SET's that should create new items;
                            # the rest are updates.
  :ratio_deletes => 0.0,    # Fraction of updates that should instead be DELETE's;
                            # the rest are modifications. (TODO)
  :ratio_advanced => 0.0,   # Fraction of modifications that should be fancy
                            # mutation commands instead of a simple SET;
                            # like append, prepend, etc. (TODO)
  :ratio_misses => 0.05,    # Fraction of GET's that should miss.
  :ratio_hot => 0.2,        # Fraction of items to have as a hot item subset.
  :ratio_hot_sets => 0.95,  # Fraction of SET's that hit the hot item subset.
  :ratio_hot_gets => 0.95,  # Fraction of GET's that hit the hot item subset.
  :bulk_load_batch => 1,    # When > 1, use batched bulk load optimization,
                            # if appropriate, flushing after this many inserts.
  :exit_after_creates => 0, # Exit after max_creates is reached.
  :num_vbuckets => 0,
  :json => 1                # Use JSON documents
}

$cur_defaults = {
  :num_items    => 0, # Number of items known to be in the bucket.
  :num_sets     => 0,
  :num_creates  => 0, # Number of sets that were creates.
  :num_deletes  => 0, # Number of sets that were deletes.
  :num_advanced => 0, # Number of sets that were advanced mutation commands.
  :num_gets     => 0
}

if ARGV.length < 2
  print("usage: rampup-client <couchdb|mongo|memcached> host:port [key=val[ key=val[ ...]]]\n")
  print("\n")
  print("optional key=val's and their defaults are...\n")
  ($cfg_defaults.keys + $cur_defaults.keys).map {|x| x.to_s}.sort.each do |k|
    print("  #{k.ljust(20)} = #{$cfg_defaults[k.to_sym] || $cur_defaults[k.to_sym]}\n")
  end
  print("\n")
  exit(-1)
end

# ----------------------------------------------------------------

def out(msg)
  print(msg)
end

# ----------------------------------------------------------------

class Cfg
  $cfg_defaults.each_pair {|k, v| attr_accessor(k) }
end

class Cur
  $cur_defaults.each_pair {|k, v| attr_accessor(k) }
end

# ----------------------------------------------------------------

cfg = Cfg.new
cur = Cur.new

[[cfg, $cfg_defaults],
 [cur, $cur_defaults]].each do |obj, defaults|
  defaults.each_pair {|k, v| obj.send((k.to_s + '=').to_sym, v) }
  (ARGV[2..-1] || []).each do |kv|
    k, v = kv.split('=')
    unless defaults[k.to_sym].nil?
      v = v.to_f if v == v.to_f.to_s
      v = v.to_i if v == v.to_i.to_s
      obj.send((k + '=').to_sym, v)
    end
  end
end

cfg.max_items = 10000 if cfg.max_items < 0 and cfg.max_creates < 0
cfg.max_items = cfg.max_creates if cfg.max_items < 0
cfg.max_creates = cfg.max_items if cfg.max_creates < 0

if cfg.ratio_sets < 1.0 or
   cfg.ratio_creates < 1.0
  cfg.bulk_load_batch = 1 # Use bulk_load_batch optimization only if 100% creates.
end

($cfg_defaults.keys).map {|x| x.to_s}.sort.each do |k|
  print("  #{k.ljust(20)} = #{cfg.send(k.to_sym)}\n")
end

($cur_defaults.keys).map {|x| x.to_s}.sort.each do |k|
  print("  #{k.ljust(20)} = #{cur.send(k.to_sym)}\n")
end

print("\n")

# ----------------------------------------------------------------

$body = 'x';
while $body.length < cfg.min_value_size
  $body = $body + Digest::MD5.hexdigest($body.length.to_s)
end

$suffix = "\"body\":\"#{$body}\"}"

def gen_doc_hash(key_num, key_str, min_value_size)
  return { "_id"     => key_str,
           "key_num" => key_num,
           "mid"     => key_str[-8..-2],
           "last"    => key_str[-1..-1],
           "body"    => $body
         }
end

def gen_doc_string(key_num, key_str, min_value_size, key_name = "key",
                   json = true)
  c = '{'
  c = '*' unless json
  return "#{c}\"#{key_name}\":\"#{key_str}\"," \
          "\"key_num\":#{key_num}," \
          "\"mid\":\"#{key_str[-8..-2]}\"," \
          "\"last\":\"#{key_str[-1..-1]}\"," +
         $suffix
end

# ----------------------------------------------------------------

def run(cfg, cur, store)
  while true
    num_ops = cur.num_gets + cur.num_sets

    return "max_ops-reached" if (cfg.max_ops > 0 and
                                 cfg.max_ops <= num_ops)

    return "max_creates-reached" if (cfg.exit_after_creates > 0 and
                                     cfg.max_creates > 0 and
                                     cfg.max_creates <= cur.num_creates)

    req_cmd, req_key_num, req_key_str, req_data,
      expect_status, expect_data =
        next_cmd(cfg, cur, store)

    cmd_beg = Time.new
    store.command(req_cmd, req_key_num, req_key_str, req_data, cfg.bulk_load_batch)
    cmd_end = Time.new
    cmd_amt = cmd_end - cmd_beg
  end

  cmd_beg = Time.new
  store.flush()
  cmd_end = Time.new
  cmd_amt = cmd_end - cmd_beg
end

def positive(x)
  return 1 if x <= 0
  return x
end

def next_cmd(cfg, cur, store)
  # This function modifies cur.
  #
  num_ops = cur.num_gets + cur.num_sets

  do_set = cfg.ratio_sets > cur.num_sets.to_f / positive(num_ops)
  if do_set
    cmd = :set
    cur.num_sets += 1

    do_set_create = (cfg.max_items > cur.num_items and
                     cfg.max_creates > cur.num_creates and
                     cfg.ratio_creates > cur.num_creates.to_f / positive(cur.num_sets))
    if do_set_create
      # Create...
      key_num = cur.num_items

      cur.num_items += 1
      cur.num_creates += 1
    else
      # Update...
      key_num = choose_key_num(cur.num_items, cfg.ratio_hot, cfg.ratio_hot_sets, cur.num_sets)

      num_updates = cur.num_sets - cur.num_creates

      do_delete = cfg.ratio_deletes > cur.num_deletes.to_f / positive(num_updates)
      if do_delete
        cur.num_deletes += 1
        cmd = :delete
      else
        num_mutations = num_updates - cur.num_deletes

        do_advanced = cfg.ratio_advanced > cur.num_advanced.to_f / positive(num_mutations)
        if do_advanced
          cur.num_advanced += 1
          cmd = :advanced
        end
      end
    end

    key_str = prepare_key(key_num, cfg.prefix)
    itm_val = store.gen_doc(key_num, key_str, cfg.min_value_size)

    return cmd, key_num, key_str, itm_val,
           :success, nil
  else
    cmd = :get
    cur.num_gets += 1

    do_get_hit = (cfg.ratio_misses * 100) < cur.num_gets.modulo(100)
    if do_get_hit
      key_num = choose_key_num(cur.num_items, cfg.ratio_hot, cfg.ratio_hot_gets, cur.num_gets)
      key_str = prepare_key(key_num, cfg.prefix)
      itm_val = store.gen_doc(key_num, key_str, cfg.min_value_size)

      return cmd, key_num, key_str, nil,
             :success, itm_val
    else
      return cmd, -1, prepare_key(-1, cfg.prefix), nil,
             :missing, "Not found"
    end
  end
end

# ----------------------------------------------------------------

def choose_key_num(num_items, ratio_hot, ratio_hot_choice, num_ops)
  hit_hot_range = (ratio_hot_choice * 100) > num_ops.modulo(100)
  if hit_hot_range
    base  = 0
    range = (ratio_hot * num_items).floor
  else
    base  = (ratio_hot * num_items).floor
    range = ((1.0 - ratio_hot) * num_items).floor
  end
  return base + num_ops.modulo(positive(range))
end

def prepare_key(key_num, prefix)
  key_hash = Digest::MD5.hexdigest(key_num.to_s)[0..15]

  return "#{prefix}-#{key_hash}" if prefix and prefix.length > 0
  return "#{key_hash}"
end

# ----------------------------------------------------------------

case ARGV[0].downcase
when 'none'
  store = Store.new
when 'couchdb'
  store = StoreCouchDB.new
when 'mongo'
  store = StoreMongo.new
when 'memcached-memcache'
  store = StoreMemCache.new
else
  store = StoreDalli.new
end

store.connect(ARGV[1], cfg)

# ----------------------------------------------------------------

print("First 5 keys...\n")
(0..5).each do |x|
  print(store.cmd_line_get(x, prepare_key(x, cfg.prefix)) + "\n")
end

print("\n")

# ----------------------------------------------------------------

run(cfg, cur, store)