<commit_before>require_relative 'connection_pool/version'
require_relative 'connection_pool/timed_stack'

# Generic connection pool class for e.g. sharing a limited number of network connections
# among many threads.  Note: Connections are eager created.
#
# Example usage with block (faster):
#
#    @pool = ConnectionPool.new { Redis.new }
#
#    @pool.with do |redis|
#      redis.lpop('my-list') if redis.llen('my-list') > 0
#    end
#
# Using optional timeout override (for that single invocation)
#
#    @pool.with(:timeout => 2.0) do |redis|
#      redis.lpop('my-list') if redis.llen('my-list') > 0
#    end
#
# Example usage replacing an existing connection (slower):
#
#    $redis = ConnectionPool.wrap { Redis.new }
#
#    def do_work
#      $redis.lpop('my-list') if $redis.llen('my-list') > 0
#    end
#
# Accepts the following options:
# - :size - number of connections to pool, defaults to 5
# - :timeout - amount of time to wait for a connection if none currently available, defaults to 5 seconds
#
class ConnectionPool
  DEFAULTS = {size: 5, timeout: 5}

  class Error < RuntimeError
  end

  def self.wrap(options, &block)
    Wrapper.new(options, &block)
  end

  def initialize(options = {}, &block)
    raise ArgumentError, 'Connection pool requires a block' unless block

    options = DEFAULTS.merge(options)

    @size = options.fetch(:size)
    @timeout = options.fetch(:timeout)

    @available = TimedStack.new(@size, &block)
    @key = :"current-#{@available.object_id}"
  end

  def with(options = {})
    success = false # hoisted
    conn = checkout(options)
    begin
      (yield conn).tap do
        success = true # means the connection wasn't interrupted
      end
    ensure
      if success
        # everything is roses, we can safely check the connection back in
        checkin
      else
        @available.discard!(pop_connection)
      end
    end
  end

  def checkout(options = {})
    conn = if stack.empty?
      timeout = options[:timeout] || @timeout
      @available.pop(timeout: timeout)
    else
      stack.last
    end

    stack.push conn
    conn
  end

  def checkin
    conn = pop_connection # mutates stack, must be on its own line
    @available.push(conn) if stack.empty?

    nil
  end

  def shutdown(&block)
    @available.shutdown(&block)<commit_msg>Remove tap usage, comment this necessary complexity<commit_after>require_relative 'connection_pool/version'
require_relative 'connection_pool/timed_stack'

# Generic connection pool class for e.g. sharing a limited number of network connections
# among many threads.  Note: Connections are eager created.
#
# Example usage with block (faster):
#
#    @pool = ConnectionPool.new { Redis.new }
#
#    @pool.with do |redis|
#      redis.lpop('my-list') if redis.llen('my-list') > 0
#    end
#
# Using optional timeout override (for that single invocation)
#
#    @pool.with(:timeout => 2.0) do |redis|
#      redis.lpop('my-list') if redis.llen('my-list') > 0
#    end
#
# Example usage replacing an existing connection (slower):
#
#    $redis = ConnectionPool.wrap { Redis.new }
#
#    def do_work
#      $redis.lpop('my-list') if $redis.llen('my-list') > 0
#    end
#
# Accepts the following options:
# - :size - number of connections to pool, defaults to 5
# - :timeout - amount of time to wait for a connection if none currently available, defaults to 5 seconds
#
class ConnectionPool
  DEFAULTS = {size: 5, timeout: 5}

  class Error < RuntimeError
  end

  def self.wrap(options, &block)
    Wrapper.new(options, &block)
  end

  def initialize(options = {}, &block)
    raise ArgumentError, 'Connection pool requires a block' unless block

    options = DEFAULTS.merge(options)

    @size = options.fetch(:size)
    @timeout = options.fetch(:timeout)

    @available = TimedStack.new(@size, &block)
    @key = :"current-#{@available.object_id}"
  end

  def with(options = {})
    # Connections can become corrupted via Timeout::Error.  Discard
    # any connection whose usage after checkout does not finish as expected.
    # See #67
    success = false
    conn = checkout(options)
    begin
      result = yield conn
      success = true # means the connection wasn't interrupted
      result
    ensure
      if success
        # everything is roses, we can safely check the connection back in
        checkin
      else
        @available.discard!(pop_connection)
      end
    end
  end

  def checkout(options = {})
    conn = if stack.empty?
      timeout = options[:timeout] || @timeout
      @available.pop(timeout: timeout)
    else
      stack.last
    end

    stack.push conn
    conn
  end

  def checkin
    conn = pop_connection # mutates stack, must be on its own line
    @available.push(conn) if stack.empty?

    nil
  end

  def shutdown(&block)
    @available.shutdown(&block)
{'fraction_non_alphanumeric': 0.1053941908713693, 'fraction_numerical': 0.004149377593360996, 'mean_word_length': 3.0208507089241032, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2449045', 'n_tokens_mistral': 1500, 'n_tokens_neox': 1447, 'n_words': 468}