Add after_refork hook (#3386)

* Add after_refork hook * Documentation for after_refork * Documentation for after_refork * Documentation for after_refork
puma · Nov 27, 2024 · e058301 · e058301
1 parent ff366fb
commit e058301
Show file tree

Hide file tree

Showing 6 changed files with 73 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -159,7 +159,7 @@ on_worker_boot do
 end
 ```
 
-In addition, there is an `on_refork` hook which is used only in [`fork_worker` mode](docs/fork_worker.md),
+In addition, there is an `on_refork` and `after_refork` hooks which are used only in [`fork_worker` mode](docs/fork_worker.md),
 when the worker 0 child process forks a grandchild worker:
 
 ```ruby
@@ -169,6 +169,13 @@ on_refork do
 end
 ```
 
+```ruby
+after_refork do
+  # Used only when fork_worker mode is enabled. Add code to run inside the Puma worker 0
+  # child process after it forks a grandchild worker.
+end
+```
+
 Importantly, note the following considerations when Ruby forks a child process: 
 
 1. File descriptors such as network sockets **are** copied from the parent to the forked
@@ -186,6 +193,7 @@ Therefore, we recommend the following:
 2. If (1) is not possible, use `before_fork` and `on_refork` to disconnect the parent's socket
    connections when forking, so that they are not accidentally copied to the child process.
 3. Use `on_worker_boot` to restart any background threads on the forked child.
+4. Use `after_refork` to restart any background threads on the parent.
 
 #### Master process lifecycle hooks
 

diff --git a/docs/fork_worker.md b/docs/fork_worker.md
@@ -24,7 +24,13 @@ The `fork_worker` option allows your application to be initialized only once for
 
 ### Usage Considerations
 
-- `fork_worker` introduces a new `on_refork` configuration hook. If you were using the `before_fork` hook previously, we generally recommend to copy its logic to `on_refork`. Note that `fork_worker` triggers the `before_fork` configuration hook *only* when initially forking the master process to worker 0, and triggers the `on_refork` hook on all subsequent forks from worker 0 to additional workers.
+- `fork_worker` introduces new `on_refork` and `after_refork` configuration hooks. Note the following:
+    - When initially forking the parent process to the worker 0 child, `before_fork` will trigger on the parent process and `on_worker_boot` will trigger on the worker 0 child as normal.
+    - When forking the worker 0 child to grandchild workers, `on_refork` and `after_refork` will trigger on the worker 0 child, and `on_worker_boot` will trigger on each grandchild worker.
+    - For clarity, `before_fork` does not trigger on worker 0, and `after_refork` does not trigger on the grandchild.
+- As a general migration guide:
+    - Copy any logic within your existing `before_fork` hook to the `on_refork` hook.
+    - Consider to copy logic from your `on_worker_boot` hook to the `after_refork` hook, if it is needed to reset the state of worker 0 after it forks.
 
 ### Limitations
 

diff --git a/lib/puma/cluster.rb b/lib/puma/cluster.rb
@@ -87,6 +87,10 @@ def spawn_workers
 
       if @options[:fork_worker] && all_workers_in_phase?
         @fork_writer << "0\n"
+
+        if worker_at(0).phase > 0
+          @fork_writer << "-2\n"
+        end
       end
     end
 

diff --git a/lib/puma/cluster/worker.rb b/lib/puma/cluster/worker.rb
@@ -88,6 +88,8 @@ def run
                   server.begin_restart(true)
                   @config.run_hooks(:before_refork, nil, @log_writer, @hook_data)
                 end
+              elsif idx == -2 # refork cycle is done
+                @config.run_hooks(:after_refork, nil, @log_writer, @hook_data)
               elsif idx == 0 # restart server
                 restart_server << true << false
               else # fork worker

diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
@@ -47,6 +47,7 @@ module Puma
   #     | on_worker_boot     | :before_worker_boot     | inside, before      |
   #     | on_worker_shutdown | :before_worker_shutdown | inside, after       |
   #     | on_refork          | :before_refork          | inside              |
+  #     | after_refork       | :after_refork           | inside              |
   #
   class DSL
     ON_WORKER_KEY = [String, Symbol].freeze
@@ -858,6 +859,27 @@ def on_refork(key = nil, &block)
       process_hook :before_refork, key, block, 'on_refork'
     end
 
+    # When `fork_worker` is enabled, code to run in Worker 0
+    # after all other workers are re-forked from this process,
+    # after the server has temporarily stopped serving requests
+    # (once per complete refork cycle).
+    #
+    # This can be used to re-open any connections to remote servers
+    # (database, Redis, ...) that were closed via on_refork.
+    #
+    # This can be called multiple times to add several hooks.
+    #
+    # @note Cluster mode with `fork_worker` enabled only.
+    #
+    # @example
+    #   after_refork do
+    #     puts 'After refork...'
+    #   end
+    #
+    def after_refork(key = nil, &block)
+      process_hook :after_refork, key, block, 'after_refork'
+    end
+
     # Provide a block to be executed just before a thread is added to the thread
     # pool. Be careful: while the block executes, thread creation is delayed, and
     # probably a request will have to wait too! The new thread will not be added to

diff --git a/test/test_integration_cluster.rb b/test/test_integration_cluster.rb
@@ -322,6 +322,35 @@ def test_fork_worker_on_refork
     refute_includes pids, get_worker_pids(1, wrkrs - 1)
   end
 
+  # use three workers to keep accepting clients
+  def test_fork_worker_after_refork
+    refork = Tempfile.new 'refork2'
+    wrkrs = 3
+    cli_server "-w #{wrkrs} test/rackup/hello_with_delay.ru", config: <<~RUBY
+      fork_worker 20
+      on_refork { File.write '#{refork.path}', 'Before refork', mode: 'a+' }
+      after_refork { File.write '#{refork.path}', '-After refork', mode: 'a+' }
+    RUBY
+
+    pids = get_worker_pids 0, wrkrs
+
+    socks = []
+    until refork.read == 'Before refork-After refork'
+      refork.rewind
+      socks << fast_connect
+      sleep 0.004
+    end
+
+    100.times {
+      socks << fast_connect
+      sleep 0.004
+    }
+
+    socks.each { |s| read_body s }
+
+    refute_includes pids, get_worker_pids(1, wrkrs - 1)
+  end
+
   def test_fork_worker_spawn
     cli_server '', config: <<~CONFIG
       workers 1