diff --git a/README.md b/README.md index f664a02..a678de2 100644 --- a/README.md +++ b/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 -------------- diff --git a/lib/resque/pool.rb b/lib/resque/pool.rb index 1841c6e..50fcc86 100644 --- a/lib/resque/pool.rb +++ b/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 diff --git a/lib/resque/pool/file_or_hash_loader.rb b/lib/resque/pool/file_or_hash_loader.rb new file mode 100644 index 0000000..8a15a95 --- /dev/null +++ b/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 diff --git a/spec/resque_pool_spec.rb b/spec/resque_pool_spec.rb index aabb58a..09233e6 100644 --- a/spec/resque_pool_spec.rb +++ b/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 diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb index 427c7e3..f0e8478 100644 --- a/spec/spec_helper.rb +++ b/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