Enable custom logic to determine worker configuration

Previously, configuration could only be loaded from a
YAML file, or a Hash instance.

This change introduces the concept of a *configuration loader*.

All of the previous configuration logic has been moved to the
FileOrHashLoader class, which will be used by default.

To provide your own logic, you set the `configuration_loader`
class variable on Resque::Pool to an object that
responds to `#call` (which can simply be a lambda/Proc).

See the README for more details.

README.md

@@ -132,6 +132,29 @@ defined by Resque 1.22 and above. See http://hone.heroku.com/resque/2012/08/21/r
 for details, overriding any command-line configuration for `TERM`. Setting `TERM_CHILD` tells
 us you know what you're doing.
 
+Custom Configuration Loader
+---------------------------
+
+If the static YAML file configuration approach does not meet you needs, you can
+specify a custom configuration loader.
+
+Set the `config_loader` class variable on Resque::Pool to an object that
+responds to `#call` (which can simply be a lambda/Proc). The class attribute
+needs to be set before starting the pool. This is usually accomplished by
+in the `resque:pool:setup` rake task, as described above.
+
+For example, if you wanted to vary the number of worker processes based on a
+value stored in Redis, you could do something like:
+
+```ruby
+task resque:pool:setup do
+Resque::Pool.config_loader = lambda {|env|
+  worker_count = Redis.current.get("pool_workers_#{env}").to_i
+  {"queueA,queueB" => worker_count }
+}
+end
+```
+
 Other Features
 --------------
 

lib/resque/pool.rb

@@ -4,6 +4,7 @@
 require 'resque/pool/version'
 require 'resque/pool/logging'
 require 'resque/pool/pooled_worker'
+require 'resque/pool/file_or_hash_loader'
 require 'erb'
 require 'fcntl'
 require 'yaml'
@@ -18,10 +19,11 @@ class Pool
     include Logging
     extend  Logging
     attr_reader :config
+    attr_reader :config_loader
     attr_reader :workers
 
-    def initialize(config)
-      init_config(config)
+    def initialize(config_loader=nil)
+      init_config(config_loader)
       @workers = Hash.new { |workers, queues| workers[queues] = {} }
       procline "(initialized)"
     end
@@ -56,10 +58,9 @@ def call_after_prefork!
     end
 
     # }}}
-    # Config: class methods to start up the pool using the default config {{{
+    # Config: class methods to start up the pool using the config loader {{{
 
-    @config_files = ["resque-pool.yml", "config/resque-pool.yml"]
-    class << self; attr_accessor :config_files, :app_name; end
+    class << self; attr_accessor :config_loader, :app_name; end
 
     def self.app_name
       @app_name ||= File.basename(Dir.pwd)
@@ -81,46 +82,32 @@ def self.single_process_group
       )
     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
+      create_configured.start.join
     end
 
-    # }}}
-    # Config: load config and config file {{{
-
-    def config_file
-      @config_file || (!@config && ::Resque::Pool.choose_config_file)
+    def self.create_configured
+      Resque::Pool.new(config_loader)
     end
 
-    def init_config(config)
-      case config
-      when String, nil
-        @config_file = config
+    # }}}
+    # Config: store loader and load config {{{
+
+    def init_config(loader)
+      case loader
+      when String, Hash, nil
+        @config_loader = FileOrHashLoader.new(loader)
       else
-        @config = config.dup
+        @config_loader = loader
       end
       load_config
     end
 
     def load_config
-      if config_file
-        @config = YAML.load(ERB.new(IO.read(config_file)).result)
-      else
-        @config ||= {}
-      end
-      environment and @config[environment] and config.merge!(@config[environment])
-      config.delete_if {|key, value| value.is_a? Hash }
+      @config = config_loader.call(environment)
     end
 
     def environment

lib/resque/pool/file_or_hash_loader.rb

@@ -0,0 +1,51 @@
+class FileOrHashLoader
+  def initialize(filename_or_hash=nil)
+    case filename_or_hash
+    when String, nil
+      @filename = filename_or_hash
+    when Hash
+      @static_config = filename_or_hash.dup
+    else
+      raise "#{self.class} cannot be initialized with #{filename_or_hash.inspect}"
+    end
+  end
+
+  def call(environment)
+    load_config_from_file(environment)
+  end
+
+  private
+
+  def load_config_from_file(environment)
+    if @static_config
+      new_config = @static_config
+    else
+      filename = config_filename
+      new_config = load_config filename
+    end
+    apply_environment new_config, environment
+  end
+
+  def apply_environment(config, environment)
+    environment and config[environment] and config.merge!(config[environment])
+    config.delete_if {|key, value| value.is_a? Hash }
+  end
+
+  def config_filename
+    @filename || choose_config_file
+  end
+
+  def load_config(filename)
+    return {} unless filename
+    YAML.load(ERB.new(IO.read(filename)).result)
+  end
+
+  CONFIG_FILES = ["resque-pool.yml", "config/resque-pool.yml"]
+  def choose_config_file
+    if ENV["RESQUE_POOL_CONFIG"]
+      ENV["RESQUE_POOL_CONFIG"]
+    else
+      CONFIG_FILES.detect { |f| File.exist?(f) }
+    end
+  end
+end

spec/resque_pool_spec.rb

@@ -1,6 +1,7 @@
 require 'spec_helper'
 
 RSpec.configure do |config|
+  config.include PoolSpecHelpers
   config.after {
     Object.send(:remove_const, :RAILS_ENV) if defined? RAILS_ENV
     ENV.delete 'RACK_ENV'
@@ -158,7 +159,7 @@ module Rails; end
 
   context "when a custom file is specified" do
     before { ENV["RESQUE_POOL_CONFIG"] = 'spec/resque-pool-custom.yml.erb' }
-    subject { Resque::Pool.new(Resque::Pool.choose_config_file) }
+    subject { Resque::Pool.new }
     it "should find the right file, and parse the ERB" do
       subject.config["foo"].should == 2
     end
@@ -175,7 +176,7 @@ module Rails; end
     }
 
     subject {
-      Resque::Pool.new(config_file_path).tap{|p| p.stub(:spawn_worker!) {} }
+      no_spawn(Resque::Pool.new(config_file_path))
     }
 
     it "should not automatically load the changes" do
@@ -191,18 +192,66 @@ module Rails; end
       File.open(config_file_path, "w"){|f| f.write "changed: 1"}
       subject.config.keys.should == ["orig"]
 
-      simulate_signal :HUP
+      simulate_signal subject, :HUP
 
       subject.config.keys.should == ["changed"]
     end
 
-    def simulate_signal(signal)
-      subject.sig_queue.clear
-      subject.sig_queue.push signal
-      subject.handle_sig_queue!
+  end
+
+end
+
+describe Resque::Pool, "the pool configuration custom loader" do
+  it "should retrieve the config based on the environment" do
+    custom_loader = double(call: Hash.new)
+    RAILS_ENV = "env"
+
+    Resque::Pool.new(custom_loader)
+
+    custom_loader.should have_received(:call).with("env")
+  end
+
+  it "should reset the config loader on HUP" do
+    custom_loader = double(call: Hash.new)
+
+    pool = no_spawn(Resque::Pool.new(custom_loader))
+    custom_loader.should have_received(:call).once
+
+    pool.sig_queue.push :HUP
+    pool.handle_sig_queue!
+    custom_loader.should have_received(:call).twice
+  end
+
+  it "can be a lambda" do
+    RAILS_ENV = "fake"
+    pool = no_spawn(Resque::Pool.new(lambda {|env|
+      {env.reverse => 1}
+    }))
+    pool.config.should == {"ekaf" => 1}
+  end
+end
+
+describe "the class-level .config_loader attribute" do
+  context "when not provided" do
+    subject { Resque::Pool.create_configured }
+
+    it "created pools use config file and hash loading logic" do
+      subject.config_loader.should be_instance_of FileOrHashLoader
     end
   end
 
+  context "when provided with a custom config loader" do
+    let(:custom_config_loader) {
+      double(call: Hash.new)
+    }
+    before(:each) { Resque::Pool.config_loader = custom_config_loader }
+    after(:each) { Resque::Pool.config_loader = nil }
+    subject { Resque::Pool.create_configured }
+
+    it "created pools use the specified config loader" do
+      subject.config_loader.should == custom_config_loader
+    end
+  end
 end
 
 describe Resque::Pool, "given after_prefork hook" do

spec/spec_helper.rb

@@ -1,3 +1,16 @@
 require 'rspec'
 $LOAD_PATH << File.expand_path("../lib", File.dirname(__FILE__))
 require 'resque/pool'
+
+module PoolSpecHelpers
+  def no_spawn(pool)
+    pool.stub(:spawn_worker!) {}
+    pool
+  end
+
+  def simulate_signal(pool, signal)
+    pool.sig_queue.clear
+    pool.sig_queue.push signal
+    pool.handle_sig_queue!
+  end
+end