Add on_thread_exit hook (#2920)

Add a hook to run when a worker thread is trimmed (exits normally).
This can be useful to clean up thread local resources that do not want to be
cleaned between every request (see clean_thread_locals for that).

Co-authored-by: Nate Berkopec <nate.berkopec@gmail.com>

lib/puma/dsl.rb

@@ -723,6 +723,23 @@ def on_refork(key = nil, &block)
       process_hook :before_refork, key, block, 'on_refork'
     end
 
+    # Code to run immediately before a thread exits. The worker does not
+    # accept new requests until this code finishes.
+    #
+    # This hook is useful for cleaning up thread local resources when a thread
+    # is trimmed.
+    #
+    # This can be called multiple times to add several hooks.
+    #
+    # @example
+    #   on_thread_exit do
+    #     puts 'On thread exit...'
+    #   end
+    def on_thread_exit(&block)
+      @options[:before_thread_exit] ||= []
+      @options[:before_thread_exit] << block
+    end
+
     # Code to run out-of-band when the worker is idle.
     # These hooks run immediately after a request has finished
     # processing and there are no busy threads on the worker.

lib/puma/thread_pool.rb

@@ -51,6 +51,7 @@ def initialize(name, options = {}, &block)
       @block = block
       @out_of_band = options[:out_of_band]
       @clean_thread_locals = options[:clean_thread_locals]
+      @before_thread_exit = options[:before_thread_exit]
       @reaping_time = options[:reaping_time]
       @auto_trim_time = options[:auto_trim_time]
 
@@ -125,6 +126,7 @@ def spawn_thread
                 @spawned -= 1
                 @workers.delete th
                 not_full.signal
+                trigger_before_thread_exit_hooks
                 Thread.exit
               end
 
@@ -162,6 +164,21 @@ def spawn_thread
 
     private :spawn_thread
 
+    def trigger_before_thread_exit_hooks
+      return unless @before_thread_exit&.any?
+
+      @before_thread_exit.each do |b|
+        begin
+          b.call
+        rescue Exception => e
+          STDERR.puts "WARNING before_thread_exit hook failed with exception (#{e.class}) #{e.message}"
+        end
+      end
+      nil
+    end
+
+    private :trigger_before_thread_exit_hooks
+
     # @version 5.0.0
     def trigger_out_of_band_hook
       return false unless @out_of_band&.any?

test/test_config.rb

@@ -483,6 +483,10 @@ def test_run_hooks_before_fork
     assert_warning_for_hooks_defined_in_single_mode :before_fork
   end
 
+  def test_run_hooks_before_thread_exit
+    assert_run_hooks :before_thread_exit, configured_with: :on_thread_exit
+  end
+
   def test_run_hooks_and_exception
     conf = Puma::Configuration.new do |c|
       c.on_restart do |a|

test/test_thread_pool.rb

@@ -167,6 +167,29 @@ def test_trim_is_ignored_if_no_waiting_threads
     assert_equal 0, pool.trim_requested
   end
 
+  def test_trim_thread_exit_hook
+    exited = Queue.new
+    options = {
+      min_threads: 0,
+      max_threads: 1,
+      before_thread_exit: [
+        proc do
+          exited << 1
+        end
+      ]
+    }
+    block = proc { }
+    pool = MutexPool.new('tst', options, &block)
+
+    pool << 1
+
+    assert_equal 1, pool.spawned
+
+    pool.trim
+    assert_equal 0, pool.spawned
+    assert_equal 1, exited.length
+  end
+
   def test_autotrim
     pool = mutex_pool(1, 2)