class Puma::ThreadPool

def <<(work)
  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 < @todo.size and @spawned < @max
      spawn_thread
    end
    @not_empty.signal
  end
end

def auto_reap!(timeout=@reaping_time)
  @reaper = Automaton.new(self, timeout, "#{@name} tp reap", :reap)
  @reaper.start!
end

def auto_trim!(timeout=@auto_trim_time)
  @auto_trim = Automaton.new(self, timeout, "#{@name} tp trim", :trim)
  @auto_trim.start!
end

def backlog
  with_mutex { @todo.size }
end

def backlog_max
  with_mutex { @backlog_max }
end

def busy_threads
  with_mutex { @spawned - @waiting + @todo.size }
end

def initialize(name, options = {}, &block)
  @not_empty = ConditionVariable.new
  @not_full = ConditionVariable.new
  @mutex = Mutex.new
  @todo = Queue.new
  @backlog_max = 0
  @spawned = 0
  @waiting = 0
  @name = name
  @min = Integer(options[:min_threads])
  @max = Integer(options[:max_threads])
  # 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)
  @block = block
  @out_of_band = options[:out_of_band]
  @out_of_band_running = false
  @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 = []
  @auto_trim = nil
  @reaper = nil
  @mutex.synchronize do
    @min.times do
      spawn_thread
      @not_full.wait(@mutex)
    end
  end
  @force_shutdown = false
  @shutdown_mutex = Mutex.new
end

def pool_capacity
  waiting + (@max - spawned)
end

def reap
  with_mutex do
    dead_workers = @workers.reject(&:alive?)
    dead_workers.each do |worker|
      worker.kill
      @spawned -= 1
    end
    @workers.delete_if do |w|
      dead_workers.include?(w)
    end
  end
end

def reset_max
  with_mutex { @backlog_max = 0 }
end

def shutdown(timeout=-1)
  threads = with_mutex do
    @shutdown = true
    @trim_requested = @spawned
    @not_empty.broadcast
    @not_full.broadcast
    @auto_trim&.stop
    @reaper&.stop
    # dup workers so that we join them all safely
    @workers.dup
  end
  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|
        elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
        t.join inner_timeout - elapsed
      end
    end
    # Wait +timeout+ seconds for threads to finish.
    join.call(timeout)
    # 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 threads are _still_ running, forcefully kill them and wait to finish.
    threads.each(&:kill)
    join.call(1)
  end
  @spawned = 0
  @workers = []
end

def spawn_thread
  @spawned += 1
  trigger_before_thread_start_hooks
  th = Thread.new(@spawned) do |spawned|
    Puma.set_thread_name '%s tp %03i' % [@name, spawned]
    todo  = @todo
    block = @block
    mutex = @mutex
    not_empty = @not_empty
    not_full = @not_full
    while true
      work = nil
      mutex.synchronize do
        while todo.empty?
          if @trim_requested > 0
            @trim_requested -= 1
            @spawned -= 1
            @workers.delete th
            not_full.signal
            trigger_before_thread_exit_hooks
            Thread.exit
          end
          @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.shift
      end
      begin
        @out_of_band_pending = true if block.call(work)
      rescue Exception => e
        STDERR.puts "Error reached top of thread-pool: #{e.message} (#{e.class})"
      end
    end
  end
  @workers << th
  th
end

def stats
  with_mutex do
    temp = @backlog_max
    @backlog_max = 0
    { backlog: @todo.size,
      running: @spawned,
      pool_capacity: @waiting + (@max - @spawned),
      busy_threads: @spawned - @waiting + @todo.size,
      backlog_max: temp
    }
  end
end

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

def trigger_before_thread_start_hooks
  return unless @before_thread_start&.any?
  @before_thread_start.each do |b|
    begin
      b[:block].call
    rescue Exception => e
      STDERR.puts "WARNING before_thread_start hook failed with exception (#{e.class}) #{e.message}"
    end
  end
  nil
end

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
end

def trim(force=false)
  with_mutex do
    free = @waiting - @todo.size
    if (force or free > 0) and @spawned - @trim_requested > @min
      @trim_requested += 1
      @not_empty.signal
    end
  end
end

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

def with_mutex(&block)
  @mutex.owned? ?
    yield :
    @mutex.synchronize(&block)
end