puma 2.2.0 → 8.0.2

data/lib/puma/thread_pool.rb

@@ -1,9 +1,81 @@
+# frozen_string_literal: true
+
 require 'thread'
 
+require_relative 'io_buffer'
+require_relative 'server_plugin_control'
+
 module Puma
-  # A simple thread pool management object.
+
+  # Add `Thread#puma_server` and `Thread#puma_server=`
+  Thread.attr_accessor(:puma_server)
+
+  # Internal Docs for A simple thread pool management object.
   #
+  # Each Puma "worker" has a thread pool to process requests.
+  #
+  # First a connection to a client is made in `Puma::Server`. It is wrapped in a
+  # `Puma::Client` instance and then passed to the `Puma::Reactor` to ensure
+  # the whole request is buffered into memory. Once the request is ready, it is passed into
+  # a thread pool via the `Puma::ThreadPool#<<` operator where it is stored in a `@todo` array.
+  #
+  # Each thread in the pool has an internal loop where it pulls a request from the `@todo` array
+  # and processes it.
   class ThreadPool
+    class ForceShutdown < RuntimeError
+    end
+
+    class ProcessorThread
+      attr_accessor :thread
+      attr_writer :marked_as_io_thread
+
+      def initialize(pool)
+        @pool = pool
+        @thread = nil
+        @marked_as_io_thread = false
+      end
+
+      def mark_as_io_thread!
+        unless @marked_as_io_thread
+          @marked_as_io_thread = true
+
+          # Immediately signal the pool that it can spawn a new thread
+          # if there's some work in the queue.
+          @pool.spawn_thread_if_needed
+        end
+      end
+
+      def marked_as_io_thread?
+        @marked_as_io_thread
+      end
+
+      def alive?
+        @thread&.alive?
+      end
+
+      def join(...)
+        @thread.join(...)
+      end
+
+      def kill(...)
+        @thread.kill(...)
+      end
+
+      def [](key)
+        @thread[key]
+      end
+
+      def raise(...)
+        @thread.raise(...)
+      end
+    end
+
+    # How long, after raising the ForceShutdown of a thread during
+    # forced shutdown mode, to wait for the thread to try and finish
+    # up its work before leaving the thread to die on the vine.
+    SHUTDOWN_GRACE_TIME = 5 # seconds
+
+    attr_reader :out_of_band_running
 
     # Maintain a minimum of +min+ and maximum of +max+ threads
     # in the pool.
@@ -11,39 +83,102 @@ module Puma
     # The block passed is the work that will be performed in each
     # thread.
     #
-    def initialize(min, max, *extra, &blk)
-      @cond = ConditionVariable.new
-      @mutex = Mutex.new
+    def initialize(name, options = {}, server: nil, &block)
+      @server = server
 
-      @todo = []
+      @not_empty = ConditionVariable.new
+      @not_full = ConditionVariable.new
+      @mutex = Mutex.new
+      @todo = Queue.new
 
+      @backlog_max = 0
       @spawned = 0
       @waiting = 0
 
-      @min = min
-      @max = max
-      @block = blk
-      @extra = extra
+      @name = name
+      @min = Integer(options[:min_threads])
+      @max = Integer(options[:max_threads])
+      @max_io_threads = Integer(options[:max_io_threads] || 0)
+
+      # Not an 'exposed' option, options[:pool_shutdown_grace_time] is used in CI
+      # to shorten @shutdown_grace_time from SHUTDOWN_GRACE_TIME. Parallel CI
+      # makes stubbing constants difficult.
+      @shutdown_grace_time = Float(options[:pool_shutdown_grace_time] || SHUTDOWN_GRACE_TIME)
+      @shutdown_debug = options[:shutdown_debug]
+      @block = block
+      @out_of_band = options[:out_of_band]
+      @out_of_band_running = false
+      @out_of_band_condvar = ConditionVariable.new
+      @before_thread_start = options[:before_thread_start]
+      @before_thread_exit = options[:before_thread_exit]
+      @reaping_time = options[:reaping_time]
+      @auto_trim_time = options[:auto_trim_time]
 
       @shutdown = false
 
       @trim_requested = 0
+      @out_of_band_pending = false
 
-      @workers = []
+      @processors = []
 
       @auto_trim = nil
+      @reaper = nil
 
       @mutex.synchronize do
-        min.times { spawn_thread }
+        @min.times do
+          spawn_thread
+          @not_full.wait(@mutex)
+        end
       end
+
+      @force_shutdown = false
+      @shutdown_mutex = Mutex.new
     end
 
-    attr_reader :spawned, :trim_requested
+    attr_reader :spawned, :trim_requested, :waiting
+    attr_accessor :min, :max
+
+    # generate stats hash so as not to perform multiple locks
+    # @return [Hash] hash containing stat info from ThreadPool
+    def stats
+      with_mutex do
+        temp = @backlog_max
+        @backlog_max = 0
+        { backlog: @todo.size,
+          running: @spawned,
+          pool_capacity: pool_capacity,
+          busy_threads: @spawned - @waiting + @todo.size,
+          io_threads: @processors.count(&:marked_as_io_thread?),
+          backlog_max: temp
+        }
+      end
+    end
+
+    def reset_max
+      with_mutex { @backlog_max = 0 }
+    end
 
     # How many objects have yet to be processed by the pool?
     #
     def backlog
-      @mutex.synchronize { @todo.size }
+      with_mutex { @todo.size }
+    end
+
+    # The maximum size of the backlog
+    #
+    def backlog_max
+      with_mutex { @backlog_max }
+    end
+
+    # @!attribute [r] pool_capacity
+    def pool_capacity
+      (waiting + (@max - spawned)).clamp(0, Float::INFINITY)
+    end
+
+    # @!attribute [r] busy_threads
+    # @version 5.0.0
+    def busy_threads
+      with_mutex { @spawned - @waiting + @todo.size }
     end
 
     # :nodoc:
@@ -53,92 +188,210 @@ module Puma
     def spawn_thread
       @spawned += 1
 
-      th = Thread.new do
+      trigger_before_thread_start_hooks
+      processor = ProcessorThread.new(self)
+      processor.thread = Thread.new(processor, @spawned) do |processor, spawned|
+        Puma.set_thread_name '%s tp %03i' % [@name, spawned]
+        # Advertise server into the thread
+        Thread.current.puma_server = @server
+
         todo  = @todo
         block = @block
         mutex = @mutex
-        cond  = @cond
-
-        extra = @extra.map { |i| i.new }
+        not_empty = @not_empty
+        not_full = @not_full
 
         while true
           work = nil
 
-          continue = true
-
           mutex.synchronize do
+            if processor.marked_as_io_thread?
+              if @processors.count { |t| !t.marked_as_io_thread? } < @max
+                # We're not at max processor threads, so the io thread can rejoin the normal population.
+                processor.marked_as_io_thread = false
+              else
+                # We're already at max threads, so we exit the extra io thread.
+                @processors.delete(processor)
+                trigger_before_thread_exit_hooks
+                Thread.exit
+              end
+            end
+
             while todo.empty?
               if @trim_requested > 0
                 @trim_requested -= 1
-                continue = false
-                break
-              end
-
-              if @shutdown
-                continue = false
-                break
+                @spawned -= 1
+                @processors.delete(processor)
+                not_full.signal
+                trigger_before_thread_exit_hooks
+                Thread.exit
               end
 
               @waiting += 1
-              cond.wait mutex
-              @waiting -= 1
+              if @out_of_band_pending && trigger_out_of_band_hook
+                @out_of_band_pending = false
+              end
+              not_full.signal
+              begin
+                not_empty.wait mutex
+              ensure
+                @waiting -= 1
+              end
             end
 
-            work = todo.pop if continue
+            work = todo.shift
           end
 
-          break unless continue
+          begin
+            @out_of_band_pending = true if block.call(processor, work)
+          rescue Exception => e
+            STDERR.puts "Error reached top of thread-pool: #{e.message} (#{e.class})"
+          end
+        end
+      end
+
+      @processors << processor
 
-          block.call(work, *extra)
+      processor
+    end
+
+    private :spawn_thread
+
+    def trigger_before_thread_start_hooks
+      return unless @before_thread_start&.any?
+
+      @before_thread_start.each do |b|
+        begin
+          b[:block].call(ServerPluginControl.new(@server))
+        rescue Exception => e
+          STDERR.puts "WARNING before_thread_start hook failed with exception (#{e.class}) #{e.message}"
         end
+      end
+      nil
+    end
 
-        mutex.synchronize do
-          @spawned -= 1
-          @workers.delete th
+    private :trigger_before_thread_start_hooks
+
+    def trigger_before_thread_exit_hooks
+      return unless @before_thread_exit&.any?
+
+      @before_thread_exit.each do |b|
+        begin
+          b[:block].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?
+
+      # we execute on idle hook when all threads are free
+      return false unless @spawned == @waiting
+      @out_of_band_running = true
+      @out_of_band.each { |b| b[:block].call }
+      true
+    rescue Exception => e
+      STDERR.puts "Exception calling out_of_band_hook: #{e.message} (#{e.class})"
+      true
+    ensure
+      @out_of_band_running = false
+      @out_of_band_condvar.broadcast
+    end
+
+    private :trigger_out_of_band_hook
+
+    def wait_while_out_of_band_running
+      return unless @out_of_band_running
 
-      @workers << th
+      with_mutex do
+        @out_of_band_condvar.wait(@mutex) while @out_of_band_running
+      end
+    end
 
-      th
+    # @version 5.0.0
+    def with_mutex(&block)
+      @mutex.owned? ?
+        yield :
+        @mutex.synchronize(&block)
     end
 
-    private :spawn_thread
+    # :nodoc:
+    #
+    # Must be called with @mutex held!
+    #
+    def can_spawn_processor?
+      io_processors_count = @processors.count(&:marked_as_io_thread?)
+      extra_io_processors_count = io_processors_count > @max_io_threads ? io_processors_count - @max_io_threads : 0
+      (@spawned - io_processors_count) < (@max - extra_io_processors_count)
+    end
 
     # Add +work+ to the todo list for a Thread to pickup and process.
     def <<(work)
-      @mutex.synchronize do
+      with_mutex do
         if @shutdown
           raise "Unable to add work while shutting down"
         end
 
         @todo << work
+        t = @todo.size
+        @backlog_max = t if t > @backlog_max
 
-        if @waiting == 0 and @spawned < @max
+        if @waiting < @todo.size and can_spawn_processor?
           spawn_thread
         end
 
-        @cond.signal
+        @not_empty.signal
       end
+      self
     end
 
-    # If too many threads are in the pool, tell one to finish go ahead
+    def spawn_thread_if_needed # :nodoc:
+      with_mutex do
+        if @waiting < @todo.size and can_spawn_processor?
+          spawn_thread
+        end
+      end
+    end
+
+    # If there are any free threads in the pool, tell one to go ahead
     # and exit. If +force+ is true, then a trim request is requested
     # even if all threads are being utilized.
     #
     def trim(force=false)
-      @mutex.synchronize do
-        if (force or @waiting > 0) and @spawned - @trim_requested > @min
+      with_mutex do
+        free = @waiting - @todo.size
+        if (force or free > 0) and @spawned - @trim_requested > @min
           @trim_requested += 1
-          @cond.signal
+          @not_empty.signal
         end
       end
     end
 
-    class AutoTrim
-      def initialize(pool, timeout)
+    # If there are dead threads in the pool make them go away while decreasing
+    # spawned counter so that new healthy threads could be created again.
+    def reap
+      with_mutex do
+        @processors, dead_processors = @processors.partition(&:alive?)
+
+        dead_processors.each do |processor|
+          processor.kill
+          @spawned -= 1
+        end
+      end
+    end
+
+    class Automaton
+      def initialize(pool, timeout, thread_name, message)
         @pool = pool
         @timeout = timeout
+        @thread_name = thread_name
+        @message = message
         @running = false
       end
 
@@ -146,8 +399,9 @@ module Puma
         @running = true
 
         @thread = Thread.new do
+          Puma.set_thread_name @thread_name
           while @running
-            @pool.trim
+            @pool.public_send(@message)
             sleep @timeout
           end
         end
@@ -159,27 +413,105 @@ module Puma
       end
     end
 
-    def auto_trim!(timeout=5)
-      @auto_trim = AutoTrim.new(self, timeout)
+    def auto_trim!(timeout=@auto_trim_time)
+      @auto_trim = Automaton.new(self, timeout, "#{@name} tp trim", :trim)
       @auto_trim.start!
     end
 
+    def auto_reap!(timeout=@reaping_time)
+      @reaper = Automaton.new(self, timeout, "#{@name} tp reap", :reap)
+      @reaper.start!
+    end
+
+    # Allows ThreadPool::ForceShutdown to be raised within the
+    # provided block if the thread is forced to shutdown during execution.
+    def with_force_shutdown
+      t = Thread.current
+      @shutdown_mutex.synchronize do
+        raise ForceShutdown if @force_shutdown
+        t[:with_force_shutdown] = true
+      end
+      yield
+    ensure
+      t[:with_force_shutdown] = false
+    end
+
     # Tell all threads in the pool to exit and wait for them to finish.
+    # Wait +timeout+ seconds then raise +ForceShutdown+ in remaining threads.
+    # Next, wait an extra +@shutdown_grace_time+ seconds then force-kill remaining
+    # threads. Finally, wait 1 second for remaining threads to exit.
     #
-    def shutdown
-      @mutex.synchronize do
+    def shutdown(timeout)
+      threads = with_mutex do
         @shutdown = true
-        @cond.broadcast
+        @trim_requested = @spawned
+        @not_empty.broadcast
+        @not_full.broadcast
+
+        @auto_trim&.stop
+        @reaper&.stop
+        # dup processors so that we join them all safely
+        @processors.dup
+      end
 
-        @auto_trim.stop if @auto_trim
+      if @shutdown_debug == true
+        shutdown_debug("Shutdown initiated")
       end
 
-      # Use this instead of #each so that we don't stop in the middle
-      # of each and see a mutated object mid #each
-      @workers.first.join until @workers.empty?
+      if timeout == -1
+        # Wait for threads to finish without force shutdown.
+        threads.each(&:join)
+      else
+        join = ->(inner_timeout) do
+          start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          threads.reject! do |t|
+            remaining = inner_timeout - (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start)
+            remaining > 0 && t.join(remaining)
+          end
+        end
+
+        # Wait +timeout+ seconds for threads to finish.
+        join.call(timeout)
+        if @shutdown_debug == :on_force && !threads.empty?
+          shutdown_debug("Shutdown timeout exceeded")
+        end
+
+        # If threads are still running, raise ForceShutdown and wait to finish.
+        @shutdown_mutex.synchronize do
+          @force_shutdown = true
+          threads.each do |t|
+            t.raise ForceShutdown if t[:with_force_shutdown]
+          end
+        end
+        join.call(@shutdown_grace_time)
+        if @shutdown_debug == :on_force && !threads.empty?
+          shutdown_debug("Shutdown grace timeout exceeded")
+        end
+
+        # If threads are _still_ running, forcefully kill them and wait to finish.
+        threads.each(&:kill)
+        join.call(1)
+      end
 
       @spawned = 0
-      @workers = []
+      @processors = []
+    end
+
+    private
+
+    def shutdown_debug(message)
+      pid = Process.pid
+      threads = Thread.list
+
+      $stdout.syswrite "#{pid}: #{message}\n"
+      $stdout.syswrite "#{pid}: === Begin thread backtrace dump ===\n"
+
+      threads.each_with_index do |thread, index|
+        $stdout.syswrite "#{pid}: Thread #{index + 1}/#{threads.size}: #{thread.inspect}\n"
+        $stdout.syswrite "#{pid}: #{(thread.backtrace || []).join("\n#{pid}: ")}\n\n"
+      end
+
+      $stdout.syswrite "#{pid}: === End thread backtrace dump ===\n"
     end
   end
 end

data/lib/puma/util.rb

@@ -1,3 +1,7 @@
+# frozen_string_literal: true
+
+require 'uri/common'
+
 module Puma
   module Util
     module_function
@@ -5,5 +9,126 @@ module Puma
     def pipe
       IO.pipe
     end
+
+    # Escapes and unescapes a URI escaped string with
+    # +encoding+. +encoding+ will be the target encoding of the string
+    # returned, and it defaults to UTF-8
+    if defined?(::Encoding)
+      def escape(s, encoding = Encoding::UTF_8)
+        URI.encode_www_form_component(s, encoding)
+      end
+
+      def unescape(s, encoding = Encoding::UTF_8)
+        URI.decode_www_form_component(s, encoding)
+      end
+    else
+      def escape(s, encoding = nil)
+        URI.encode_www_form_component(s, encoding)
+      end
+
+      def unescape(s, encoding = nil)
+        URI.decode_www_form_component(s, encoding)
+      end
+    end
+    module_function :unescape, :escape
+
+    DEFAULT_SEP = /[&;] */n
+
+    # Stolen from Mongrel, with some small modifications:
+    # Parses a query string by breaking it up at the '&'
+    # and ';' characters.  You can also use this to parse
+    # cookies by changing the characters used in the second
+    # parameter (which defaults to '&;').
+    def parse_query(qs, d = nil, &unescaper)
+      unescaper ||= method(:unescape)
+
+      params = {}
+
+      (qs || '').split(d ? /[#{d}] */n : DEFAULT_SEP).each do |p|
+        next if p.empty?
+        k, v = p.split('=', 2).map(&unescaper)
+
+        if cur = params[k]
+          if cur.class == Array
+            params[k] << v
+          else
+            params[k] = [cur, v]
+          end
+        else
+          params[k] = v
+        end
+      end
+
+      params
+    end
+
+    # A case-insensitive Hash that preserves the original case of a
+    # header when set.
+    class HeaderHash < Hash
+      def self.new(hash={})
+        HeaderHash === hash ? hash : super(hash)
+      end
+
+      def initialize(hash={})
+        super()
+        @names = {}
+        hash.each { |k, v| self[k] = v }
+      end
+
+      def each
+        super do |k, v|
+          yield(k, v.respond_to?(:to_ary) ? v.to_ary.join("\n") : v)
+        end
+      end
+
+      # @!attribute [r] to_hash
+      def to_hash
+        hash = {}
+        each { |k,v| hash[k] = v }
+        hash
+      end
+
+      def [](k)
+        super(k) || super(@names[k.downcase])
+      end
+
+      def []=(k, v)
+        canonical = k.downcase
+        delete k if @names[canonical] && @names[canonical] != k # .delete is expensive, don't invoke it unless necessary
+        @names[k] = @names[canonical] = k
+        super k, v
+      end
+
+      def delete(k)
+        canonical = k.downcase
+        result = super @names.delete(canonical)
+        @names.delete_if { |name,| name.downcase == canonical }
+        result
+      end
+
+      def include?(k)
+        @names.include?(k) || @names.include?(k.downcase)
+      end
+
+      alias_method :has_key?, :include?
+      alias_method :member?, :include?
+      alias_method :key?, :include?
+
+      def merge!(other)
+        other.each { |k, v| self[k] = v }
+        self
+      end
+
+      def merge(other)
+        hash = dup
+        hash.merge! other
+      end
+
+      def replace(other)
+        clear
+        other.each { |k, v| self[k] = v }
+        self
+      end
+    end
   end
 end