From ac59329f9faf948ace7666167143bbfe50e1c915 Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Mon, 11 Dec 2023 11:02:30 +0100
Subject: [PATCH] Add an integration test to ensure we're resilient to mold
 crashing

This is limited to the mold exiting before it acknowledged to the
monitor process. This allows to sanity check the new process.
---
 docs/CONFIGURATION.md              |  5 ++++
 lib/pitchfork/refork_condition.rb  |  8 +++++-
 test/integration/test_reforking.rb | 45 ++++++++++++++++++++++++++++++
 test/unit/test_flock.rb            | 13 +++++++++
 4 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/docs/CONFIGURATION.md b/docs/CONFIGURATION.md
index 29c4a544..315e7009 100644
--- a/docs/CONFIGURATION.md
+++ b/docs/CONFIGURATION.md
@@ -317,6 +317,11 @@ That is the case for instance of many SQL databases protocols.
 This is also the callback in which memory optimizations, such as
 heap compaction should be done.
 
+This callback is also a good place to check for potential corruption
+issues caused by forking. If you detect something wrong, you can
+call `Process.exit`, and this mold won't be used, another one will be
+spawned later. e.g. you can check `Socket.getaddrinfo` still works, etc.
+
 ### `after_worker_fork`
 
 ```ruby
diff --git a/lib/pitchfork/refork_condition.rb b/lib/pitchfork/refork_condition.rb
index 042e8df4..2aaf289a 100644
--- a/lib/pitchfork/refork_condition.rb
+++ b/lib/pitchfork/refork_condition.rb
@@ -2,6 +2,12 @@
 
 module Pitchfork
   module ReforkCondition
+    @backoff_delay = 10.0
+
+    class << self
+      attr_accessor :backoff_delay
+    end
+
     class RequestsCount
       def initialize(request_counts)
         @limits = request_counts
@@ -31,7 +37,7 @@ def backoff?
         end
       end
 
-      def backoff!(delay = 10.0)
+      def backoff!(delay = ReforkCondition.backoff_delay)
         @backoff_until = Pitchfork.time_now + delay
       end
     end
diff --git a/test/integration/test_reforking.rb b/test/integration/test_reforking.rb
index 028a18e9..51e073f2 100644
--- a/test/integration/test_reforking.rb
+++ b/test/integration/test_reforking.rb
@@ -102,6 +102,51 @@ def Process.fork
       assert_exited(pid, 1, timeout: 5)
     end
 
+    def test_exiting_mold
+      addr, port = unused_port
+
+      pid = spawn_server(app: File.join(ROOT, "test/integration/env.ru"), config: <<~CONFIG)
+        Pitchfork::ReforkCondition.backoff_delay = 0.0
+
+        listen "#{addr}:#{port}"
+        worker_processes 2
+        spawn_timeout 2
+        refork_after [5, 5]
+        after_mold_fork do |_server, mold|
+          if mold.generation > 0
+            if File.exist?("crashed-once.txt")
+              $stderr.puts "[mold success]"
+            else
+              File.write("crashed-once.txt", "1")
+              $stderr.puts "[mold crashing]"
+              exit 1
+            end
+          end
+        end
+      CONFIG
+
+      assert_healthy("http://#{addr}:#{port}")
+      assert_stderr "worker=0 gen=0 ready"
+      assert_stderr "worker=1 gen=0 ready", timeout: 5
+
+      7.times do
+        assert_equal true, healthy?("http://#{addr}:#{port}")
+      end
+
+      assert_stderr(/mold pid=\d+ gen=1 spawned/)
+      assert_stderr("[mold crashing]")
+      assert_stderr(/mold pid=\d+ gen=1 reaped/)
+
+      10.times do
+        assert_equal true, healthy?("http://#{addr}:#{port}")
+      end
+
+      assert_stderr "worker=0 gen=1 ready", timeout: 15
+      assert_stderr "worker=1 gen=1 ready"
+
+      assert_clean_shutdown(pid)
+    end
+
     def test_fork_unsafe
       addr, port = unused_port
 
diff --git a/test/unit/test_flock.rb b/test/unit/test_flock.rb
index a5f76fcf..d65bcd37 100644
--- a/test/unit/test_flock.rb
+++ b/test/unit/test_flock.rb
@@ -58,4 +58,17 @@ def test_at_fork
     Process.wait(pid)
     assert_equal true, @flock.try_lock
   end
+
+  def test_unlock_on_crash
+    assert_predicate @flock, :try_lock
+
+    pid = fork do
+      exit!(0)
+    end
+    @flock.at_fork
+    _, status = Process.waitpid2(pid)
+    assert_predicate status, :success?
+
+    assert_predicate @flock, :try_lock
+  end
 end