Active Record V4 connection pool performance patch

gistfile1.rb

require 'thread_safe'
require 'semantic_logger'
require 'gene_pool'
# Require original code since it contains more than just this class
require 'active_record/connection_adapters/abstract/connection_pool'

module ActiveRecord
  # Not Used. Connection are no longer timed out
  # A warning is logged when the connection time is exceeded
  class ConnectionTimeoutError < ConnectionNotEstablished
  end

  module ConnectionAdapters
    # Connection pool base class for managing Active Record database
    # connections.
    #
    # == Introduction
    #
    # A connection pool synchronizes thread access to a limited number of
    # database connections. The basic idea is that each thread checks out a
    # database connection from the pool, uses that connection, and checks the
    # connection back in. ConnectionPool is completely thread-safe, and will
    # ensure that a connection cannot be used by two threads at the same time,
    # as long as ConnectionPool's contract is correctly followed. It will also
    # handle cases in which there are more threads than connections: if all
    # connections have been checked out, and a thread tries to checkout a
    # connection anyway, then ConnectionPool will wait until some other thread
    # has checked in a connection.
    #
    # == Obtaining (checking out) a connection
    #
    # Connections can be obtained and used from a connection pool in several
    # ways:
    #
    # 1. Simply use ActiveRecord::Base.connection as with Active Record 2.1 and
    #    earlier (pre-connection-pooling). Eventually, when you're done with
    #    the connection(s) and wish it to be returned to the pool, you call
    #    ActiveRecord::Base.clear_active_connections!. This will be the
    #    default behavior for Active Record when used in conjunction with
    #    Action Pack's request handling cycle.
    # 2. Manually check out a connection from the pool with
    #    ActiveRecord::Base.connection_pool.checkout. You are responsible for
    #    returning this connection to the pool when finished by calling
    #    ActiveRecord::Base.connection_pool.checkin(connection).
    # 3. Use ActiveRecord::Base.connection_pool.with_connection(&block), which
    #    obtains a connection, yields it as the sole argument to the block,
    #    and returns it to the pool after the block completes.
    #
    # Connections in the pool are actually AbstractAdapter objects (or objects
    # compatible with AbstractAdapter's interface).
    #
    # == Options
    #
    # There are several connection-pooling-related options that you can add to
    # your database connection configuration:
    #
    # * +pool+: number indicating size of connection pool (default 5)
    # * +checkout_timeout+: number of seconds to block and wait for a connection
    #   before giving up and raising a timeout error (default 5 seconds).
    # * +reaping_frequency+: frequency in seconds to periodically run the
    #   Reaper, which attempts to find and close dead connections, which can
    #   occur if a programmer forgets to close a connection at the end of a
    #   thread or a thread dies unexpectedly. (Default nil, which means don't
    #   run the Reaper).
    # * +dead_connection_timeout+: number of seconds from last checkout
    #   after which the Reaper will consider a connection reapable. (default
    #   5 seconds).
    class ConnectionPool
      include SemanticLogger::Loggable

      attr_accessor :automatic_reconnect, :checkout_timeout, :dead_connection_timeout
      attr_reader :spec, :size, :reaper


      # TODO Remove later
      attr_reader :pool

      # Creates a new ConnectionPool object. +spec+ is a ConnectionSpecification
      # object which describes database connection information (e.g. adapter,
      # host name, username, password, etc), as well as the maximum size for
      # this ConnectionPool.
      #
      # The default ConnectionPool maximum size is 5.
      def initialize(spec)
        @spec                    = spec
        @checkout_timeout        = spec.config[:checkout_timeout] || 5
        # Log a warning if the checkout takes longer than checkout_warning seconds
        @checkout_warning        = spec.config[:checkout_warning] || 0.1
        @dead_connection_timeout = spec.config[:dead_connection_timeout] || 5
        @reaper                  = Reaper.new self, spec.config[:reaping_frequency]

        # default max pool size to 5
        @size                    = (spec.config[:pool] && spec.config[:pool].to_i) || 5

        # The hash of reserved connections mapped to threads
        @reserved_connections    = ThreadSafe::Hash.new
        @automatic_reconnect     = true

        @pool = GenePool.new(
          name:          self.class.name,
          pool_size:     @size,
          timeout:       @checkout_timeout,
          warn_timeout:  @checkout_warning,
          logger:        logger,
          # Prevent GenePool from automatically closing the connection when it is removed
          close_proc:    nil,
          timeout_class: ConnectionTimeoutError
        ) do
          raise ConnectionNotEstablished unless @automatic_reconnect

          # Create a new database connection
          # Log a warning if it takes longer than 100ms to create the connection
          logger.benchmark_warn("Create Database Connection", min_duration: 100) do
            conn = Base.send(spec.adapter_method, spec.config)
            conn.pool = self
            conn
          end
        end
        @reaper.run
      end

      # Retrieve the connection associated with the current thread, or call
      # #checkout to obtain one if necessary.
      #
      # #connection can be called any number of times; the connection is
      # held in a hash keyed by the thread id.
      def connection
        @reserved_connections[current_connection_id] ||= checkout
      end

      # Is there an open connection that is being used for the current thread?
      def active_connection?
        @reserved_connections.fetch(current_connection_id) {
          return false
        }.in_use?
      end

      # Signal that the thread is finished with the current connection.
      # #release_connection releases the connection-thread association
      # and returns the connection to the pool.
      def release_connection(with_id = current_connection_id)
        conn = @reserved_connections.delete(with_id)
        checkin conn if conn
      end

      # If a connection already exists yield it to the block. If no connection
      # exists checkout a connection, yield it to the block, and checkin the
      # connection when finished.
      def with_connection
        connection_id = current_connection_id
        fresh_connection = true unless active_connection?
        yield connection
      ensure
        release_connection(connection_id) if fresh_connection
      end

      # Return the connections as an array
      def connections
        conns = []
        @pool.each {|c| conns << c}
        conns
      end

      # Returns true if a connection has already been opened.
      def connected?
        @pool.size > 0
      end

      # Disconnects all connections in the pool, and clears the pool.
      def disconnect!
        @pool.each do |conn|
          checkin conn
          conn.disconnect!
        end
      end

      # Clears the cache which maps classes.
      def clear_reloadable_connections!
        @pool.each do |conn|
          if conn.requires_reloading?
            conn.disconnect!
            @pool.remove(conn)
          end
        end
      end

      def clear_stale_cached_connections! # :nodoc:
        reap
      end
      deprecate :clear_stale_cached_connections! => "Please use #reap instead"

      # Check-out a database connection from the pool, indicating that you want
      # to use it. You should call #checkin when you no longer need this.
      #
      # This is done by either returning and leasing existing connection, or by
      # creating a new connection and leasing it.
      #
      # If all connections are leased and the pool is at capacity (meaning the
      # number of currently leased connections is greater than or equal to the
      # size limit set), an ActiveRecord::ConnectionTimeoutError exception will be raised.
      #
      # Returns: an AbstractAdapter object.
      #
      # TODO No longer Raises:
      # - ConnectionTimeoutError: no connection can be obtained from the pool.
      def checkout
        logger.benchmark_warn("#checkout", min_duration: (@checkout_warning / 1000.0)) do
          conn = @pool.checkout

          # Mark connection as in-use and set last used to now
          conn.lease

          conn.run_callbacks :checkout do
            # Connection must be verified every time since it may have been returned
            # in a bad state
            conn.verify!
          end
          conn
        end
      end

      # Check-in a database connection back into the pool, indicating that you
      # no longer need this connection.
      #
      # +conn+: an AbstractAdapter object, which was obtained by earlier by
      # calling +checkout+ on this pool.
      def checkin(conn)
        logger.benchmark_warn("#checkin", min_duration: (@checkout_warning / 1000.0)) do
          conn.run_callbacks :checkin do
            conn.expire
          end

          release_reservation(conn)
          @pool.checkin(conn)
        end
      end

      # Remove a connection from the connection pool.  The connection will
      # remain open and active but will no longer be managed by this pool.
      def remove(conn)
        release_reservation(conn)
        @pool.remove(conn)
      end

      # Removes dead connections from the pool.  A dead connection can occur
      # if a programmer forgets to close a connection at the end of a thread
      # or a thread dies unexpectedly.
      def reap
        stale = Time.now - @dead_connection_timeout
        @pool.each do |conn|
          remove(conn) if conn.in_use? && stale > conn.last_use && !conn.active?
        end
      end

      private

      # Releases a connection from it's reserved thread
      def release_reservation(conn)
        thread_id = if @reserved_connections[current_connection_id] == conn
          current_connection_id
        else
          @reserved_connections.keys.find { |k| @reserved_connections[k] == conn }
        end

        @reserved_connections.delete(thread_id) if thread_id
      end

      def current_connection_id #:nodoc:
        Base.connection_id ||= Thread.current.object_id
      end
    end
  end
end

In order to use this patch, include the following gems in your rails app:

• gene_pool
• rails_semantic_logger