From c35bf33406e76c3d027dbde623f071bcbe6095bc Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Thu, 30 Nov 2023 16:50:41 +0100
Subject: [PATCH] Add a `before_fork` callback called before forking new molds
 and new workers.

Fix: https://github.com/Shopify/pitchfork/issues/76
---
 docs/CONFIGURATION.md                  | 10 ++++++++++
 lib/pitchfork/configurator.rb          |  5 +++++
 lib/pitchfork/http_server.rb           |  5 ++++-
 test/integration/test_configuration.rb | 16 ++++++++++++++++
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md
index da2a36b5..29c4a544 100644
--- a/docs/CONFIGURATION.md
+++ b/docs/CONFIGURATION.md
@@ -284,6 +284,16 @@ after_monitor_ready do |server|
 end
 ```
 
+### `before_fork`
+
+Called by the mold before forking a new workers, and by workers before they spawn a new mold.
+
+```ruby
+before_fork do |server|
+  server.logger.info("About to fork, closing connections!")
+end
+```
+
 ### `after_mold_fork`
 
 ```ruby
diff --git a/lib/pitchfork/configurator.rb b/lib/pitchfork/configurator.rb
index 001e5c5e..cd2290d7 100644
--- a/lib/pitchfork/configurator.rb
+++ b/lib/pitchfork/configurator.rb
@@ -37,6 +37,7 @@ class Configurator
       :timeout => 22,
       :logger => default_logger,
       :worker_processes => 1,
+      :before_fork => nil,
       :after_worker_fork => lambda { |server, worker|
         server.logger.info("worker=#{worker.nr} gen=#{worker.generation} pid=#{$$} spawned")
       },
@@ -133,6 +134,10 @@ def logger(obj)
       set[:logger] = obj
     end
 
+    def before_fork(*args, &block)
+      set_hook(:before_fork, block_given? ? block : args[0], 1)
+    end
+
     def after_worker_fork(*args, &block)
       set_hook(:after_worker_fork, block_given? ? block : args[0])
     end
diff --git a/lib/pitchfork/http_server.rb b/lib/pitchfork/http_server.rb
index 9babc604..08a3e8f7 100644
--- a/lib/pitchfork/http_server.rb
+++ b/lib/pitchfork/http_server.rb
@@ -75,7 +75,7 @@ def extend_deadline(extra_time)
 
     # :stopdoc:
     attr_accessor :app, :timeout, :soft_timeout, :cleanup_timeout, :spawn_timeout, :worker_processes,
-                  :after_worker_fork, :after_mold_fork,
+                  :before_fork, :after_worker_fork, :after_mold_fork,
                   :listener_opts, :children,
                   :orig_app, :config, :ready_pipe,
                   :default_middleware, :early_hints
@@ -563,6 +563,7 @@ def spawn_worker(worker, detach:)
       # reason it gets stuck before reaching the worker loop,
       # the monitor process will kill it.
       worker.update_deadline(@spawn_timeout)
+      @before_fork&.call(self)
       Pitchfork.fork_sibling do
         worker.pid = Process.pid
 
@@ -870,6 +871,8 @@ def worker_loop(worker)
     def spawn_mold(current_generation)
       return false unless @promotion_lock.try_lock
 
+      @before_fork&.call(self)
+
       begin
         Pitchfork.fork_sibling do
           mold = Worker.new(nil, pid: Process.pid, generation: current_generation)
diff --git a/test/integration/test_configuration.rb b/test/integration/test_configuration.rb
index 1bea2af9..17332cc8 100644
--- a/test/integration/test_configuration.rb
+++ b/test/integration/test_configuration.rb
@@ -26,6 +26,22 @@ def test_after_request_complete
     assert_clean_shutdown(pid)
   end
 
+  def test_before_fork
+    addr, port = unused_port
+
+    pid = spawn_server(app: File.join(ROOT, "test/integration/env.ru"), config: <<~CONFIG)
+      listen "#{addr}:#{port}"
+      worker_processes 1
+
+      before_fork do |server|
+        $stderr.puts "[before_fork]"
+      end
+    CONFIG
+
+    assert_healthy("http://#{addr}:#{port}")
+    assert_stderr("[before_fork]")
+    assert_clean_shutdown(pid)
+  end
   def test_before_worker_exit
     addr, port = unused_port