# Copyright 2010
# Beyond Perspective Solutions, Inc.
# http://www.bpsinc.jp/

# Function1: Find out Unicorn process ID(master/children)
# Function2: Get memory usage of each unicorn process(master/children)
# Function3: Send USR2 signal to unicorn master for reloading code
# Function4: Send QUIT signal to unicorn child for restart
module UnicornManager

  # get unicorn_pids to 
  def self.get_unicorn_pids(rails_root)
    unicorn_master_pid_file = File.expand_path("tmp/pids/unicorn.pid", rails_root)
    unicorn_master_pid = File.read(unicorn_master_pid_file).to_i

    # works on only Linux ??
    ps_output = `ps w --ppid #{unicorn_master_pid}`
    ps_output_lines = ps_output.split("\n")

    children = []
    ps_output_lines.each do |line|
      chunks = line.strip.split(/\s+/, 5)
      pid, pcmd = chunks[0], chunks[4]
      next if pid !~ /\A\d+\z/ or pcmd !~ /worker/
      children << pid.to_i
    end

    return { :master => unicorn_master_pid, :children => children }
  end

  # send HUP to master
  def self.reload_config(rails_root)
    pids = get_unicorn_pids(rails_root)
    Process.kill :HUP, pids[:master]
  end

  # send USR2 to original master
  def self.reload_code(rails_root)
    # find old master
    pids = get_unicorn_pids(rails_root)
    Process.kill :USR2, pids[:master]
  end

  # send QUIT to all children
  def self.restart_all_child(rails_root)
    pids = get_unicorn_pids(rails_root)
    pids[:children].each do |child|
      restart_child child
    end
  end

  # send QUIT to child
  def self.restart_child(child_pid)
    Process.kill :QUIT, child_pid.to_i
  end

  # get memory usage
  # returns: { :master => {123 => 1000000}, :children => { 124 => ..., 125 => ..., ... } }
  def self.get_memories(rails_root)
    pids = get_unicorn_pids(rails_root)
    mem = { :master => {pids[:master] => nil}, :children => {} }
    ps_output = `ps axuw`
    ps_output_lines = ps_output.split("\n")
    ps_output_lines.each do |line|
      chunks = line.strip.split(/\s+/, 11)
      pid, pmem_rss, pcmd = chunks.values_at(1, 5, 10)
      pmem = pmem_rss.to_i * 1024
      pid = pid.to_i

      if pid == pids[:master] then
        mem[:master][pid] = pmem
      elsif pids[:children].include?(pid) then
        mem[:children][pid] = pmem
      end
    end
    
    return mem
  end

end

if $0 == __FILE__ then
  require "pp"
  mf_rails_root = "/var/www/music-fly.net/webservice2"

  # Sample1: get process IDs and dump
  #pp UnicornManager::get_unicorn_pids(mf_rails_root)

  # Sample2: reload code (master restart)
  #UnicornManager::reload_code(mf_rails_root)

  # Sample3: get memory usage and make fat process restart 
  memory = UnicornManager::get_memories(mf_rails_root)
  pp memory

  memory[:children].each do |pid, mem_size|
    if 1024*1024*30 < mem_size then
      puts "sending QUIT to child #{pid}"
      Process.kill :QUIT, pid
    end
  end

  # check if the fat processes are certainly restarted
  puts "waiting..."
  sleep 3
  memory = UnicornManager::get_memories(mf_rails_root)
  pp memory

end