Merge pull request resque#680 from mrzor/hooks

Enable registering of multiple Resque hooks (2)

Author: steveklabnik

.gitignore

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

lib/resque.rb

@@ -85,41 +85,46 @@ 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).
@@ -403,5 +408,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

@@ -405,12 +405,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

test/worker_test.rb

@@ -5,10 +5,6 @@
     Resque.redis = Resque.redis # reset state in Resque object
     Resque.redis.flushall
 
-    Resque.before_first_fork = nil
-    Resque.before_fork = nil
-    Resque.after_fork = nil
-
     @worker = Resque::Worker.new(:jobs)
     Resque::Job.create(:jobs, SomeJob, 20, '/tmp')
   end