allow atfork_* hooks inside app blocks for ease-of-management

This should allow users to more-easily enable/disable apps
and their dependent atfork_* hooks.

3 files changed, 20 insertions, 3 deletions

diff --git a/Documentation/yahns_config.txt b/Documentation/yahns_config.txt
index e58e578..3679494 100644
--- a/Documentation/yahns_config.txt
+++ b/Documentation/yahns_config.txt
@@ -176,6 +176,14 @@ Ruby it is running under.
 
 ## APP-BLOCK DIRECTIVES
 
+* atfork_prepare, atfork_parent, atfork_child
+
+    These are identical to the methods defined in WORKER_PROCESSES-LEVEL
+    DIRECTIVES, however they are available inside the app blocks for
+    convenience in case it is easier to organize per-app hooks.
+
+    Default: (none)
+
 * check_client_connection BOOLEAN
 
     When enabled, yahns will check the client connection by writing
@@ -467,6 +475,9 @@ Ruby it is running under.
 
 ## WORKER_PROCESSES-LEVEL DIRECTIVES
 
+Note: all of the atfork_* hooks described here are available inside the
+"app" blocks, too.
+
 * atfork_prepare &BLOCK
 
     This &BLOCK is executed in the parent before fork(2) operation.
diff --git a/lib/yahns/config.rb b/lib/yahns/config.rb
index a46868b..29f5746 100644
--- a/lib/yahns/config.rb
+++ b/lib/yahns/config.rb
@@ -22,9 +22,10 @@ class Yahns::Config # :nodoc:
       return var if @block == nil
       msg = "#{var} must be called outside of #{@block.type}"
     else
-      return var if @block && ctx == @block.type
+      ctx = Array(ctx)
+      return var if @block && ctx.include?(@block.type)
       msg = @block ? "may not be used inside a #{@block.type} block" :
-                     "must be used with a #{ctx} block"
+                     "must be used with a #{ctx.join(' or ')} block"
     end
     raise ArgumentError, msg
   end
@@ -91,7 +92,7 @@ class Yahns::Config # :nodoc:
   %w(atfork_prepare atfork_parent atfork_child).each do |fn|
     eval(
     %Q(def #{fn}(*args, &blk);) <<
-    %Q(  _check_in_block(:worker_processes, :#{fn});) <<
+    %Q(  _check_in_block([:worker_processes,:app], :#{fn});) <<
     %Q(  _add_hook(:#{fn}, block_given? ? blk : args[0]);) <<
     %Q(end)
     )
diff --git a/test/test_server.rb b/test/test_server.rb
index e6ddefc..6b9f00f 100644
--- a/test/test_server.rb
+++ b/test/test_server.rb
@@ -408,6 +408,7 @@ class TestServer < Testcase
         app(:rack, ru) {
           listen "#{host}:#{port}"
           persistent_connections false
+          atfork_child { warn "INFO hihi from app.atfork_child" }
         }
         worker_processes(1) do
           atfork_child { puts "af #$$ worker is running" }
@@ -430,6 +431,8 @@ class TestServer < Testcase
     lines = out.readlines.map!(&:chomp!)
     out.close!
 
+    assert_match %r{INFO hihi from app\.atfork_child}, File.read(err.path)
+
     assert_equal 3, lines.size
     assert_equal("af #{pid} parent about to spawn", lines.shift)
 
@@ -451,6 +454,7 @@ class TestServer < Testcase
         app(:rack, ru) {
           listen "#{host}:#{port}"
           persistent_connections false
+          atfork_child { |nr| warn "INFO hihi.#{nr} from app.atfork_child" }
         }
         worker_processes(1) do
           atfork_child { |nr| puts "af.#{nr} #$$ worker is running" }
@@ -473,6 +477,7 @@ class TestServer < Testcase
     lines = out.readlines.map!(&:chomp!)
     out.close!
 
+    assert_match %r{INFO hihi\.0 from app\.atfork_child}, File.read(err.path)
     assert_equal 3, lines.size
     assert_equal("af.0 #{pid} parent about to spawn", lines.shift)