Ticket #2162: connection_pool_test.rb

require 'abstract_unit'
require 'fixtures/topic'
require 'monitor'

# todo change configuration to use a hash within a hash

module ActiveRecord
  class Base
    def self.adapter_pool_connection(config)
      ConnectionAdapters::AdapterPool.new(config)
    end
    
    def self.mock_connection(config)
      ConnectionAdapters::MockAdapter.new(config)
    end
  end
  
  class ConnectionAdapters::MockAdapter < ::ActiveRecord::ConnectionAdapters::AbstractAdapter
    attr_writer :active
    
    def self.live_connection_count
      @@live_connection_count ||= 0
    end
    
    def initialize(config)
      @@live_connection_count = self.class.live_connection_count + 1
      @active = true
    end
    
    def pretend_to_do_something(&block)
      yield if block_given?
    end
    
    def raise_connection_failed
      raise ActiveRecord::ConnectionFailed
    end
    
    def active?
      @active
    end
  end

  class ConnectionAdapters::AdapterPool
    include MonitorMixin
  
    def initialize(config)
      super()
      @config = config
      @max_signal = new_cond
      start_reaper_thread 
    end
    
    def kind_of?(klass)
      klass == ConnectionAdapters::AdapterPool || method_missing(:kind_of?, klass)
    end
    
    # Pass all undefined methods through to a local instance of the underlying
    # database adapter
    
    def method_missing(method, *args, &block)
      use_local_connection do |connection|
        connection.send(method, *args, &block)
      end
    end
 
    def current_local_connection
      Thread.current[:current_local_connection]
    end    

    def idle_connections
      @idle_connections ||= []
    end
    
    def all_connections
      @all_connections ||= []
    end

    # ensure there aren't too many/too few idle connections

    def rebalance_pool
      synchronize do
        while max_idle && max_idle < idle_connections.size
          remove_connection(idle_connections.shift)
        end
        
        while min_idle && min_idle > idle_connections.size
          idle_connections << add_connection
        end
      end
    end
    
    # remove any stale connections from the pool
    
    def refresh_pool
      synchronize do
        rejected = idle_connections.reject {|c| c.active? }
        rejected.each {|c| remove_connection c }
      end
    end
    
    def refresh_and_rebalance_pool
      synchronize do
        refresh_pool
        rebalance_pool
      end
    end
    
    def reaper_thread
      @reaper_thread
    end
    
    private
      def max_idle
        @config[:max_idle]
      end
      
      def min_idle
        @config[:min_idle]
      end
      
      def max_count
        @config[:max_count]
      end
      
      def reaper_siesta_period
        @config[:reaper_siesta] || 30
      end
    
      def add_connection
        connection = ActiveRecord::Base.send("#{@config[:underlying_adapter]}_connection", @config)
        all_connections << connection
        connection
      end
      
      def remove_connection(connection)
        # @TODO we should clean this up in some more meaningful way before discarding it?
        synchronize do
          all_connections.delete(connection)
          idle_connections.delete(connection)
        end
        
        if current_local_connection == connection
          Thread.current[:current_local_connection] = nil
        end
      end
      
      # Uses a local connection for a block of work

      def use_local_connection(&block)
        if !current_local_connection
          acquire_and_use_local_connection(&block)
        else
          yield current_local_connection
        end
      end
      
      # Aquire a new local connection and use it for the given block of work
      
      def acquire_and_use_local_connection(&block)
        begin
          Thread.current[:current_local_connection]=(acquire_connection)
          yield current_local_connection
        rescue ActiveRecord::ConnectionFailed
          # if the connection has failed, then reset it (so it isn't returned to the pool)
          remove_connection(current_local_connection)
          raise 
        ensure
          if current_local_connection
            # return the connection back to the pool
            synchronize do
              idle_connections << Thread.current[:current_local_connection]
              @max_signal.signal
            end
          end
          Thread.current[:current_local_connection]= nil
        end
      end
      
      def acquire_connection
        synchronize do
          unless result = idle_connections.pop
            if max_count.nil? || all_connections.size < max_count
              result = add_connection
            else
              @max_signal.wait
            end
          end

          result || acquire_connection
        end
      end
      
      def start_reaper_thread
        unless @config[:prevent_reaper_thread]
          @reaper_thread = Thread.new do
            Thread.current[:continue] = true
            while Thread.current[:continue]
              sleep(reaper_siesta_period)
              refresh_and_rebalance_pool
            end
          end
        end    
      end
  end
end

class AdapterPoolWithLiveConnectionTest < Test::Unit::TestCase
  include MonitorMixin
  
  fixtures :topics

  def setup
    super
    ActiveRecord::Base.allow_concurrency = false
    @original_config = ActiveRecord::Base.remove_connection
    @pooled_config = @original_config.dup
    @pooled_config[:underlying_adapter]= @pooled_config[:adapter]
    @pooled_config[:adapter]= "adapter_pool"
    @pooled_config[:prevent_reaper_thread]= true
    ActiveRecord::Base.establish_connection(@pooled_config)
  end
  
  def teardown
    ActiveRecord::Base.establish_connection(@original_config)
  end
  
  def test_basic_connection_works
    assert_kind_of(ActiveRecord::ConnectionAdapters::AdapterPool, ActiveRecord::Base.connection)
    assert_equal 2, Topic.count
  end
end

class AdapterPoolWithMockConnectionTest < Test::Unit::TestCase
  def setup
    super
    @original_config = ActiveRecord::Base.remove_connection
    @pooled_config = @original_config.dup
    @pooled_config[:underlying_adapter]= "mock"
    @pooled_config[:adapter]= "adapter_pool"
    @pooled_config[:prevent_reaper_thread]= true
    ActiveRecord::Base.establish_connection(@pooled_config)
  end
  
  def teardown
    ActiveRecord::Base.establish_connection(@original_config)
  end
  
  def test_basic_connection_works
    assert_equal(0, live_connection_count)
    
    assert_kind_of(ActiveRecord::ConnectionAdapters::AdapterPool, ActiveRecord::Base.connection)
    
    assert_equal(1, live_connection_count)
    assert_equal(1, idle_connection_count)
    
    ActiveRecord::Base.connection.pretend_to_do_something
    
    assert_equal(1, live_connection_count)
    assert_equal(1, idle_connection_count)
    
    assert ActiveRecord::Base.connection.reaper_thread.nil?
  end
  
  def test_transaction_support
    ActiveRecord::Base.transaction {
      assert_equal(0, idle_connection_count)
      first = current_local_connection
      
      ActiveRecord::Base.transaction {
        assert_equal(0, idle_connection_count)
        assert_equal(1, all_connection_count)
        assert_same first, current_local_connection
      }
    }
    
    assert_equal(1, all_connection_count)
    assert_equal(1, idle_connection_count)
  end
  
  def test_connections_are_pooled_once_used
    ActiveRecord::Base.transaction {@first = current_local_connection}
    ActiveRecord::Base.transaction {@second = current_local_connection}
    assert_same(@first, @second)
  end
 
  def test_connections_are_local_to_current_thread
    ActiveRecord::Base.connection.pretend_to_do_something do
      @first = current_local_connection
      assert_equal(1, all_connection_count)
      assert_equal(0, idle_connection_count)

      with_new_local_connection do   
        @second = current_local_connection
        assert_not_same @first, @second
        assert_equal(2, all_connection_count)
        assert_equal(0, idle_connection_count)
      end

      assert_equal(1, idle_connection_count)
    end
    
    assert_equal(2, idle_connection_count)
  end
  
  def test_connection_failures_invalidate_connection
    ActiveRecord::Base.transaction {@first = current_local_connection}
    assert_raise(ActiveRecord::ConnectionFailed) {ActiveRecord::Base.connection.raise_connection_failed}
    assert_equal(0, idle_connection_count)
    assert_equal(0, all_connection_count)
    ActiveRecord::Base.transaction {@second = current_local_connection}
    assert_not_same(@first, @second, "Original should no longer be available")
  end
  
  def test_rebalance_respects_min_idle
    @pooled_config[:min_idle] = 10
    ActiveRecord::Base.establish_connection(@pooled_config)
    ActiveRecord::Base.connection.rebalance_pool
    assert_equal(10, idle_connection_count)
  end
  
  def test_rebalance_respects_max_idle
    @pooled_config[:max_idle] = 5
    ActiveRecord::Base.establish_connection(@pooled_config)
    
    with_new_local_connection do
      with_new_local_connection do
        with_new_local_connection do
          with_new_local_connection do
            with_new_local_connection do
              with_new_local_connection do
                with_new_local_connection do
                  with_new_local_connection do
                    assert_equal(8, all_connection_count)
                  end
                end
              end
            end
          end
        end
      end
    end
    
    assert_equal(8, idle_connection_count)
    ActiveRecord::Base.connection.rebalance_pool
    assert_equal(5, idle_connection_count)
  end
  
  def test_refresh_culls_inactive_connections
    with_new_local_connection do
      # 1 inactive connection
      current_local_connection.active = false
      with_new_local_connection do
        with_new_local_connection do
          with_new_local_connection do
            # 2 inactive connections 
            current_local_connection.active = false
            with_new_local_connection do
              with_new_local_connection do
                with_new_local_connection do
                  # 3 inactive connections
                  current_local_connection.active = false
                  with_new_local_connection do
                    # 4 inactive connections
                    current_local_connection.active = false
                    assert_equal(8, all_connection_count)
                  end
                end
              end
            end
          end
        end
      end
    end
    
    assert_equal(8, idle_connection_count)
    ActiveRecord::Base.connection.refresh_pool
    assert_equal(4, idle_connection_count)
  end

  def test_max_total_setting
    @pooled_config[:max_count] = 5
    ActiveRecord::Base.establish_connection(@pooled_config)
    threads = []
    last_exception = nil
    
    test = self
    completed = 0
    
    110.times do
      Thread.new do
        begin
          ActiveRecord::Base.connection.pretend_to_do_something do
            100.times {Thread.pass}
          end
        rescue Exception => e
          last_exception = e
        ensure
          completed = completed + 1
        end
      end
    end

    # ensure most threads have completed before continuing

    while completed < 100
      Thread.pass
    end
    
    # join last few remaining threads
    
    Thread.list.each {|t|
      t.join unless t == Thread.current
    }
    
    # check no exceptions were caught
    
    assert last_exception.nil?, last_exception.to_s
    
    # check no more than 5 connections were used, and that all 5 are now currently idle
    
    assert_equal 5, all_connection_count
    assert_equal 5, idle_connection_count
  end
  
  def test_reaper_thread_refreshes_connections
    @pooled_config[:prevent_reaper_thread]= false
    ActiveRecord::Base.establish_connection(@pooled_config)
    with_new_local_connection do
      # 1 inactive connection
      current_local_connection.active = false
      with_new_local_connection do
        with_new_local_connection do
          with_new_local_connection do
            # 2 inactive connections 
            current_local_connection.active = false
            with_new_local_connection do
              with_new_local_connection do
                with_new_local_connection do
                  # 3 inactive connections
                  current_local_connection.active = false
                  with_new_local_connection do
                    # 4 inactive connections
                    current_local_connection.active = false
                    assert_equal(8, all_connection_count)
                  end
                end
              end
            end
          end
        end
      end
    end
    
    ActiveRecord::Base.connection.reaper_thread[:continue] = false
    ActiveRecord::Base.connection.reaper_thread.wakeup
    ActiveRecord::Base.connection.reaper_thread.join
    assert_equal(4, idle_connection_count)
  end
  
  def test_reaper_thread_rebalances_connections
    @pooled_config[:max_idle] = 5
    @pooled_config[:prevent_reaper_thread]= false
    ActiveRecord::Base.establish_connection(@pooled_config)
    with_new_local_connection do
      with_new_local_connection do
        with_new_local_connection do
          with_new_local_connection do
            with_new_local_connection do
              with_new_local_connection do
                with_new_local_connection do
                  with_new_local_connection do
                    assert_equal(8, all_connection_count)
                  end
                end
              end
            end
          end
        end
      end
    end

    ActiveRecord::Base.connection.reaper_thread[:continue] = false
    ActiveRecord::Base.connection.reaper_thread.wakeup
    ActiveRecord::Base.connection.reaper_thread.join
    assert_equal(5, idle_connection_count)    
  end

  private
    def with_new_local_connection(&block)
      Thread.new do 
        ActiveRecord::Base.connection.pretend_to_do_something do   
          yield
        end
      end.join
    end
  
    def current_local_connection
      ActiveRecord::Base.connection.current_local_connection
    end
    
    def live_connection_count
      ActiveRecord::ConnectionAdapters::MockAdapter.live_connection_count
    end
    
    def idle_connection_count
      ActiveRecord::Base.connection.idle_connections.size
    end
    
    def all_connection_count
      ActiveRecord::Base.connection.all_connections.size
    end
end