#!/usr/bin/ruby
# restart-unicorn: Graceful restart for Unicorn
#    depends on Linux's proc(5)
#
# MIT License: Copyright(c)2011 MATSUYAMA Kengo

require 'scanf'
require 'timeout'

class ProcFS
  PROC_DIR = '/proc'

  # cf. proc(5)
  PROC_STAT_FORMAT = [
    %w[pid %d],
    %w[comm %s],
    %w[state %c],
    %w[ppid %d],
    %w[pgrp %d],
    %w[session %d],
  ]

  STAT_PARAMS = PROC_STAT_FORMAT.map(&:first)
  STAT_FORMAT = PROC_STAT_FORMAT.map(&:last).join(' ')

  Stat = Struct.new(*STAT_PARAMS.map(&:to_sym))

  attr_reader :pid

  STAT_PARAMS.each do |name|
    next if method_defined?(name)
    define_method(name) { stat && stat.send(name) }
  end


  def self.all
    Dir.entries(PROC_DIR).grep(/\A\d+\z/).map {|pid| ProcFS.new(pid)}
  end

  def self.detect(options = {}, &block)
    if timeout = options[:timeout]
      Timeout.timeout(timeout) do
        all.detect(&block) or (sleep 0.2 and redo)
      end
    else
      all.detect(&block)
    end
  end

  def initialize(pid)
    @pid = pid.to_i
  end

  def stat
    @stat ||= Stat.new(*read('stat').scanf(STAT_FORMAT))
  rescue
    nil
  end

  def cmdline
    @cmdline ||= read('cmdline').gsub(/\0/, ' ').strip
  rescue
    nil
  end

  private

  def read(name)
    File.read(File.join(PROC_DIR, @pid.to_s, name))
  end
end


if ARGV.size != 1
  puts "USAGE: #{File.basename($0)} PIDFILE"
  exit
end

pid_file = ARGV.shift

shutdown_old_master_process = proc do
  old_pid_file = pid_file + '.oldbin'
  if File.exist?(old_pid_file)
    Process.kill(:QUIT, File.read(old_pid_file).to_i)
    File.unlink(old_pid_file)
  end
end

# shutdown an old master process. (if exists)
shutdown_old_master_process.call

# find a current master process.
pid = File.read(pid_file).to_i
master_proccess = ProcFS.new(pid)

unless master_proccess.stat
  STDERR.puts "Master process #{pid} is missing"
  exit 1
end

# send SIGUSR2 to the master.
Process.kill(:USR2, pid)

begin
  # find a new master.
  new_master_process = ProcFS.detect(:timeout => 5) do |process|
    process.ppid == pid && process.cmdline =~ /\Aunicorn(?:_rails)? master/
  end
rescue Exception
  # can't find a new master process.
  STDERR.puts "#{$!.class}: #{$!.to_s}"
  STDERR.puts "Can't find a new master process"
  exit 1
end

begin
  # find a new worker process.
  new_worker_process = ProcFS.detect(:timeout => 5) do |process|
    process.ppid == new_master_process.pid && process.cmdline =~ /\Aunicorn(?:_rails)? worker/
  end

  # wait for the worker to sleep.
  timeout(60) do
    count = 2
    begin
      sleep 0.2
      stat = ProcFS.new(new_worker_process.pid).stat or raise "New worker process #{new_worker_process.pid} was died"
    end until stat.state == 'S' && (count -= 1).zero?
  end

  # everything is ok. kill the old master process.
  shutdown_old_master_process.call
rescue Exception
  # something was failed. kill the new master process.
  STDERR.puts "#{$!.class}: #{$!.to_s}"
  STDERR.puts "Kill the new master process #{new_master_process.pid}"
  Process.kill(:QUIT, new_master_process.pid)
  exit 1
end