resque-pool-diet 0.3.1

data/features/support/aruba_daemon_support.rb

@@ -0,0 +1,76 @@
+require 'aruba/cucumber'
+require 'aruba/api'
+require 'aruba/process'
+
+module Aruba
+
+  module Api
+
+    # this is a horrible hack, to make sure that it's done what it needs to do
+    # before we do our next step
+    def keep_trying(timeout=10, tries=0)
+      puts "Try: #{tries}" if @announce_env
+      yield
+    rescue RSpec::Expectations::ExpectationNotMetError
+      if tries < timeout
+        sleep 1
+        tries += 1
+        retry
+      else
+        raise
+      end
+    end
+
+    def run_background(cmd)
+      @background = run(cmd)
+    end
+
+    def send_signal(cmd, signal)
+      announce_or_puts "$ kill -#{signal} #{processes[cmd].pid}" if @announce_env
+      processes[cmd].send_signal signal
+    end
+
+    def background_pid
+      @pid_from_pidfile || @background.pid
+    end
+
+    # like all_stdout, but doesn't stop processes first
+    def interactive_stdout
+      only_processes.inject("") { |out, ps| out << ps.stdout(@aruba_keep_ansi) }
+    end
+
+    # like all_stderr, but doesn't stop processes first
+    def interactive_stderr
+      only_processes.inject("") { |out, ps| out << ps.stderr(@aruba_keep_ansi) }
+    end
+
+    # like all_output, but doesn't stop processes first
+    def interactive_output
+      interactive_stdout << interactive_stderr
+    end
+
+    def interpolate_background_pid(string)
+      interpolated = string.gsub('$PID', background_pid.to_s)
+      announce_or_puts interpolated if @announce_env
+      interpolated
+    end
+
+    def kill_all_processes!
+    #  stop_processes!
+    #rescue
+    #  processes.each {|cmd,process| send_signal(cmd, 'KILL') }
+    #  raise
+    end
+
+  end
+
+  class Process
+    def pid
+      @process.pid
+    end
+    def send_signal signal
+      @process.send :send_signal, signal
+    end
+  end
+
+end

data/features/support/env.rb

@@ -0,0 +1 @@
+ENV["RAILS_ENV"] = "test"

data/lib/resque/pool/cli.rb

@@ -0,0 +1,120 @@
+require 'trollop'
+require 'resque/pool'
+require 'fileutils'
+
+module Resque
+  class Pool
+    module CLI
+      extend self
+
+      def run
+        opts = parse_options
+        daemonize if opts[:daemon]
+        manage_pidfile opts[:pidfile]
+        redirect opts
+        setup_environment opts
+        start_pool
+      end
+
+      def parse_options
+        opts = Trollop::options do
+          version "resque-pool #{VERSION} (c) nicholas a. evans"
+          banner <<-EOS
+resque-pool is the best way to manage a group (pool) of resque workers
+
+When daemonized, stdout and stderr default to resque-pool.stdxxx.log files in
+the log directory and pidfile defaults to resque-pool.pid in the current dir.
+
+Usage:
+   resque-pool [options]
+where [options] are:
+          EOS
+          opt :config, "Alternate path to config file", :type => String, :short => "-c"
+          opt :appname, "Alternate appname",         :type => String,    :short => "-a"
+          opt :daemon, "Run as a background daemon", :default => false,  :short => "-d"
+          opt :stdout, "Redirect stdout to logfile", :type => String,    :short => '-o'
+          opt :stderr, "Redirect stderr to logfile", :type => String,    :short => '-e'
+          opt :nosync, "Don't sync logfiles on every write"
+          opt :pidfile, "PID file location",         :type => String,    :short => "-p"
+          opt :environment, "Set RAILS_ENV/RACK_ENV/RESQUE_ENV", :type => String, :short => "-E"
+        end
+        if opts[:daemon]
+          opts[:stdout]  ||= "log/resque-pool.stdout.log"
+          opts[:stderr]  ||= "log/resque-pool.stderr.log"
+          opts[:pidfile] ||= "tmp/pids/resque-pool.pid"
+        end
+        opts
+      end
+
+      def daemonize
+        raise 'First fork failed' if (pid = fork) == -1
+        exit unless pid.nil?
+        Process.setsid
+        raise 'Second fork failed' if (pid = fork) == -1
+        exit unless pid.nil?
+      end
+
+      def manage_pidfile(pidfile)
+        return unless pidfile
+        pid = Process.pid
+        if File.exist? pidfile
+          if process_still_running? pidfile
+            raise "Pidfile already exists at #{pidfile} and process is still running."
+          else
+            File.delete pidfile
+          end
+        else
+          FileUtils.mkdir_p File.dirname(pidfile)
+        end
+        File.open pidfile, "w" do |f|
+          f.write pid
+        end
+        at_exit do
+          if Process.pid == pid
+            File.delete pidfile
+          end
+        end
+      end
+
+      def process_still_running?(pidfile)
+        old_pid = open(pidfile).read.strip.to_i
+        Process.kill 0, old_pid
+        true
+      rescue Errno::ESRCH
+        false
+      rescue Errno::EPERM
+        true
+      rescue ::Exception => e
+        $stderr.puts "While checking if PID #{old_pid} is running, unexpected #{e.class}: #{e}"
+        true
+      end
+
+      def redirect(opts)
+        $stdin.reopen  '/dev/null'        if opts[:daemon]
+        # need to reopen as File, or else Resque::Pool::Logging.reopen_logs! won't work
+        out = File.new(opts[:stdout], "a") if opts[:stdout] && !opts[:stdout].empty?
+        err = File.new(opts[:stderr], "a") if opts[:stderr] && !opts[:stderr].empty?
+        $stdout.reopen out if out
+        $stderr.reopen err if err
+        $stdout.sync = $stderr.sync = true unless opts[:nosync]
+      end
+
+      def setup_environment(opts)
+        Resque::Pool.app_name = opts[:appname]    if opts[:appname]
+        ENV["RACK_ENV"] = ENV["RAILS_ENV"] = ENV["RESQUE_ENV"] = opts[:environment] if opts[:environment]
+        Resque::Pool.log "Resque Pool running in #{ENV["RAILS_ENV"] || "development"} environment"
+        ENV["RESQUE_POOL_CONFIG"] = opts[:config] if opts[:config]
+      end
+
+      def start_pool
+        require 'rake'
+        require 'resque/pool/tasks'
+        Rake.application.init
+        Rake.application.load_rakefile
+        Rake.application["resque:pool"].invoke
+      end
+
+    end
+  end
+end
+

data/lib/resque/pool/logging.rb

@@ -0,0 +1,65 @@
+module Resque
+  class Pool
+    module Logging
+      extend self
+
+      # more than a little bit complicated...
+      # copied this from Unicorn.
+      def self.reopen_logs!
+        log "Flushing logs"
+        [$stdout, $stderr].each do |fd|
+          if fd.instance_of? File
+            # skip if the file is the exact same inode and device
+            orig_st = fd.stat
+            begin
+              cur_st = File.stat(fd.path)
+              next if orig_st.ino == cur_st.ino && orig_st.dev == cur_st.dev
+            rescue Errno::ENOENT
+            end
+            # match up the encoding
+            open_arg = 'a'
+            if fd.respond_to?(:external_encoding) && enc = fd.external_encoding
+              open_arg << ":#{enc.to_s}"
+              enc = fd.internal_encoding and open_arg << ":#{enc.to_s}"
+            end
+            # match up buffering (does reopen reset this?)
+            sync = fd.sync
+            # sync to disk
+            fd.fsync
+            # reopen, and set ruby buffering appropriately
+            fd.reopen fd.path, open_arg
+            fd.sync = sync
+            log "Reopened logfile: #{fd.path}"
+          end
+        end
+      end
+
+      # Given a string, sets the procline ($0)
+      # Procline is always in the format of:
+      #   resque-pool-master: STRING
+      def procline(string)
+        $0 = "resque-pool-master#{app}: #{string}"
+      end
+
+      # TODO: make this use an actual logger
+      def log(message)
+        puts "resque-pool-manager#{app}[#{Process.pid}]: #{message}"
+        #$stdout.fsync
+      end
+
+      # TODO: make this use an actual logger
+      def log_worker(message)
+        puts "resque-pool-worker#{app}[#{Process.pid}]: #{message}"
+        #$stdout.fsync
+      end
+
+      # Include optional app name in procline
+      def app
+        app_name   = self.respond_to?(:app_name)       && self.app_name
+        app_name ||= self.class.respond_to?(:app_name) && self.class.app_name
+        app_name ? "[#{app_name}]" : ""
+      end
+
+    end
+  end
+end

data/lib/resque/pool/pooled_worker.rb

@@ -0,0 +1,21 @@
+require 'resque/worker'
+
+class Resque::Pool
+  module PooledWorker
+    def shutdown_with_pool
+      shutdown_without_pool || Process.ppid == 1
+    end
+
+    def self.included(base)
+      base.instance_eval do
+        alias_method :shutdown_without_pool, :shutdown?
+        alias_method :shutdown?, :shutdown_with_pool
+      end
+    end
+
+  end
+end
+
+Resque::Worker.class_eval do
+  include Resque::Pool::PooledWorker
+end

data/lib/resque/pool/tasks.rb

@@ -0,0 +1,20 @@
+# -*- encoding: utf-8 -*-
+require 'resque/tasks'
+
+namespace :resque do
+
+  # resque worker config (not pool related).  e.g. hoptoad, rails environment
+  task :setup
+
+  namespace :pool do
+     # resque pool config.  e.g. after_prefork connection handling
+    task :setup
+  end
+
+  desc "Launch a pool of resque workers"
+  task :pool => %w[resque:setup resque:pool:setup] do
+    require 'resque/pool'
+    Resque::Pool.run
+  end
+
+end

data/lib/resque/pool/version.rb

@@ -0,0 +1,5 @@
+module Resque
+  class Pool
+    VERSION = "0.3.1"
+  end
+end

data/lib/resque/pool.rb

@@ -0,0 +1,342 @@
+# -*- encoding: utf-8 -*-
+require 'resque'
+require 'resque/worker'
+require 'resque/pool/version'
+require 'resque/pool/logging'
+require 'resque/pool/pooled_worker'
+require 'fcntl'
+require 'yaml'
+
+module Resque
+  class Pool
+    SIG_QUEUE_MAX_SIZE = 5
+    DEFAULT_WORKER_INTERVAL = 5
+    QUEUE_SIGS = [ :QUIT, :INT, :TERM, :USR1, :USR2, :CONT, :HUP, :WINCH, ]
+    CHUNK_SIZE = (16 * 1024)
+
+    include Logging
+    extend  Logging
+    attr_reader :config
+    attr_reader :workers
+
+    def initialize(config)
+      init_config(config)
+      @workers = Hash.new { |workers, queues| workers[queues] = {} }
+      procline "(initialized)"
+    end
+
+    # Config: after_prefork {{{
+
+    # The `after_prefork` hook will be run in workers if you are using the
+    # preforking master worker to save memory. Use this hook to reload
+    # database connections and so forth to ensure that they're not shared
+    # among workers.
+    #
+    # Call with a block to set the hook.
+    # Call with no arguments to return the hook.
+    def self.after_prefork(&block)
+      block ? (@after_prefork = block) : @after_prefork
+    end
+
+    # Set the after_prefork proc.
+    def self.after_prefork=(after_prefork)
+      @after_prefork = after_prefork
+    end
+
+    def call_after_prefork!
+      self.class.after_prefork && self.class.after_prefork.call
+    end
+
+    # }}}
+    # Config: class methods to start up the pool using the default config {{{
+
+    @config_files = ["resque-pool.yml", "config/resque-pool.yml"]
+    class << self; attr_accessor :config_files, :app_name; end
+
+    def self.app_name
+      @app_name ||= File.basename(Dir.pwd)
+    end
+
+    def self.choose_config_file
+      if ENV["RESQUE_POOL_CONFIG"]
+        ENV["RESQUE_POOL_CONFIG"]
+      else
+        @config_files.detect { |f| File.exist?(f) }
+      end
+    end
+
+    def self.run
+      if GC.respond_to?(:copy_on_write_friendly=)
+        GC.copy_on_write_friendly = true
+      end
+      Resque::Pool.new(choose_config_file).start.join
+    end
+
+    # }}}
+    # Config: load config and config file {{{
+
+    def config_file
+      @config_file || (!@config && ::Resque::Pool.choose_config_file)
+    end
+
+    def init_config(config)
+      case config
+      when String, nil
+        @config_file = config
+      else
+        @config = config.dup
+      end
+      load_config
+    end
+
+    def load_config
+      if config_file
+        @config = YAML.load_file(config_file)
+      else
+        @config ||= {}
+      end
+      environment and @config[environment] and config.merge!(@config[environment])
+      config.delete_if {|key, value| value.is_a? Hash }
+    end
+
+    def environment
+      if defined? RAILS_ENV
+        RAILS_ENV
+      else
+        ENV['RACK_ENV'] || ENV['RAILS_ENV'] || ENV['RESQUE_ENV']
+      end
+    end
+
+    # }}}
+
+    # Sig handlers and self pipe management {{{
+
+    def self_pipe; @self_pipe ||= [] end
+    def sig_queue; @sig_queue ||= [] end
+
+    def init_self_pipe!
+      self_pipe.each { |io| io.close rescue nil }
+      self_pipe.replace(IO.pipe)
+      self_pipe.each { |io| io.fcntl(Fcntl::F_SETFD, Fcntl::FD_CLOEXEC) }
+    end
+
+    def init_sig_handlers!
+      QUEUE_SIGS.each { |sig| trap_deferred(sig) }
+      trap(:CHLD)     { |_| awaken_master }
+    end
+
+    def awaken_master
+      begin
+        self_pipe.last.write_nonblock('.') # wakeup master process from select
+      rescue Errno::EAGAIN, Errno::EINTR
+        # pipe is full, master should wake up anyways
+        retry
+      end
+    end
+
+    class QuitNowException < Exception; end
+    # defer a signal for later processing in #join (master process)
+    def trap_deferred(signal)
+      trap(signal) do |sig_nr|
+        if @waiting_for_reaper && [:INT, :TERM].include?(signal)
+          log "Recieved #{signal}: short circuiting QUIT waitpid"
+          raise QuitNowException
+        end
+        if sig_queue.size < SIG_QUEUE_MAX_SIZE
+          sig_queue << signal
+          awaken_master
+        else
+          log "ignoring SIG#{signal}, queue=#{sig_queue.inspect}"
+        end
+      end
+    end
+
+    def reset_sig_handlers!
+      QUEUE_SIGS.each {|sig| trap(sig, "DEFAULT") }
+    end
+
+    def handle_sig_queue!
+      case signal = sig_queue.shift
+      when :USR1, :USR2, :CONT
+        log "#{signal}: sending to all workers"
+        signal_all_workers(signal)
+      when :HUP
+        log "HUP: reload config file and reload logfiles"
+        load_config
+        Logging.reopen_logs!
+        log "HUP: gracefully shutdown old children (which have old logfiles open)"
+        signal_all_workers(:QUIT)
+        log "HUP: new children will inherit new logfiles"
+        maintain_worker_count
+      when :WINCH
+        log "WINCH: gracefully stopping all workers"
+        @config = {}
+        maintain_worker_count
+      when :QUIT
+        log "QUIT: graceful shutdown, waiting for children"
+        signal_all_workers(:QUIT)
+        reap_all_workers(0) # will hang until all workers are shutdown
+        :break
+      when :INT
+        log "INT: immediate shutdown (graceful worker shutdown)"
+        signal_all_workers(:QUIT)
+        :break
+      when :TERM
+        log "TERM: immediate shutdown (and immediate worker shutdown)"
+        signal_all_workers(:TERM)
+        :break
+      end
+    end
+
+    # }}}
+    # start, join, and master sleep {{{
+
+    def start
+      procline("(starting)")
+      init_self_pipe!
+      init_sig_handlers!
+      maintain_worker_count
+      procline("(started)")
+      log "started manager"
+      report_worker_pool_pids
+      self
+    end
+
+    def report_worker_pool_pids
+      if workers.empty?
+        log "Pool is empty"
+      else
+        log "Pool contains worker PIDs: #{all_pids.inspect}"
+      end
+    end
+
+    def join
+      loop do
+        reap_all_workers
+        break if handle_sig_queue! == :break
+        if sig_queue.empty?
+          master_sleep
+          maintain_worker_count
+        end
+        procline("managing #{all_pids.inspect}")
+      end
+      procline("(shutting down)")
+      #stop # gracefully shutdown all workers on our way out
+      log "manager finished"
+      #unlink_pid_safe(pid) if pid
+    end
+
+    def master_sleep
+      begin
+        ready = IO.select([self_pipe.first], nil, nil, 1) or return
+        ready.first && ready.first.first or return
+        loop { self_pipe.first.read_nonblock(CHUNK_SIZE) }
+      rescue Errno::EAGAIN, Errno::EINTR
+      end
+    end
+
+    # }}}
+    # worker process management {{{
+
+    def reap_all_workers(waitpid_flags=Process::WNOHANG)
+      @waiting_for_reaper = waitpid_flags == 0
+      begin
+        loop do
+          # -1, wait for any child process
+          wpid, status = Process.waitpid2(-1, waitpid_flags)
+          break unless wpid
+
+          if worker = delete_worker(wpid)
+            log "Reaped resque worker[#{status.pid}] (status: #{status.exitstatus}) queues: #{worker.queues.join(",")}"
+          else
+            # this died before it could be killed, so it's not going to have any extra info
+            log "Tried to reap worker [#{status.pid}], but it had already died. (status: #{status.exitstatus})"
+          end
+        end
+      rescue Errno::ECHILD, QuitNowException
+      end
+    end
+
+    # TODO: close any file descriptors connected to worker, if any
+    def delete_worker(pid)
+      worker = nil
+      workers.detect do |queues, pid_to_worker|
+        worker = pid_to_worker.delete(pid)
+      end
+      worker
+    end
+
+    def all_pids
+      workers.map {|q,workers| workers.keys }.flatten
+    end
+
+    def signal_all_workers(signal)
+      all_pids.each do |pid|
+        Process.kill signal, pid
+      end
+    end
+
+    # }}}
+    # ???: maintain_worker_count, all_known_queues {{{
+
+    def maintain_worker_count
+      all_known_queues.each do |queues|
+        delta = worker_delta_for(queues)
+        spawn_missing_workers_for(queues) if delta > 0
+        quit_excess_workers_for(queues)   if delta < 0
+      end
+    end
+
+    def all_known_queues
+      config.keys | workers.keys
+    end
+
+    # }}}
+    # methods that operate on a single grouping of queues {{{
+    # perhaps this means a class is waiting to be extracted
+
+    def spawn_missing_workers_for(queues)
+      worker_delta_for(queues).times do |nr|
+        spawn_worker!(queues)
+      end
+    end
+
+    def quit_excess_workers_for(queues)
+      delta = -worker_delta_for(queues)
+      pids_for(queues)[0...delta].each do |pid|
+        Process.kill("QUIT", pid)
+      end
+    end
+
+    def worker_delta_for(queues)
+      config.fetch(queues, 0) - workers.fetch(queues, []).size
+    end
+
+    def pids_for(queues)
+      workers[queues].keys
+    end
+
+    def spawn_worker!(queues)
+      worker = create_worker(queues)
+      pid = fork do
+        log_worker "Starting worker #{worker}"
+        call_after_prefork!
+        reset_sig_handlers!
+        #self_pipe.each {|io| io.close }
+        worker.work(ENV['INTERVAL'] || DEFAULT_WORKER_INTERVAL) # interval, will block
+      end
+      workers[queues][pid] = worker
+    end
+
+    def create_worker(queues)
+      queues = queues.to_s.split(',')
+      worker = ::Resque::Worker.new(*queues)
+      worker.verbose = ENV['LOGGING'] || ENV['VERBOSE']
+      worker.very_verbose = ENV['VVERBOSE']
+      worker
+    end
+
+    # }}}
+
+  end
+end