Catch exceptions in hooks defined by users

History.md

@@ -17,6 +17,7 @@
   * Preserve `BUNDLE_GEMFILE` env var when using `prune_bundler` (#1893)
   * Send 408 request timeout even when queue requests is disabled (#2119)
   * Rescue IO::WaitReadable instead of EAGAIN for blocking read (#2121)
+  * Rescue exceptions in hooks defined by users (on_worker_boot, after_worker_fork etc) (#1551)
 
 * Refactor
   * Remove unused loader argument from Plugin initializer (#2095)

lib/puma/cluster.rb

@@ -137,7 +137,7 @@ def spawn_workers
 
       diff.times do
         idx = next_worker_index
-        @launcher.config.run_hooks :before_worker_fork, idx
+        @launcher.config.run_hooks :before_worker_fork, idx, @launcher.events
 
         pid = fork { worker(idx, master) }
         if !pid
@@ -149,7 +149,7 @@ def spawn_workers
         debug "Spawned worker: #{pid}"
         @workers << Worker.new(idx, pid, @phase, @options)
 
-        @launcher.config.run_hooks :after_worker_fork, idx
+        @launcher.config.run_hooks :after_worker_fork, idx, @launcher.events
       end
 
       if diff > 0
@@ -268,7 +268,7 @@ def worker(index, master)
 
       # Invoke any worker boot hooks so they can get
       # things in shape before booting the app.
-      @launcher.config.run_hooks :before_worker_boot, index
+      @launcher.config.run_hooks :before_worker_boot, index, @launcher.events
 
       server = start_server
 
@@ -310,7 +310,7 @@ def worker(index, master)
 
       # Invoke any worker shutdown hooks so they can prevent the worker
       # exiting until any background operations are completed
-      @launcher.config.run_hooks :before_worker_shutdown, index
+      @launcher.config.run_hooks :before_worker_shutdown, index, @launcher.events
     ensure
       @worker_write << "t#{Process.pid}\n" rescue nil
       @worker_write.close
@@ -486,7 +486,7 @@ def run
 
       @master_read, @worker_write = read, @wakeup
 
-      @launcher.config.run_hooks :before_fork, nil
+      @launcher.config.run_hooks :before_fork, nil, @launcher.events
       GC.compact if GC.respond_to?(:compact)
 
       spawn_workers

lib/puma/configuration.rb

@@ -275,8 +275,15 @@ def load_plugin(name)
       @plugins.create name
     end
 
-    def run_hooks(key, arg)
-      @options.all_of(key).each { |b| b.call arg }
+    def run_hooks(key, arg, events)
+      @options.all_of(key).each do |b|
+        begin
+          b.call arg
+        rescue => e
+          events.log "WARNING hook #{key} failed with exception (#{e.class}) #{e.message}"
+          events.debug e.backtrace.join("\n")
+        end
+      end
     end
 
     def self.temp_path

lib/puma/launcher.rb

@@ -237,7 +237,7 @@ def reload_worker_directory
     end
 
     def restart!
-      @config.run_hooks :on_restart, self
+      @config.run_hooks :on_restart, self, @events
 
       if Puma.jruby?
         close_binder_listeners

test/test_config.rb

@@ -197,6 +197,112 @@ def test_config_raise_exception_on_sigterm
     conf.options[:raise_exception_on_sigterm] = true
     assert_equal conf.options[:raise_exception_on_sigterm], true
   end
+
+  def test_run_hooks_on_restart_hook
+    require 'puma/events'
+
+    messages = []
+    conf = Puma::Configuration.new do |c|
+      c.on_restart do |a|
+        messages << "on_restart is called with #{a}"
+      end
+    end
+    conf.load
+
+    conf.run_hooks :on_restart, 'ARG', Puma::Events.strings
+    assert_equal messages, ['on_restart is called with ARG']
+  end
+
+  def test_run_hooks_before_worker_fork
+    require 'puma/events'
+
+    messages = []
+    conf = Puma::Configuration.new do |c|
+      c.on_worker_fork do |a|
+        messages << "before_worker_fork is called with #{a}"
+      end
+    end
+    conf.load
+
+    conf.run_hooks :before_worker_fork, 'ARG', Puma::Events.strings
+    assert_equal messages, ['before_worker_fork is called with ARG']
+  end
+
+  def test_run_hooks_after_worker_fork
+    require 'puma/events'
+
+    messages = []
+    conf = Puma::Configuration.new do |c|
+      c.after_worker_fork do |a|
+        messages << "after_worker_fork is called with #{a}"
+      end
+    end
+    conf.load
+
+    conf.run_hooks :after_worker_fork, 'ARG', Puma::Events.strings
+    assert_equal messages, ['after_worker_fork is called with ARG']
+  end
+
+  def test_run_hooks_before_worker_boot
+    require 'puma/events'
+
+    messages = []
+    conf = Puma::Configuration.new do |c|
+      c.on_worker_boot do |a|
+        messages << "before_worker_boot is called with #{a}"
+      end
+    end
+    conf.load
+
+    conf.run_hooks :before_worker_boot, 'ARG', Puma::Events.strings
+    assert_equal messages, ['before_worker_boot is called with ARG']
+  end
+
+  def test_run_hooks_before_worker_shutdown
+    require 'puma/events'
+
+    messages = []
+    conf = Puma::Configuration.new do |c|
+      c.on_worker_shutdown do |a|
+        messages << "before_worker_shutdown is called with #{a}"
+      end
+    end
+    conf.load
+
+    conf.run_hooks :before_worker_shutdown, 'ARG', Puma::Events.strings
+    assert_equal messages, ['before_worker_shutdown is called with ARG']
+  end
+
+  def test_run_hooks_before_fork
+    require 'puma/events'
+
+    messages = []
+    conf = Puma::Configuration.new do |c|
+      c.before_fork do |a|
+        messages << "before_fork is called with #{a}"
+      end
+    end
+    conf.load
+
+    conf.run_hooks :before_fork, 'ARG', Puma::Events.strings
+    assert_equal messages, ['before_fork is called with ARG']
+  end
+
+  def test_run_hooks_and_exception
+    require 'puma/events'
+
+    conf = Puma::Configuration.new do |c|
+      c.on_restart do |a|
+        raise RuntimeError, 'Error from hook'
+      end
+    end
+    conf.load
+    events = Puma::Events.strings
+
+    conf.run_hooks :on_restart, 'ARG', events
+    expected = /WARNING hook on_restart failed with exception \(RuntimeError\) Error from hook/
+    assert_match expected, events.stdout.string
+  end
 end
 
 # Thread unsafe modification of ENV