# Rake task to launch multiple Resque workers in development/production with simple management included

require 'resque/tasks' # Require Resque tasks

namespace :workers do

# = $ rake workers:start
#
# Launch multiple Resque workers with the Rails environment loaded,
# so they have access to your models, etc.
#
# Each worker is being run in their own separate process (_not_ thread).
#
# On clean shutdown (SIGINT / Ctrl-C, SIGQUIT, SIGTERM), the task will clean
# after itself: kill its workers and delete PID files, when appropriate. It
# will deal fine with already dead workers.
#
#
# Default options like COUNT can and should be over-ridden when invoking, of course:
#
#    $ rake workers:start COUNT=10 QUEUE=my_queue
#
#
# To daemonize, simply run with nohup:
#
#    $ nohup rake workers:start > log/workers.log 2>&1 &
#
#
# You can and should set up your monitoring tool to watch for process with PID
# from `cat tmp/pids/resque/master.pid`.
#
# For proper monitoring of _individual_ workers, use provided examples for God or Monit:
# http://github.com/defunkt/resque/blob/master/examples/.
#
#
# A task for killing all workers on the machine (`rake workers:killall`) is also provided,
# for pruning orphaned workers etc.
#
desc "Run and manage group of Resque workers with some default options"
task :start => :environment do

  # - CONFIGURATION ----
  ENV['QUEUE']   ||= '*'
  ENV['COUNT']   ||= '3'
  # --------------------

  def queue
    ENV['QUEUE']
  end

  def count
    ENV['COUNT']
  end

  def queue_pathname
    queue.gsub(/[:\*<>\|\"\/\\\?]/, "_")
  end

  def Process.exists?(pid)
    kill(0, pid.to_i) rescue false
  end

  def pid_directory
    @pid_directory ||= Rails.root.join('tmp', 'pids', "resque")
  end

  def pid_directory_for_group
    @pid_directory_for_group ||= Rails.root.join('tmp', 'pids', "resque", queue_pathname)
  end

  def group_master_pid
    File.read pid_directory.join("#{queue_pathname}.pid").to_s rescue nil
  end

  def group?
    !group_master_pid || group_master_pid.to_s == Process.pid.to_s
  end

  def group_running?
    Process.exists?(group_master_pid)
  end

  def kill_worker(pid)
    %w{QUIT TERM KILL}.each do |signal|
      if Signal.list.has_key?(signal)
        begin
          Process.kill(signal, pid)
        rescue Errno::EINVAL
          next
        end

        puts "Killed worker with PID #{pid}"
        break
      end
    end
  rescue Errno::ESRCH => e
    puts " STALE worker with PID #{pid}"
  end

  def kill_workers
    @pids.each { |pid| kill_worker(pid) }
  end

  def kill_workers_and_remove_pids_for_group
    Dir.glob(pid_directory_for_group.join('worker_*.pid').to_s) do |pidfile|
      begin
        pid = pidfile[/(\d+)\.pid/, 1].to_i
        kill_worker(pid)
      ensure
        FileUtils.rm pidfile, :force => true
      end
    end
    if group_master_pid
      FileUtils.rm    pid_directory.join("#{queue_pathname}.pid").to_s
      FileUtils.rm_rf pid_directory_for_group.to_s
    end
  end

  def shutdown
    puts "\n*** Exiting"
    if group?
      kill_workers_and_remove_pids_for_group
    else
      kill_workers
    end
    exit(0)
  end

  # Clean up after dead group from before -- and become one
  unless group_running?
    puts "--- Cleaning up after previous group (PID: #{group_master_pid})"
    kill_workers_and_remove_pids_for_group 
  end

  # Handle exit
  trap('INT')  { shutdown }

  begin
    trap('QUIT') { shutdown }
  rescue ArgumentError
    warn "SIGQUIT not supported"
  end

  trap('TERM') { shutdown }
  trap('KILL') { shutdown }
  trap('SIGKILL') { shutdown }

  puts "=== Launching #{ENV['COUNT']} worker(s) on '#{ENV['QUEUE']}' queue(s) with PID #{Process.pid}"

  # Launch workers in separate processes, saving their PIDs
  @pids = []
  ENV['COUNT'].to_i.times do
    if Process.respond_to?(:fork)
      @pids << Process.fork { Rake::Task['resque:work'].invoke }
    else
      @pids << Process.spawn("rake environment resque:work", chdir: Rails.root)
    end
  end

  if group?
    # Make sure we have directory for pids
    FileUtils.mkdir_p pid_directory.to_s
    FileUtils.mkdir_p pid_directory_for_group.to_s
    # Create PID files for workers
    File.open( pid_directory.join("#{queue_pathname}.pid").to_s, 'w' ) do |f| f.write Process.pid end
    @pids.each do |pid|
      File.open( pid_directory_for_group.join("worker_#{pid}.pid").to_s, 'w' ) { |f| f.write pid }
    end
    # Stay in foreground, if any of our workers dies, we'll get killed so Monit/God etc can come to the resq^Hcue
    Process.wait
  else
    # Stay in foreground, if any of our workers dies, continue running
    Process.waitall
  end
end

desc "Kill ALL workers on this machine"
task :kilall do
  require 'resque'
  Resque::Worker.all.each do |worker|
    puts "Shutting down worker #{worker}"
    host, pid, queues = worker.id.split(':')
    kill_worker(pid.to_i)
  end
end

end