redis key distribution profiling script

profile.rb

# profile.rb - A script that analyzes the memory usage of a redis instance from it's "dump.rdb"
#
# To run this script, edit the key_groups variable in the stats method, and then run
# `ruby profile.rb dump.rdb` in the command line.
# Running this script will:
# * Dump the contents of the input rdb file in CSV format to "/tmp/memory.csv"
# * Load the contents of "/tmp/memory.csv" to the table "mem" in the "redis_mem" database
# * Print the stats of each key group to STDOUT
#
# Dependencies:
# * The python library "redis-rdb-tools" https://github.com/sripathikrishnan/redis-rdb-tools
# * The postgresql database, with the database "redis_mem" created
# * The "pg" and "sequel" gems

require 'sequel'
DB = Sequel.connect('postgres://localhost/redis_mem', max_connections: 20)

# Run the rdb tool, and dump the results to the "redis_mem" database under the table "mem"
def load_rdb rdb_file
  raise unless File.exists?(rdb_file)
  puts "Running rdb..."
  `rdb -c memory #{rdb_file} > /tmp/memory.csv`

  puts "Load /tmp/memory.csv to database"
  DB.run "drop table mem"
  DB.run "create table mem (database int, type varchar(32), key varchar(256), size_in_bytes int, encoding varchar(32), num_elements int, len_largest_element int);"
  `psql redis_mem -c "\\copy mem from /tmp/memory.csv with csv header"`
  
  num_rows = DB['select count(1) from mem'].all[0][:count]
  puts "Loaded #{num_rows} into database."
end

class KeyGroup
  attr_accessor :desc, :regex, :num_keys, :mem
  def initialize desc, regex
    @desc = desc; @regex = regex
  end
end

# Run queries on the database to compute the stats
def stats
  key_groups = [
    KeyGroup.new('User liked', 'user:\d+:liked_collages'),
    KeyGroup.new('Friends feed', 'fb_auth:\d+:friend_collages'),
    KeyGroup.new('Views', 'collage::views:v3'),
    KeyGroup.new('Series updated_at', 'cb:series:updated_at'),
    KeyGroup.new('Collage liked by', 'collage:\d+:liked_by'),
    KeyGroup.new('Followed feed', 'user:\d+:followed_collages'),
    KeyGroup.new('Collage number of likes', 'collage::likes:v3'),
    KeyGroup.new('Echoes', 'collage:\d+:echoes_originating'),
  ]
  
  threads = []
  key_groups.each do |kg|
    threads << Thread.new do
      res = DB["select sum(size_in_bytes), count(1) from mem where key similar to ?", kg.regex].all[0]
      kg.mem = res[:sum]; kg.num_keys = res[:count]
      puts "Completed #{kg.inspect}"
    end
  end
  threads.each{ |t| t.join }
  key_groups
end

def print_key_groups key_groups
  puts %|description,regex,"memory usage","number of keys"|
  key_groups.each do |kg|
    puts %|"#{kg.desc}","#{kg.regex}",#{kg.mem},#{kg.num_keys}|
  end
end

def run_main
  # ARGV[0] is expected to be the path of the rdb file
  load_rdb ARGV[0]
  key_groups = stats
  print_key_groups key_groups
end

if File.identical?(__FILE__, $0)
  run_main
end