From b06763082feeb361714c67ebf0575a63c58d3f87 Mon Sep 17 00:00:00 2001
From: Bill Ruddock <bill.ruddock@gmail.com>
Date: Tue, 30 Aug 2022 16:45:33 +0100
Subject: [PATCH] Add on_thread_exit hook

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).
---
 lib/puma/dsl.rb          | 17 +++++++++++++++++
 lib/puma/server.rb       |  1 +
 lib/puma/thread_pool.rb  | 13 +++++++++++++
 test/test_config.rb      |  4 ++++
 test/test_thread_pool.rb | 18 ++++++++++++++++++
 5 files changed, 53 insertions(+)

diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index 7026cbc7ca..5336f9bf9f 100644
--- a/lib/puma/dsl.rb
+++ b/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.
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index b5bfff5670..fa6cb3a226 100644
--- a/lib/puma/server.rb
+++ b/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))
diff --git a/lib/puma/thread_pool.rb b/lib/puma/thread_pool.rb
index e3e1915a9e..eb96e618c7 100644
--- a/lib/puma/thread_pool.rb
+++ b/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?
diff --git a/test/test_config.rb b/test/test_config.rb
index 65297941ce..df3d5bdbd9 100644
--- a/test/test_config.rb
+++ b/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|
diff --git a/test/test_thread_pool.rb b/test/test_thread_pool.rb
index 8021e7be7c..5f785f81dc 100644
--- a/test/test_thread_pool.rb
+++ b/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)