pool.rb

module RDB

  class Pool
    include Singleton

    def initialize
      @pool_size = (ENV['RDB_POOL_SIZE'] || 10).to_i
      @pool_timeout = (ENV['RDB_POOL_TIMEOUT'] || 5).to_i
    end

    def the_pool
      @pool ||= ConnectionPool.new(size: @pool_size, timeout: @pool_timeout) { self.class.connect }
    end

    def shutdown
      if ENV['DEBUG_RDB'] || RDB.logger.level == 0
        msg = "++ Shutting down connection due to failure".red
        puts msg if ENV['DEBUG_RDB']
        RDB.logger.debug msg if RDB.logger.level == 0
      end

      temp_pool = @pool
      @pool = nil
      temp_pool.shutdown { |conn| conn.close } if temp_pool.present?
      the_pool
    end

    # @todo: probably a better way to delegate this...
    def with
      the_pool.with do |conn|
        begin
          yield(conn) if block_given?
        rescue StandardError => e
          self.shutdown if is_connection_error_exception?(e)
          raise e
        end
      end
    end

    # Taken from NoBrainer
    # https://github.com/nviennot/nobrainer/blob/master/lib/no_brainer/query_runner/reconnect.rb
    def is_connection_error_exception?(e)
      case e
      when Errno::ECONNREFUSED, Errno::EHOSTUNREACH, Errno::EPIPE,
        Errno::ECONNRESET, Errno::ETIMEDOUT, IOError, Timeout::Error
        true
      when RethinkDB::RqlError
        e.message =~ /lost contact/ ||
        e.message =~ /(P|p)rimary .* not available/||
        e.message =~ /Connection.*closed/ ||
        e.message =~ /execution expired/ ||
        e.message =~ /Connection.*refused/
      else
        false
      end
    end

    def stats
      the_pool.stats
    end

    def self.connection_params
      @db_options ||= begin
        db_uri ||= ENV['RETHINKDB_URL'] ||  ENV['RETHINKDB_URI']
        uri = URI.parse(db_uri)
        raise "RETHINKDB connection url is not a rethinkdb url!" unless uri.scheme.downcase == 'rethinkdb'
        result = {
          host:     uri.host || 'localhost',
          port:     uri.port || 28015,
          db:       uri.path[1..-1] || 'test',
          username: uri.user || 'admin',
          password: uri.password,
        }

        result = result.compact
        result.symbolize_keys!
        result.freeze
      end
    end

    def self.connect
      if ENV['DEBUG_RDB'] || RDB.logger.level == 0
        msg = "++ Adding connection to: rethinkdb://#{self.connection_params[:host]}:#{self.connection_params[:port]}/#{self.connection_params[:db]}".light_blue
        puts msg if ENV['DEBUG_RDB']
        RDB.logger.debug msg if RDB.logger.level == 0
      end
      RethinkDB::RQL.new.connect(self.connection_params)
    end
  end
end