#!/usr/bin/env ruby

require 'fileutils'

class Sidekiqctl
  DEFAULT_KILL_TIMEOUT = 10

  attr_reader :stage, :pidfile, :kill_timeout

  def self.print_usage
    puts "#{File.basename($0)} - stop a Sidekiq process from the command line."
    puts
    puts "Usage: #{File.basename($0)} <command> <pidfile> <kill_timeout>"
    puts " where <command> is either 'quiet' or 'stop'"
    puts "       <pidfile> is path to a pidfile"
    puts "       <kill_timeout> is number of seconds to wait until Sidekiq exits"
    puts "       (default: #{Sidekiqctl::DEFAULT_KILL_TIMEOUT}), after which Sidekiq will be KILL'd"
    puts
    puts "Be sure to set the kill_timeout LONGER than Sidekiq's -t timeout.  If you want"
    puts "to wait 60 seconds for jobs to finish, use `sidekiq -t 60` and `sidekiqctl stop"
    puts " path_to_pidfile 61`"
    puts
  end

  def initialize(stage, pidfile, timeout)
    @stage = stage
    @pidfile = pidfile
    @kill_timeout = timeout

    done('No pidfile given', :error) if !pidfile
    done("Pidfile #{pidfile} does not exist", :warn) if !File.exist?(pidfile)
    done('Invalid pidfile content', :error) if pid == 0

    fetch_process

    begin
      send(stage)
    rescue NoMethodError
      done "Invalid command: #{stage}", :error
    end
  end

  def fetch_process
    Process.kill(0, pid)
  rescue Errno::ESRCH
    done "Process doesn't exist", :error
  # We were not allowed to send a signal, but the process must have existed
  # when Process.kill() was called.
  rescue Errno::EPERM
    return pid
  end

  def done(msg, error = nil)
    puts msg
    exit(exit_signal(error))
  end

  def exit_signal(error)
    (error == :error) ? 1 : 0
  end

  def pid
    @pid ||= File.read(pidfile).to_i
  end

  def quiet
    `kill -TSTP #{pid}`
  end

  def stop
    `kill -TERM #{pid}`
    kill_timeout.times do
      begin
        Process.kill(0, pid)
      rescue Errno::ESRCH
        FileUtils.rm_f pidfile
        done 'Sidekiq shut down gracefully.'
      rescue Errno::EPERM
        done 'Not permitted to shut down Sidekiq.'
      end
      sleep 1
    end
    `kill -9 #{pid}`
    FileUtils.rm_f pidfile
    done 'Sidekiq shut down forcefully.'
  end
  alias_method :shutdown, :stop
end

if ARGV.length < 2
  Sidekiqctl.print_usage
else
  stage = ARGV[0]
  pidfile = ARGV[1]
  timeout = ARGV[2].to_i
  timeout = Sidekiqctl::DEFAULT_KILL_TIMEOUT if timeout == 0

  Sidekiqctl.new(stage, pidfile, timeout)
end