From 92bae748c29d73c065726a16fa23f69505a5360b Mon Sep 17 00:00:00 2001
From: Kamil Trzcinski <ayufan@ayufan.eu>
Date: Thu, 23 Apr 2020 19:44:47 +0200
Subject: [PATCH] Improve OOB hook

This moves the hook to be executed outside of busy loop
considering the thread to be truly idle at a time of execution
of the hook
---
 History.md              |  1 +
 lib/puma/server.rb      |  7 +------
 lib/puma/thread_pool.rb | 25 ++++++++++++++++++++++++-
 3 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/History.md b/History.md
index d7f8562537..6d5a6d440c 100644
--- a/History.md
+++ b/History.md
@@ -38,6 +38,7 @@
   * Fixed a few minor concurrency bugs in ThreadPool that may have affected non-GVL Rubies (#2220)
   * Fix `out_of_band` hook never executed if the number of worker threads is > 1 (#2177)
   * Fix ThreadPool#shutdown timeout accuracy (#2221)
+  * Call `out_of_band` hook outside of busy loop (#2230)
 
 * Refactor
   * Remove unused loader argument from Plugin initializer (#2095)
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index b5269f0ddd..094b00db7c 100644
--- a/lib/puma/server.rb
+++ b/lib/puma/server.rb
@@ -226,6 +226,7 @@ def run(background=true)
         end
       end
 
+      @thread_pool.out_of_band_hook = @options[:out_of_band]
       @thread_pool.clean_thread_locals = @options[:clean_thread_locals]
 
       if @queue_requests
@@ -436,12 +437,6 @@ def process_client(client, buffer)
         rescue StandardError => e
           @events.unknown_error self, e, "Client"
         end
-
-        if @options[:out_of_band]
-          @thread_pool.with_mutex do
-            @options[:out_of_band].each(&:call) if @thread_pool.busy_threads == 1
-          end
-        end
       end
     end
 
diff --git a/lib/puma/thread_pool.rb b/lib/puma/thread_pool.rb
index d6a27c41de..d73556ef84 100644
--- a/lib/puma/thread_pool.rb
+++ b/lib/puma/thread_pool.rb
@@ -65,6 +65,7 @@ def initialize(min, max, *extra, &block)
 
     attr_reader :spawned, :trim_requested, :waiting
     attr_accessor :clean_thread_locals
+    attr_accessor :out_of_band_hook
 
     def self.clean_thread_locals
       Thread.current.keys.each do |key| # rubocop: disable Performance/HashEachMethods
@@ -117,6 +118,7 @@ def spawn_thread
 
               @waiting += 1
               not_full.signal
+              trigger_out_of_band_hook
               not_empty.wait mutex
               @waiting -= 1
             end
@@ -143,9 +145,30 @@ def spawn_thread
 
     private :spawn_thread
 
+    def trigger_out_of_band_hook
+      return unless out_of_band_hook && out_of_band_hook.any?
+
+      with_mutex do |mutex|
+        # we execute on idle hook when all threads are free
+        return unless @spawned == @waiting
+
+        # we unlock thread for the duration of out of band hook
+        # to allow other requests to be accepted
+        mutex.unlock
+
+        begin
+          out_of_band_hook.each(&:call)
+        rescue Exception => e
+          STDERR.puts "Exception calling out_of_band_hook: #{e.message} (#{e.class})"
+        ensure
+          mutex.lock
+        end
+      end
+    end
+
     def with_mutex(&block)
       @mutex.owned? ?
-        yield :
+        yield(@mutex) :
         @mutex.synchronize(&block)
     end