traffic-creator

#!/usr/bin/env ruby
#
# Author:: Daniel DeLeo (<dan@opscode.com>)
# Author:: Seth Falcon (<seth@opscode.com>)
# Author:: Chris Walters (<cw@opscode.com>)
# Copyright:: Copyright (c) 2010-2011 Opscode, Inc.
# License:: Apache License, Version 2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

require "rubygems"

$:.unshift(File.expand_path(File.dirname(__FILE__) + '/../lib'))

require 'pp'
require 'bunny'
require 'yajl'
require 'uuidtools'
require 'word_salad'

require 'chef_search/expander/configuration'

Chef::Expander.init_config(ARGV)

MESSAGES_TO_SEND = 10_000

NUM_RAND_KEY_PAIRS = 50
NUM_RAND_VALUE_PAIRS = 50

PERSISTENT_MESSAGES = true

KEYS = NUM_RAND_VALUE_PAIRS.words

SAMPLE_NODES = []
Dir.glob(File.expand_path(File.dirname(__FILE__)) + '/../data/*_node.json') do |node_file|
  SAMPLE_NODES << Yajl::Parser.parse(IO.read(node_file))
end

NUM_NODES = SAMPLE_NODES.size
puts "Read #{NUM_NODES} sample nodes"

puts "Using rabbitmq config #{Chef::Expander.config.amqp_config.inspect}"

puts "connecting to rabbitmq"
amqp_client = Bunny.new(Chef::Expander.config.amqp_config)
amqp_client.start

puts 'declaring queues'
queues = {}
0.upto(1023) do |vnode|
  queues[vnode] = amqp_client.queue("vnode-#{vnode}", :durable => true)
end
  
def add_rand_keys(node)
  rand_key_vals = Hash[*((2 * NUM_RAND_KEY_PAIRS).words)]
  rand_vals = Hash[*(KEYS.zip(NUM_RAND_VALUE_PAIRS.words)).flatten]
  node.merge(rand_key_vals.merge(rand_vals))
end

puts "sending #{MESSAGES_TO_SEND} messages"
start_time = Time.now
sent_messages = 0
1.upto(MESSAGES_TO_SEND) do
  node = SAMPLE_NODES[rand(NUM_NODES)]
  node = add_rand_keys(node)
  index_data = {:action => :add}
  index_data[:payload] = {:item => node}
  index_data[:payload][:type] = :node
  index_data[:payload][:database] = :testdb
  index_data[:payload][:enqueued_at] = Time.now.utc.to_i

  id = node["name"]
  vnode = rand(1024)
  index_data[:payload][:id] = id

  puts "queue: vnode-#{vnode} (#{sent_messages} / #{MESSAGES_TO_SEND})"
  amqp_client.tx_select if PERSISTENT_MESSAGES
  queues[vnode].publish(Yajl::Encoder.encode(index_data), :persistent => PERSISTENT_MESSAGES)
  amqp_client.tx_commit if PERSISTENT_MESSAGES
  sent_messages += 1
end
end_time = Time.now

total_time = end_time - start_time
rate = MESSAGES_TO_SEND.to_f / total_time
puts "done (#{total_time}s, #{rate} msg/s)"