Enable registering of multiple Resque hooks.

Calling `Resque.before_fork {}` would register a hook, but if any other gem,
plugin or file also tried to set a before fork hook, the original would be
replaced by the new hook. An example of this is the rpm_contrib gem which allows
New Relic instrumentation of Resque jobs. The gem defines a `before_first_fork`
hook and an `after_fork` hook. This results in conflicts if a user wants to
define their own hooks of the same type.

This fix allows multiple hooks to be registered and executes all of them
instead of simply executing the last registered hook. I've consolidated
the related tests into a single file and added some tests to ensure that the
hooks are called, called at the right time, and that multiple hooks can be
registered.

Author: panthomakos

.gitignore

@@ -1,3 +1,4 @@
 Gemfile.lock
 doc/
 test/dump.rdb
+test/dump-cluster.rdb

lib/resque.rb

@@ -85,47 +85,51 @@ def redis_id
   # changes you make will be permanent for the lifespan of the
   # worker.
   #
-  # Call with a block to set the hook.
-  # Call with no arguments to return the hook.
+  # Call with a block to register a hook.
+  # Call with no arguments to return all registered hooks.
   def before_first_fork(&block)
-    block ? (@before_first_fork = block) : @before_first_fork
+    block ? register_hook(:before_first_fork, block) : hooks(:before_first_fork)
   end
 
-  # Set a proc that will be called in the parent process before the
-  # worker forks for the first time.
-  attr_writer :before_first_fork
+  # Register a before_first_fork proc.
+  def before_first_fork=(block)
+    register_hook(:before_first_fork, block)
+  end
 
   # The `before_fork` hook will be run in the **parent** process
   # before every job, so be careful- any changes you make will be
   # permanent for the lifespan of the worker.
   #
-  # Call with a block to set the hook.
-  # Call with no arguments to return the hook.
+  # Call with a block to register a hook.
+  # Call with no arguments to return all registered hooks.
   def before_fork(&block)
-    block ? (@before_fork = block) : @before_fork
+    block ? register_hook(:before_fork, block) : hooks(:before_fork)
   end
 
-  # Set the before_fork proc.
-  attr_writer :before_fork
+  # Register a before_fork proc.
+  def before_fork=(block)
+    register_hook(:before_fork, block)
+  end
 
   # The `after_fork` hook will be run in the child process and is passed
   # the current job. Any changes you make, therefore, will only live as
   # long as the job currently being processed.
   #
-  # Call with a block to set the hook.
-  # Call with no arguments to return the hook.
+  # Call with a block to register a hook.
+  # Call with no arguments to return all registered hooks.
   def after_fork(&block)
-    block ? (@after_fork = block) : @after_fork
+    block ? register_hook(:after_fork, block) : hooks(:after_fork)
   end
 
-  # Set the after_fork proc.
-  attr_writer :after_fork
+  # Register an after_fork proc.
+  def after_fork=(block)
+    register_hook(:after_fork, block)
+  end
 
   # The `before_pause` hook will be run in the parent process before the
   # worker has paused processing (via #pause_processing or SIGUSR2).
   def before_pause(&block)
-    @before_pause = block if block_given?
-    @before_pause
+    block ? register_hook(:before_pause, block) : hooks(:before_pause)
   end
 
   # Set the after_pause proc.
@@ -134,8 +138,7 @@ def before_pause(&block)
   # The `after_pause` hook will be run in the parent process after the
   # worker has paused (via SIGCONT).
   def after_pause(&block)
-    @after_pause = block if block_given?
-    @after_pause
+    block ? register_hook(:after_pause, block) : hooks(:after_pause)
   end
 
   # Set the after_continue proc.
@@ -403,5 +406,29 @@ def keys
       key.sub("#{redis.namespace}:", '')
     end
   end
+
+  private
+
+  # Register a new proc as a hook. If the block is nil this is the
+  # equivalent of removing all hooks of the given name.
+  #
+  # `name` is the hook that the block should be registered with.
+  def register_hook(name, block)
+    return clear_hooks(name) if block.nil?
+
+    @hooks ||= {}
+    @hooks[name] ||= []
+    @hooks[name] << block
+  end
+
+  # Clear all hooks given a hook name.
+  def clear_hooks(name)
+    @hooks && @hooks[name] = []
+  end
+
+  # Retrieve all hooks of a given name.
+  def hooks(name)
+    (@hooks && @hooks[name]) || []
+  end
 end
 

lib/resque/worker.rb

@@ -425,12 +425,14 @@ def register_worker
 
     # Runs a named hook, passing along any arguments.
     def run_hook(name, *args)
-      return unless hook = Resque.send(name)
-      msg = "Running #{name} hook"
+      return unless hooks = Resque.send(name)
+      msg = "Running #{name} hooks"
       msg << " with #{args.inspect}" if args.any?
       log msg
 
-      args.any? ? hook.call(*args) : hook.call
+      hooks.each do |hook|
+        args.any? ? hook.call(*args) : hook.call
+      end
     end
 
     # Unregisters ourself as a worker. Useful when shutting down.

test/resque_hook_test.rb

@@ -0,0 +1,120 @@
+require 'test_helper'
+
+describe "Resque Hooks" do
+  before do
+    Resque.redis.flushall
+
+    Resque.before_first_fork = nil
+    Resque.before_fork = nil
+    Resque.after_fork = nil
+
+    @worker = Resque::Worker.new(:jobs)
+
+    $called = false
+
+    class CallNotifyJob
+      def self.perform
+        $called = true
+      end
+    end
+  end
+
+  it 'retrieving hooks if none have been set' do
+    assert_equal [], Resque.before_first_fork
+    assert_equal [], Resque.before_fork
+    assert_equal [], Resque.after_fork
+  end
+
+  it 'it calls before_first_fork once' do
+    counter = 0
+
+    Resque.before_first_fork { counter += 1 }
+    2.times { Resque::Job.create(:jobs, CallNotifyJob) }
+
+    assert_equal(0, counter)
+    @worker.work(0)
+    assert_equal(1, counter)
+  end
+
+  it 'it calls before_fork before each job' do
+    counter = 0
+
+    Resque.before_fork { counter += 1 }
+    2.times { Resque::Job.create(:jobs, CallNotifyJob) }
+
+    assert_equal(0, counter)
+    @worker.work(0)
+    assert_equal(2, counter)
+  end
+
+  it 'it calls after_fork after each job' do
+    counter = 0
+
+    Resque.after_fork { counter += 1 }
+    2.times { Resque::Job.create(:jobs, CallNotifyJob) }
+
+    assert_equal(0, counter)
+    @worker.work(0)
+    assert_equal(2, counter)
+  end
+
+  it 'it calls before_first_fork before forking' do
+    Resque.before_first_fork { assert(!$called) }
+
+    Resque::Job.create(:jobs, CallNotifyJob)
+    @worker.work(0)
+  end
+
+  it 'it calls before_fork before forking' do
+    Resque.before_fork { assert(!$called) }
+
+    Resque::Job.create(:jobs, CallNotifyJob)
+    @worker.work(0)
+  end
+
+  it 'it calls after_fork after forking' do
+    Resque.after_fork { assert($called) }
+
+    Resque::Job.create(:jobs, CallNotifyJob)
+    @worker.work(0)
+  end
+
+  it 'it registeres multiple before_first_forks' do
+    first = false
+    second = false
+
+    Resque.before_first_fork { first = true }
+    Resque.before_first_fork { second = true }
+    Resque::Job.create(:jobs, CallNotifyJob)
+
+    assert(!first && !second)
+    @worker.work(0)
+    assert(first && second)
+  end
+
+  it 'it registers multiple before_forks' do
+    first = false
+    second = false
+
+    Resque.before_fork { first = true }
+    Resque.before_fork { second = true }
+    Resque::Job.create(:jobs, CallNotifyJob)
+
+    assert(!first && !second)
+    @worker.work(0)
+    assert(first && second)
+  end
+
+  it 'it registers multiple after_forks' do
+    first = false
+    second = false
+
+    Resque.after_fork { first = true }
+    Resque.after_fork { second = true }
+    Resque::Job.create(:jobs, CallNotifyJob)
+
+    assert(!first && !second)
+    @worker.work(0)
+    assert(first && second)
+  end
+end