Add on_thread_exit hook

lib/puma/dsl.rb

@@ -641,6 +641,23 @@ def on_refork(&block)
       @options[:before_refork] << block
     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/server.rb

@@ -241,6 +241,7 @@ def run(background=true, thread_name: 'srv')
 
       @thread_pool.out_of_band_hook = @options[:out_of_band]
       @thread_pool.clean_thread_locals = @options[:clean_thread_locals]
+      @thread_pool.before_thread_exit_hook = @options[:before_thread_exit]
 
       if @queue_requests
         @reactor = Reactor.new(@io_selector_backend, &method(:reactor_wakeup))

lib/puma/thread_pool.rb

@@ -70,6 +70,7 @@ def initialize(name, min, max, *extra, &block)
     attr_reader :spawned, :trim_requested, :waiting
     attr_accessor :clean_thread_locals
     attr_accessor :out_of_band_hook # @version 5.0.0
+    attr_accessor :before_thread_exit_hook
 
     def self.clean_thread_locals
       Thread.current.keys.each do |key| # rubocop: disable Style/HashEachMethods
@@ -121,6 +122,7 @@ def spawn_thread
                 @spawned -= 1
                 @workers.delete th
                 not_full.signal
+                trigger_before_thread_exit_hook
                 Thread.exit
               end
 
@@ -158,6 +160,17 @@ def spawn_thread
 
     private :spawn_thread
 
+    def trigger_before_thread_exit_hook
+      return unless before_thread_exit_hook && before_thread_exit_hook.any?
+
+      before_thread_exit_hook.each(&:call)
+      nil
+    rescue Exception => e
+      STDERR.puts "Exception calling before_thread_exit_hook: #{e.message} (#{e.class})"
+    end
+
+    private :trigger_before_thread_exit_hook
+
     # @version 5.0.0
     def trigger_out_of_band_hook
       return false unless out_of_band_hook && out_of_band_hook.any?

test/test_config.rb

@@ -430,6 +430,10 @@ def test_run_hooks_before_fork
     assert_run_hooks :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

@@ -159,6 +159,24 @@ def test_trim_is_ignored_if_no_waiting_threads
     assert_equal 0, pool.trim_requested
   end
 
+  def test_trim_thread_exit_hook
+    pool = mutex_pool(0, 1)
+    exited = Queue.new
+    pool.before_thread_exit_hook = [
+      proc do
+        exited << 1
+      end
+    ]
+
+    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)