em-pg-client 0.3.2 → 0.3.3

data/lib/pg/em/connection_pool.rb

@@ -78,14 +78,30 @@ module PG
       # - +:lazy+ = false - should lazy allocate first connection
       # - +:connection_class+ = {PG::EM::Client}
       #
+      # For convenience the given block will be set as the +on_connect+ option.
+      #
+      # @yieldparam pg [Client] connected client instance on each newly
+      #                         created connection
+      # @yieldparam is_async [Boolean] always +true+ in a connection pool
+      #                      context
+      # @yieldparam is_reset [Boolean] always +false+ unless
+      #                      +async_autoreconnect+ options is +true+ and
+      #                      was actually re-connecting
+      #
       # @raise [PG::Error]
       # @raise [ArgumentError]
-      def initialize(options = {})
+      # @see Client#on_connect
+      def initialize(options = {}, &on_connect)
         @available = []
         @pending = []
         @allocated = {}
+        @max_size = DEFAULT_SIZE
         @connection_class = Client
 
+        if block_given?
+          options = {on_connect: on_connect}.merge(options)
+        end
+
         lazy = false
         @options = options.reject do |key, value|
           case key.to_sym
@@ -101,9 +117,7 @@ module PG
           end
         end
 
-        @max_size ||= DEFAULT_SIZE
-
-        raise ArgumentError, "#{self.class}.new: pool size must be > 1" if @max_size < 1
+        raise ArgumentError, "#{self.class}.new: pool size must be >= 1" if @max_size < 1
 
         # allocate first connection, unless we are lazy
         hold unless lazy
@@ -179,6 +193,10 @@ module PG
       #   @return [Boolean] asynchronous auto re-connect status
       #   Set {Client#async_autoreconnect}  on all present and future connections
       #   in this pool or read value from options
+      # @!attribute [rw] on_connect
+      #   @return [Proc<Client,is_async,is_reset>] connect hook
+      #   Set {Client#on_connect} on all present and future connections
+      #   in this pool or read value from options
       # @!attribute [rw] on_autoreconnect
       #   @return [Proc<Client, Error>] auto re-connect hook
       #   Set {Client#on_autoreconnect} on all present and future connections
@@ -186,6 +204,7 @@ module PG
       %w[connect_timeout
          query_timeout
          async_autoreconnect
+         on_connect
          on_autoreconnect].each do |name|
         class_eval <<-EOD, __FILE__, __LINE__
           def #{name}=(value)
@@ -194,11 +213,24 @@ module PG
             @available.each(&b)
             @allocated.each_value(&b)
           end
-
-          def #{name}
-            @options[:#{name}] || @options['#{name}']
-          end
         EOD
+        if name.start_with?('on_')
+          class_eval <<-EOD, __FILE__, __LINE__
+            def #{name}(&hook)
+              if block_given?
+                self.#{name} = hook
+              else
+                @options[:#{name}] || @options['#{name}']
+              end
+            end
+          EOD
+        else
+          class_eval <<-EOD, __FILE__, __LINE__
+            def #{name}
+              @options[:#{name}] || @options['#{name}']
+            end
+          EOD
+        end
         DeferredOptions.class_eval <<-EOD, __FILE__, __LINE__
           def #{name}=(value)
             self[:#{name}=] = value
@@ -327,7 +359,7 @@ module PG
             begin
               id = fiber.object_id
               # mark allocated pool for proper #size value
-              # the connecting process will yield from fiber
+              # the connection is made asynchronously
               @allocated[id] = opts = DeferredOptions.new
               conn = @connection_class.new(@options)
             ensure

data/spec/em_client_autoreconnect.rb

@@ -285,6 +285,51 @@ describe 'pg-em autoreconnect with on_autoreconnect' do
     end
   end
 
+  it "should execute on_connect before on_autoreconnect after server restart" do
+    @client.on_connect.should be_nil
+    run_on_connect = false
+    @client.on_connect = proc do |client, is_async, is_reset|
+      client.should be_an_instance_of PG::EM::Client
+      is_async.should be_true
+      is_reset.should be_true
+      client.query_defer('SELECT pg_database_size(current_database());').callback {
+        run_on_connect = true
+      }
+    end
+    @client.on_autoreconnect = proc do |client, ex|
+      run_on_connect.should be_true
+      @on_autoreconnect.call(client, ex)
+    end
+    system($pgserver_cmd_stop).should be_true
+    system($pgserver_cmd_start).should be_true
+    @tested_proc.call
+  end
+
+  it "should skip on_autoreconnect when on_connect failed after server restart" do
+    run_on_connect = false
+    run_on_autoreconnect = false
+    @client.on_connect = proc do |client, is_async, is_reset|
+      client.should be_an_instance_of PG::EM::Client
+      is_async.should be_true
+      is_reset.should be_true
+      client.query_defer('SELLECT 1;').errback {
+        run_on_connect = true
+      }
+    end
+    @client.on_autoreconnect = proc do |client, ex|
+      run_on_autoreconnect = true
+    end
+    system($pgserver_cmd_stop).should be_true
+    system($pgserver_cmd_start).should be_true
+    @client.exec_prepared_defer('get_db_size') do |ex|
+      ex.should be_an_instance_of PG::SyntaxError
+      @client.status.should be PG::CONNECTION_OK
+      run_on_connect.should be_true
+      run_on_autoreconnect.should be_false
+      EM.stop
+    end.should be_a_kind_of ::EM::DefaultDeferrable
+  end
+
   before(:all) do
     @tested_proc = proc do
       @client.exec_prepared_defer('get_db_size') do |result|

data/spec/em_client_on_connect.rb

@@ -0,0 +1,171 @@
+$:.unshift "lib"
+gem 'eventmachine', '~> 1.0.0'
+gem 'pg', ENV['EM_PG_CLIENT_TEST_PG_VERSION']
+require 'eventmachine'
+require 'em-synchrony'
+require 'pg/em'
+module PGSpecMacros
+  def test_blocking
+    client = subject.new(options)
+    client.should be_an_instance_of subject
+    client.on_connect.should be on_connect
+    client.exec_prepared(query_name) do |result|
+      result.should be_an_instance_of PG::Result
+      result[0]['pg_database_size'].to_i.should be > 0
+    end
+    client.reset
+    client.exec_prepared(query_name) do |result|
+      result.should be_an_instance_of PG::Result
+      result[0]['pg_database_size'].to_i.should be > 0
+    end
+  end
+
+  def test_blocking_error
+    expect do
+      subject.new(options)
+    end.to raise_error(on_connect_exception)
+    client = subject.new
+    client.should be_an_instance_of subject
+    client.on_connect = on_connect
+    expect do
+      client.reset
+    end.to raise_error(on_connect_exception)
+  end
+end
+
+RSpec.configure do |config|
+  config.include(PGSpecMacros)
+end
+
+shared_context 'test on_connect' do
+
+  it "should invoke on_connect after deferrable connect and reset" do
+    EM.run do
+      subject.connect_defer(options) do |client|
+        client.should be_an_instance_of subject
+        client.on_connect.should be on_connect
+        client.exec_prepared_defer(query_name) do |result|
+          result.should be_an_instance_of PG::Result
+          result[0]['pg_database_size'].to_i.should be > 0
+          client.reset_defer do |client|
+            client.should be_an_instance_of subject
+            client.exec_prepared_defer(query_name) do |result|
+              result.should be_an_instance_of PG::Result
+              result[0]['pg_database_size'].to_i.should be > 0
+            end.should be_a_kind_of ::EM::DefaultDeferrable
+            EM.stop
+          end.should be_a_kind_of ::EM::DefaultDeferrable
+        end.should be_a_kind_of ::EM::DefaultDeferrable
+      end.should be_a_kind_of ::EM::DefaultDeferrable
+    end
+  end
+
+  it "should invoke on_connect after synchrony connect and reset" do
+    EM.synchrony do
+      test_blocking
+      EM.stop
+    end
+  end
+
+end
+
+shared_context 'test on_connect error' do
+
+  it "should fail on_connect with exception after deferrable connect and reset" do
+    EM.run do
+      subject.connect_defer(options) do |ex|
+        ex.should be_an_instance_of on_connect_exception
+        subject.connect_defer do |client|
+          client.should be_an_instance_of subject
+          client.on_connect = on_connect
+          client.reset_defer do |ex|
+            ex.should be_an_instance_of on_connect_exception
+            EM.stop
+          end.should be_a_kind_of ::EM::DefaultDeferrable
+        end.should be_a_kind_of ::EM::DefaultDeferrable
+      end.should be_a_kind_of ::EM::DefaultDeferrable
+    end
+  end
+
+  it "should fail on_connect with exception after synchrony connect and reset" do
+    EM.synchrony do
+      test_blocking_error
+      EM.stop
+    end
+  end
+end
+
+shared_context 'test blocking' do
+  it "should invoke on_connect after blocking connect and reset" do
+    test_blocking
+  end
+end
+
+shared_context 'test blocking on_connect error' do
+  it "should fail on_connect with exception after blocking connect and reset" do
+    test_blocking_error
+  end
+end
+
+describe 'on_connect option' do
+  subject          { PG::EM::Client }
+  let(:query_name) { 'get_db_size' }
+  let(:query)      { 'SELECT pg_database_size(current_database());' }
+  let(:options)    { {on_connect: on_connect} }
+  let(:sleep_query){ 'SELECT pg_sleep(0.1)'}
+  let(:on_connect_exception) { Class.new(StandardError) }
+
+
+  describe 'with deferrable on_connect' do
+    let(:on_connect) { proc {|client, is_async|
+      is_async.should be_true
+      PG::EM::FeaturedDeferrable.new.tap do |df|
+        client.exec_defer(sleep_query).callback do
+          df.bind_status client.prepare_defer(query_name, query)
+        end.errback { df.fail on_connect_exception }
+      end
+    } }
+
+    include_context 'test on_connect'
+  end
+
+  describe 'with synchrony on_connect' do
+    let(:on_connect) { proc {|client, is_async|
+      is_async.should be_true
+      was_async = false
+      EM.next_tick { was_async = true }
+      client.exec(sleep_query)
+      client.prepare(query_name, query)
+      was_async.should be_true
+    } }
+
+    include_context 'test on_connect'
+  end
+
+  describe 'with blocking on_connect' do
+    let(:on_connect) { proc {|client, is_async|
+      is_async.should be_false
+      client.prepare(query_name, query)
+    } }
+
+    include_context 'test blocking'
+  end
+
+  describe 'with error raised in on_connect' do
+    let(:on_connect) { proc {|client|
+      raise on_connect_exception
+    } }
+
+    include_context 'test on_connect error'
+    include_context 'test blocking on_connect error'
+  end
+
+  describe 'with on_connect deferrable failure' do
+    let(:on_connect) { proc {|client|
+      EM::DefaultDeferrable.new.tap {|df| df.fail on_connect_exception.new }
+    } }
+
+    include_context 'test on_connect error'
+  end
+
+end

data/spec/em_connection_pool.rb

@@ -0,0 +1,198 @@
+$:.unshift "lib"
+gem 'eventmachine', '~> 1.0.0'
+gem 'pg', ENV['EM_PG_CLIENT_TEST_PG_VERSION']
+require 'eventmachine'
+require 'em-synchrony'
+require 'em-synchrony/fiber_iterator'
+require 'pg/em/connection_pool'
+
+shared_context 'test on_connect' do
+
+  it 'should call prepared statement concurrently with synchrony' do
+    results = []
+    pool = subject.new(options)
+    pool.max_size.should eq concurrency
+    pool.size.should eq 1
+    start = Time.now
+    EM::Synchrony::FiberIterator.new((1..concurrency), concurrency).each do |index|
+      pool.exec_prepared(query_name) do |result|
+        result.should be_an_instance_of PG::Result
+        result[0]['pg_database_size'].to_i.should be > 0
+      end
+      pool.query(sleep_query).should be_an_instance_of PG::Result
+      results << index
+    end
+    delta = Time.now - start
+    delta.should be_between(sleep_interval, sleep_interval * concurrency / 2)
+    results.sort.should eq (1..concurrency).to_a
+    pool.size.should eq concurrency
+    EM.stop
+  end
+
+  it 'should call prepared statement concurrently with deferrable' do
+    results = []
+    subject.connect_defer(options) do |pool|
+      pool.max_size.should eq concurrency
+      pool.size.should eq 1
+      start = Time.now
+      concurrency.times do |index|
+        pool.exec_prepared_defer(query_name) do |result|
+          result.should be_an_instance_of PG::Result
+          result[0]['pg_database_size'].to_i.should be > 0
+          pool.query_defer(sleep_query) do |result|
+            result.should be_an_instance_of PG::Result
+            results << index
+            if results.length == concurrency
+              delta = Time.now - start
+              delta.should be_between(sleep_interval, sleep_interval * concurrency / 2)
+              results.sort.should eq (0...concurrency).to_a
+              pool.size.should eq concurrency
+              EM.stop
+            end
+          end
+        end
+      end
+    end
+  end
+end
+
+shared_context 'test on_connect error' do
+
+  it "should fail on_connect with exception after synchrony connect" do
+    expect do
+      subject.new(options)
+    end.to raise_error(on_connect_exception)
+    client = subject.new(options.merge(lazy: true))
+    client.should be_an_instance_of subject
+    expect do
+      client.query(sleep_query)
+    end.to raise_error(on_connect_exception)
+    EM.stop
+  end
+
+  it "should fail on_connect with exception after deferrable connect" do
+    subject.connect_defer(options) do |ex|
+      ex.should be_an_instance_of on_connect_exception
+      pool = subject.new(options.merge(lazy: true))
+      pool.should be_an_instance_of subject
+      pool.query_defer(sleep_query) do |ex|
+        ex.should be_an_instance_of on_connect_exception
+        EM.stop
+      end.should be_an_instance_of PG::EM::FeaturedDeferrable
+    end.should be_an_instance_of PG::EM::FeaturedDeferrable
+  end
+
+end
+
+describe 'connection pool' do
+  subject { PG::EM::ConnectionPool }
+
+  describe 'on_connect' do
+    let(:query_name)     { 'get_db_size' }
+    let(:query)          { 'SELECT pg_database_size(current_database());' }
+    let(:concurrency)    { 10 }
+    let(:options)        { {size: concurrency, on_connect: on_connect} }
+    let(:sleep_interval) { 0.1 }
+    let(:sleep_query)    { "SELECT pg_sleep(#{sleep_interval})"}
+    let(:on_connect_exception) { Class.new(StandardError) }
+    let(:on_connect)     { proc {} }
+
+    around(:each) do |testcase|
+      EM.synchrony do
+        begin
+          testcase.call
+        end
+      end
+    end
+
+    it 'should setup block as on_connect client option' do
+      connect_hook = false
+      pool = subject.new { connect_hook = true }
+      connect_hook.should be_true
+      pool.should be_an_instance_of subject
+      pool.on_connect.should be_an_instance_of Proc
+      EM.stop
+    end
+
+    it 'should prefer on_connect from options' do
+      connect_hook = false
+      pool = subject.new(options) { connect_hook = true }
+      connect_hook.should be_false
+      pool.should be_an_instance_of subject
+      pool.on_connect.should be on_connect
+      EM.stop
+    end
+
+    describe 'with deferrable on_connect' do
+      let(:on_connect)     { proc {|client, is_async|
+        is_async.should be_true
+        PG::EM::FeaturedDeferrable.new.tap do |df|
+          client.exec_defer(sleep_query).callback do
+            df.bind_status client.prepare_defer(query_name, query)
+          end.errback { df.fail on_connect_exception }
+        end
+      } }
+
+      include_context 'test on_connect'
+    end
+
+    describe 'with synchrony on_connect' do
+      let(:on_connect) { proc {|client, is_async|
+        is_async.should be_true
+        was_async = false
+        EM.next_tick { was_async = true }
+        client.exec(sleep_query)
+        client.prepare(query_name, query)
+        was_async.should be_true
+      } }
+
+      include_context 'test on_connect'
+    end
+
+    describe 'with error raised in on_connect' do
+      let(:on_connect) { proc {|client|
+        raise on_connect_exception
+      } }
+
+      include_context 'test on_connect error'
+    end
+
+    describe 'with on_connect deferrable failure' do
+      let(:on_connect) { proc {|client|
+        EM::DefaultDeferrable.new.tap {|df| df.fail on_connect_exception.new }
+      } }
+
+      include_context 'test on_connect error'
+    end
+  end
+
+  describe '#transaction' do
+    let(:concurrency)    { 2 }
+    let(:options)        { {size: concurrency} }
+    let(:query)          { 'SELECT pg_database_size(current_database());' }
+
+    around(:each) do |testcase|
+      EM.synchrony do
+        begin
+          @pool = subject.new(options)
+          testcase.call
+        end
+      end
+    end
+
+    it 'should lock transaction connection to fiber' do
+      @pool.transaction do |pg|
+        @pool.hold {|c| c.should be pg }
+        Fiber.new do
+          @pool.size.should eq 1
+          @pool.hold {|c| c.should_not be pg }
+          @pool.size.should eq 2
+          EM.stop
+        end.resume
+        @pool.hold {|c| c.should be pg }
+        @pool.size.should eq 2
+      end
+    end
+
+  end
+end