refactor hook calling a bit

Author: defunkt

lib/resque/helpers.rb

@@ -6,16 +6,6 @@ def redis
       Resque.redis
     end
 
-    #Direct access to the before_first_fork proc
-    def before_first_fork
-      Resque.before_first_fork  
-    end
-
-    #Direct access to the after_fork proc
-    def after_fork
-      Resque.after_fork  
-    end
- 
     # Given a Ruby object, returns a string suitable for storage in a
     # queue.
     def encode(object)

lib/resque/worker.rb

@@ -112,6 +112,7 @@ def work(interval = 5, &block)
 
         if not @paused and job = reserve
           log "got: #{job.inspect}"
+          run_hook :before_fork
 
           if @child = fork
             rand # Reseeding
@@ -142,7 +143,7 @@ def process(job = nil)
       return unless job ||= reserve
 
       begin
-        call_after_fork
+        run_hook :after_fork, job
         working_on job
         job.perform
       rescue Object => e
@@ -203,7 +204,7 @@ def startup
       enable_gc_optimizations
       register_signal_handlers
       prune_dead_workers
-      call_before_first_fork
+      run_hook :before_first_fork
       register_worker
     end
 
@@ -308,18 +309,14 @@ def register_worker
       started!
     end
 
-    #Call any before_first_fork procs, if any
-    def call_before_first_fork
-      return unless before_first_fork
-      log "Calling before_first_fork"
-      before_first_fork.call
-    end
+    # Runs a named hook, passing along any arguments.
+    def run_hook(name, *args)
+      return unless hook = Resque.send(name)
+      msg = "Running #{name} hook"
+      msg << " with #{args.inspect}" if args.any?
+      log msg
 
-    #Call any before_first_fork procs, if any
-    def call_after_fork
-      return unless after_fork
-      log "Calling after_fork"
-      after_fork.call
+      args.any? ? hook.call(*args) : hook.call
     end
 
     # Unregisters ourself as a worker. Useful when shutting down.

test/worker_test.rb

@@ -4,6 +4,10 @@
   setup do
     Resque.redis.flush_all
 
+    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
@@ -243,15 +247,31 @@
 
   test "Will call a before_first_fork hook only once" do
     Resque.redis.flush_all
-    $BEFORE_FORK_CALLED = false
-    Resque.before_first_fork = Proc.new { $BEFORE_FORK_CALLED = true }
+    $BEFORE_FORK_CALLED = 0
+    Resque.before_first_fork = Proc.new { $BEFORE_FORK_CALLED += 1 }
     workerA = Resque::Worker.new(:jobs)
     Resque::Job.create(:jobs, SomeJob, 20, '/tmp')
 
+    assert_equal 0, $BEFORE_FORK_CALLED
+
+    workerA.work(0)
+    assert_equal 1, $BEFORE_FORK_CALLED
+
+    # TODO: Verify it's only run once. Not easy.
+#     workerA.work(0)
+#     assert_equal 1, $BEFORE_FORK_CALLED
+  end
+
+  test "Will call a before_fork hook before forking" do
+    Resque.redis.flush_all
+    $BEFORE_FORK_CALLED = false
+    Resque.before_fork = Proc.new { $BEFORE_FORK_CALLED = true }
+    workerA = Resque::Worker.new(:jobs)
+
     assert !$BEFORE_FORK_CALLED
+    Resque::Job.create(:jobs, SomeJob, 20, '/tmp')
     workerA.work(0)
     assert $BEFORE_FORK_CALLED
-    Resque.before_first_fork = nil
   end
 
   test "Will call an after_fork hook after forking" do
@@ -264,6 +284,5 @@
     Resque::Job.create(:jobs, SomeJob, 20, '/tmp')
     workerA.work(0)
     assert $AFTER_FORK_CALLED
-    Resque.after_fork = nil
   end
 end