1 files changed, 283 insertions, 78 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index d69f02d504..7a3d9bfd3e 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -11,9 +11,6 @@ module ActiveRecord
 
   # Raised when a connection pool is full and another connection is requested
   class PoolFullError < ConnectionNotEstablished
-    def initialize size, timeout
-      super("Connection pool of size #{size} and timeout #{timeout}s is full")
-    end
   end
 
   module ConnectionAdapters
@@ -57,13 +54,146 @@ module ActiveRecord
     #
     # == Options
     #
-    # There are two connection-pooling-related options that you can add to
+    # 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)
-    # * +wait_timeout+: number of seconds to block and wait for a connection
+    # * +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
+      # Threadsafe, fair, FIFO queue.  Meant to be used by ConnectionPool
+      # with which it shares a Monitor.  But could be a generic Queue.
+      #
+      # The Queue in stdlib's 'thread' could replace this class except
+      # stdlib's doesn't support waiting with a timeout.
+      class Queue
+        def initialize(lock = Monitor.new)
+          @lock = lock
+          @cond = @lock.new_cond
+          @num_waiting = 0
+          @queue = []
+        end
+
+        # Test if any threads are currently waiting on the queue.
+        def any_waiting?
+          synchronize do
+            @num_waiting > 0
+          end
+        end
+
+        # Return the number of threads currently waiting on this
+        # queue.
+        def num_waiting
+          synchronize do
+            @num_waiting
+          end
+        end
+
+        # Add +element+ to the queue.  Never blocks.
+        def add(element)
+          synchronize do
+            @queue.push element
+            @cond.signal
+          end
+        end
+
+        # If +element+ is in the queue, remove and return it, or nil.
+        def delete(element)
+          synchronize do
+            @queue.delete(element)
+          end
+        end
+
+        # Remove all elements from the queue.
+        def clear
+          synchronize do
+            @queue.clear
+          end
+        end
+
+        # Remove the head of the queue.
+        #
+        # If +timeout+ is not given, remove and return the head the
+        # queue if the number of available elements is strictly
+        # greater than the number of threads currently waiting (that
+        # is, don't jump ahead in line).  Otherwise, return nil.
+        #
+        # If +timeout+ is given, block if it there is no element
+        # available, waiting up to +timeout+ seconds for an element to
+        # become available.
+        #
+        # Raises:
+        # - ConnectionTimeoutError if +timeout+ is given and no element
+        # becomes available after +timeout+ seconds,
+        def poll(timeout = nil)
+          synchronize do
+            if timeout
+              no_wait_poll || wait_poll(timeout)
+            else
+              no_wait_poll
+            end
+          end
+        end
+
+        private
+
+        def synchronize(&block)
+          @lock.synchronize(&block)
+        end
+
+        # Test if the queue currently contains any elements.
+        def any?
+          !@queue.empty?
+        end
+
+        # A thread can remove an element from the queue without
+        # waiting if an only if the number of currently available
+        # connections is strictly greater than the number of waiting
+        # threads.
+        def can_remove_no_wait?
+          @queue.size > @num_waiting
+        end
+
+        # Removes and returns the head of the queue if possible, or nil.
+        def remove
+          @queue.shift
+        end
+
+        # Remove and return the head the queue if the number of
+        # available elements is strictly greater than the number of
+        # threads currently waiting.  Otherwise, return nil.
+        def no_wait_poll
+          remove if can_remove_no_wait?
+        end
+
+        # Waits on the queue up to +timeout+ seconds, then removes and
+        # returns the head of the queue.
+        def wait_poll(timeout)
+          @num_waiting += 1
+
+          t0 = Time.now
+          elapsed = 0
+          loop do
+            @cond.wait(timeout - elapsed)
+
+            return remove if any?
+
+            elapsed = Time.now - t0
+            raise ConnectionTimeoutError if elapsed >= timeout
+          end
+        ensure
+          @num_waiting -= 1
+        end
+      end
+
       # Every +frequency+ seconds, the reaper will call +reap+ on +pool+.
       # A reaper instantiated with a nil frequency will never reap the
       # connection pool.
@@ -91,7 +221,7 @@ module ActiveRecord
 
       include MonitorMixin
 
-      attr_accessor :automatic_reconnect, :timeout
+      attr_accessor :automatic_reconnect, :checkout_timeout, :dead_connection_timeout
       attr_reader :spec, :connections, :size, :reaper
 
       # Creates a new ConnectionPool object. +spec+ is a ConnectionSpecification
@@ -108,7 +238,8 @@ module ActiveRecord
         # The cache of reserved connections mapped to threads
         @reserved_connections = {}
 
-        @timeout = spec.config[:wait_timeout] || 5
+        @checkout_timeout = spec.config[:checkout_timeout] || 5
+        @dead_connection_timeout = spec.config[:dead_connection_timeout]
         @reaper  = Reaper.new self, spec.config[:reaping_frequency]
         @reaper.run
 
@@ -117,6 +248,16 @@ module ActiveRecord
 
         @connections         = []
         @automatic_reconnect = true
+
+        @available = Queue.new self
+      end
+
+      # Hack for tests to be able to add connections.  Do not call outside of tests
+      def insert_connection_for_test!(c) #:nodoc:
+        synchronize do
+          @connections << c
+          @available.add c
+        end
       end
 
       # Retrieve the connection associated with the current thread, or call
@@ -125,21 +266,28 @@ module ActiveRecord
       # #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
+        synchronize do
+          @reserved_connections[current_connection_id] ||= checkout
+        end
       end
 
-      # Check to see if there is an active connection in this connection
-      # pool.
+      # Is there an open connection that is being used for the current thread?
       def active_connection?
-        active_connections.any?
+        synchronize do
+          @reserved_connections.fetch(current_connection_id) {
+            return false
+          }.in_use?
+        end
       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
+        synchronize do
+          conn = @reserved_connections.delete(with_id)
+          checkin conn if conn
+        end
       end
 
       # If a connection already exists yield it to the block. If no connection
@@ -167,6 +315,7 @@ module ActiveRecord
             conn.disconnect!
           end
           @connections = []
+          @available.clear
         end
       end
 
@@ -181,22 +330,17 @@ module ActiveRecord
           @connections.delete_if do |conn|
             conn.requires_reloading?
           end
-        end
-      end
-
-      # Verify active connections and remove and disconnect connections
-      # associated with stale threads.
-      def verify_active_connections! #:nodoc:
-        synchronize do
-          @connections.each do |connection|
-            connection.verify!
+          @available.clear
+          @connections.each do |conn|
+            @available.add conn
           end
         end
       end
 
       def clear_stale_cached_connections! # :nodoc:
+        reap
       end
-      deprecate :clear_stale_cached_connections!
+      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.
@@ -213,23 +357,10 @@ module ActiveRecord
       # Raises:
       # - PoolFullError: no connection can be obtained from the pool.
       def checkout
-        # Checkout an available connection
         synchronize do
-          # Try to find a connection that hasn't been leased, and lease it
-          conn = connections.find { |c| c.lease }
-
-          # If all connections were leased, and we have room to expand,
-          # create a new connection and lease it.
-          if !conn && connections.size < size
-            conn = checkout_new_connection
-            conn.lease
-          end
-
-          if conn
-            checkout_and_verify conn
-          else
-            raise PoolFullError.new(size, timeout)
-          end
+          conn = acquire_connection
+          conn.lease
+          checkout_and_verify(conn)
         end
       end
 
@@ -243,6 +374,10 @@ module ActiveRecord
           conn.run_callbacks :checkin do
             conn.expire
           end
+
+          release conn
+
+          @available.add conn
         end
       end
 
@@ -251,13 +386,13 @@ module ActiveRecord
       def remove(conn)
         synchronize do
           @connections.delete conn
+          @available.delete conn
 
           # FIXME: we might want to store the key on the connection so that removing
           # from the reserved hash will be a little easier.
-          thread_id = @reserved_connections.keys.find { |k|
-            @reserved_connections[k] == conn
-          }
-          @reserved_connections.delete thread_id if thread_id
+          release conn
+
+          @available.add checkout_new_connection if @available.any_waiting?
         end
       end
 
@@ -266,7 +401,7 @@ module ActiveRecord
       # or a thread dies unexpectedly.
       def reap
         synchronize do
-          stale = Time.now - @timeout
+          stale = Time.now - @dead_connection_timeout
           connections.dup.each do |conn|
             remove conn if conn.in_use? && stale > conn.last_use && !conn.active?
           end
@@ -275,12 +410,49 @@ module ActiveRecord
 
       private
 
+      # Acquire a connection by one of 1) immediately removing one
+      # from the queue of available connections, 2) creating a new
+      # connection if the pool is not at capacity, 3) waiting on the
+      # queue for a connection to become available.
+      #
+      # Raises:
+      # - PoolFullError if a connection could not be acquired (FIXME:
+      #   why not ConnectionTimeoutError?
+      def acquire_connection
+        if conn = @available.poll
+          conn
+        elsif @connections.size < @size
+          checkout_new_connection
+        else
+          t0 = Time.now
+          begin
+            @available.poll(@checkout_timeout)
+          rescue ConnectionTimeoutError
+            msg = 'could not obtain a database connection within %0.3f seconds (waited %0.3f seconds)' %
+              [@checkout_timeout, Time.now - t0]
+            raise PoolFullError, msg
+          end
+        end
+      end
+
+      def release(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 new_connection
-        ActiveRecord::Base.send(spec.adapter_method, spec.config)
+        ActiveRecord::Model.send(spec.adapter_method, spec.config)
       end
 
       def current_connection_id #:nodoc:
-        ActiveRecord::Base.connection_id ||= Thread.current.object_id
+        ActiveRecord::Model.connection_id ||= Thread.current.object_id
       end
 
       def checkout_new_connection
@@ -298,10 +470,6 @@ module ActiveRecord
         end
         c
       end
-
-      def active_connections
-        @connections.find_all { |c| c.in_use? }
-      end
     end
 
     # ConnectionHandler is a collection of ConnectionPool objects. It is used
@@ -328,43 +496,40 @@ module ActiveRecord
     # ActiveRecord::Base.connection_handler. Active Record models use this to
     # determine that connection pool that they should use.
     class ConnectionHandler
-      attr_reader :connection_pools
+      def initialize
+        @owner_to_pool = Hash.new { |h,k| h[k] = {} }
+        @class_to_pool = Hash.new { |h,k| h[k] = {} }
+      end
 
-      def initialize(pools = {})
-        @connection_pools = pools
-        @class_to_pool    = {}
+      def connection_pools
+        owner_to_pool.values.compact
       end
 
-      def establish_connection(name, spec)
-        @connection_pools[spec] ||= ConnectionAdapters::ConnectionPool.new(spec)
-        @class_to_pool[name] = @connection_pools[spec]
+      def establish_connection(owner, spec)
+        @class_to_pool.clear
+        owner_to_pool[owner] = ConnectionAdapters::ConnectionPool.new(spec)
       end
 
       # Returns true if there are any active connections among the connection
       # pools that the ConnectionHandler is managing.
       def active_connections?
-        connection_pools.values.any? { |pool| pool.active_connection? }
+        connection_pools.any?(&:active_connection?)
       end
 
       # Returns any connections in use by the current thread back to the pool,
       # and also returns connections to the pool cached by threads that are no
       # longer alive.
       def clear_active_connections!
-        @connection_pools.each_value {|pool| pool.release_connection }
+        connection_pools.each(&:release_connection)
       end
 
       # Clears the cache which maps classes.
       def clear_reloadable_connections!
-        @connection_pools.each_value {|pool| pool.clear_reloadable_connections! }
+        connection_pools.each(&:clear_reloadable_connections!)
       end
 
       def clear_all_connections!
-        @connection_pools.each_value {|pool| pool.disconnect! }
-      end
-
-      # Verify active connections.
-      def verify_active_connections! #:nodoc:
-        @connection_pools.each_value {|pool| pool.verify_active_connections! }
+        connection_pools.each(&:disconnect!)
       end
 
       # Locate the connection of the nearest super class. This can be an
@@ -387,21 +552,61 @@ module ActiveRecord
       # connection and the defined connection (if they exist). The result
       # can be used as an argument for establish_connection, for easily
       # re-establishing the connection.
-      def remove_connection(klass)
-        pool = @class_to_pool.delete(klass.name)
-        return nil unless pool
-
-        @connection_pools.delete pool.spec
-        pool.automatic_reconnect = false
-        pool.disconnect!
-        pool.spec.config
+      def remove_connection(owner)
+        if pool = owner_to_pool.delete(owner)
+          @class_to_pool.clear
+          pool.automatic_reconnect = false
+          pool.disconnect!
+          pool.spec.config
+        end
       end
 
+      # Retrieving the connection pool happens a lot so we cache it in @class_to_pool.
+      # This makes retrieving the connection pool O(1) once the process is warm.
+      # When a connection is established or removed, we invalidate the cache.
+      #
+      # Ideally we would use #fetch here, as class_to_pool[klass] may sometimes be nil.
+      # However, benchmarking (https://gist.github.com/3552829) showed that #fetch is
+      # significantly slower than #[]. So in the nil case, no caching will take place,
+      # but that's ok since the nil case is not the common one that we wish to optimise
+      # for.
       def retrieve_connection_pool(klass)
-        pool = @class_to_pool[klass.name]
-        return pool if pool
-        return nil if ActiveRecord::Model == klass
-        retrieve_connection_pool klass.active_record_super
+        class_to_pool[klass] ||= begin
+          until pool = pool_for(klass)
+            klass = klass.superclass
+            break unless klass < Model::Tag
+          end
+
+          class_to_pool[klass] = pool || pool_for(ActiveRecord::Model)
+        end
+      end
+
+      private
+
+      def owner_to_pool
+        @owner_to_pool[Process.pid]
+      end
+
+      def class_to_pool
+        @class_to_pool[Process.pid]
+      end
+
+      def pool_for(owner)
+        owner_to_pool.fetch(owner) {
+          if ancestor_pool = pool_from_any_process_for(owner)
+            # A connection was established in an ancestor process that must have
+            # subsequently forked. We can't reuse the connection, but we can copy
+            # the specification and establish a new connection with it.
+            establish_connection owner, ancestor_pool.spec
+          else
+            owner_to_pool[owner] = nil
+          end
+        }
+      end
+
+      def pool_from_any_process_for(owner)
+        owner_to_pool = @owner_to_pool.values.find { |v| v[owner] }
+        owner_to_pool && owner_to_pool[owner]
       end
     end