Redis Data Pump

redis_pump.rb

#!/usr/bin/ruby
require 'redis'
require 'securerandom'
require 'benchmark'

HOST = ENV['HOST'] || 'localhost'
PORT = ENV['PORT'] || '60812'
DB = 0
LIST_SIZE_PER_PROCESS=300000
NO_OF_PROCESS=5
TEST_TIMEOUT= ENV['TEST_TIMEOUT'] || 5

def pump(process_no, list_size = LIST_SIZE_PER_PROCESS)
  redis = Redis.new(host: HOST, port: PORT, db: DB)
  list_size.times.each_slice(2000) do |slice|
    redis.zadd('org_pump_data', slice.map { |score| [score, SecureRandom.hex] })
  end
  puts "Pump #{process_no} Completed - #{redis.zcard('org_pump_data')}"
end

def fork_process(n = NO_OF_PROCESS)
  puts "Pump forked with #{n} processes."
  Redis.new(host: HOST, port: PORT, db: DB).del('org_pump_data')
  n.times { |index| fork { pump(index) } }
end

def test
  puts 'Start test'
  puts 'Preparing data for test'
  copy_data_redis = Redis.new(host: HOST, port: PORT, db: DB, timeout: 30)
  copy_data_redis.del('pump_test')
  copy_data_redis.zunionstore('pump_test', ['org_pump_data'])
  puts "Test Data Size: #{copy_data_redis.zcard('pump_test')}, Source Size: #{copy_data_redis.zcard('org_pump_data')}"

  commands_executed = []
  Thread.new do
    monitor_redis = Redis.new(host: HOST, port: PORT, db: DB)
    puts 'monitor started'
    monitor_redis.monitor { |data| commands_executed << data }
  end

  puts 'Running test to remove the create data set within expected timeout,'
  sleep 2
  time_taken = Benchmark.measure do
    puts Redis.new(host: HOST, port: PORT, db: DB, timeout: TEST_TIMEOUT).zremrangebyscore('pump_test', 0, (Time.now + 1000000000000).to_i)
  end.real

  puts "Time taken: #{time_taken}, Test Data Size: #{copy_data_redis.zcard('pump_test')}"
  sleep 2
  puts 'Commands Executed In Redis:'  
  commands_executed.each_with_index { |command, index| puts "\t #{index}) #{command}"}
end

if ARGV[0] == 'test'
  test
elsif ARGV[0] == 'pump'
  fork_process
else
  puts 'usage: >> ruby redis_pump.rb test or >> ruby redis_pump.rb pump'
end