Checkpoint, better pool class

Author: andrewvc

elasticsearch-transport/elasticsearch-transport.gemspec

@@ -48,7 +48,7 @@ Gem::Specification.new do |s|
   s.add_development_dependency "patron" unless defined? JRUBY_VERSION
   s.add_development_dependency "typhoeus", '~> 0.6'
   s.add_development_dependency "net-http-persistent"
-  s.add_development_dependency "manticore", '~> 0.3.5' if defined? JRUBY_VERSION
+  s.add_development_dependency "manticore", '~> 0.5.2' if defined? JRUBY_VERSION
   s.add_development_dependency "hashie"
 
   # Prevent unit test failures on Ruby 1.8

elasticsearch-transport/lib/elasticsearch/transport/transport/base.rb

@@ -33,7 +33,7 @@ def initialize(arguments={}, &block)
           @hosts       = arguments[:hosts]   || []
           @options     = arguments[:options] || {}
           @block       = block
-          @connections = __build_connections
+          @connections = Connections::Collection.new(:connections => __build_connections)
 
           @serializer  = options[:serializer] || ( options[:serializer_class] ? options[:serializer_class].new(self) : DEFAULT_SERIALIZER_CLASS.new(self) )
           @protocol    = options[:protocol] || DEFAULT_PROTOCOL

elasticsearch-transport/lib/elasticsearch/transport/transport/connections/collection.rb

@@ -1,6 +1,14 @@
 require 'concurrent'
 require 'thread'
 
+
+# 1 manticore instance
+# Our own connection manager
+#   Pick the least in-use connection
+#   attempt to resurrect a dead connection if none are alive
+#   in the background periodically attempt to resurrect dead connections
+#
+
 module Elasticsearch
   module Transport
     module Transport
@@ -25,9 +33,12 @@ class Collection
           #
           def initialize(arguments={})
             selector_class = arguments[:selector_class] || DEFAULT_SELECTOR
-            @connections = Concurrent::Map.new
-            @connections   = arguments[:connections]    || []
-            @selector      = arguments[:selector]       || selector_class.new(arguments.merge(:connections => self))
+
+            @connection_checkouts = {}
+            (arguments[:connections] || []).each do |conn|
+              @connection_checkouts[conn] = 0
+            end
+
             @state_mutex = Mutex.new
           end
 
@@ -37,14 +48,18 @@ def initialize(arguments={})
           def update(new_connections)
             new_connections.each {|connection| add_connection(connection)}
 
-            @connections.each do |connection,_|
-              @connections.delete(connection) unless new_connections.include?(connection)
-              connection.close
+            @state_mutex.synchronize do
+              @connections.each do |connection,_|
+                @connections.delete(connection) unless new_connections.include?(connection)
+                connection.close
+              end
             end
           end
 
           def add_connection(connection)
-            @connections.put_if_absent(connection, true)
+            @state_mutex.synchronize do
+              @connections[connection] = 0 unless @connections[connection]
+            end
           end
           alias :<< :add_connection
 
@@ -53,15 +68,19 @@ def add_connection(connection)
           # @return [Array]
           #
           def hosts
-            @connections.to_a.map { |c| c.host }
+            @state_mutex.synchronize do
+              @connections.to_a.map { |c| c.host }
+            end
           end
 
           # Returns an Array of alive connections.
           #
           # @return [Array]
           #
           def connections
-            @connections.keys.reject { |c| c.dead? }
+            @state_mutex.synchronize do
+              @connections.keys.reject { |c| c.dead? }
+            end
           end
           alias :alive :connections
 
@@ -70,15 +89,26 @@ def connections
           # @return [Array]
           #
           def dead
-            @connections.keys.select { |c| c.dead? }
+            @state_mutex.synchronize do
+              dead_unsafe
+            end
           end
 
           # Returns an Array of all connections, both dead and alive
           #
           # @return [Array]
           #
           def all
-            @connections.keys
+            @state_mutex.synchronize do
+              @connections.keys
+            end
+          end
+
+          def with_connection(options={})
+            connection = get_connection(options)
+            yield(connection)
+          ensure
+            return_connection(connection) if connection
           end
 
           # Returns a connection.
@@ -89,10 +119,31 @@ def all
           # @return [Connection]
           #
           def get_connection(options={})
-            if connections.empty? && (dead_connection = dead.sort { |a,b| a.failures <=> b.failures }.first)
-              dead_connection.alive!
+            @state_mutex.synchronize do
+              if @connections.empty?
+                if (dead_connection = dead_unsafe.keys.sort { |a,b| a.failures <=> b.failures }.first)
+                  dead_connection.alive!
+                end
+              end
+
+              # It'd be great if we had a concurrent version of a ConcurrentSkipListMap but we don't :(
+              min_kv = @connections.min_by {|k,v| v}
+              if min_kv
+                connection, count = min_kv
+                @connections[connection] = count + 1
+                connection
+              else
+                nil
+              end
+            end
+          end
+
+          def return_connection(connection)
+            @state_mutex.synchronize do
+              if (count = @connections[connection])
+                @connections[connection] = count - 1
+              end
             end
-            selector.select(options)
           end
 
           def each(&block)
@@ -106,8 +157,12 @@ def [](key)
           def size
             connections.size
           end
-        end
+          end
 
+          private
+          def dead_unsafe
+            @connections.keys.select { |c| c.dead? }
+          end
       end
     end
   end

elasticsearch-transport/lib/elasticsearch/transport/transport/connections/selector.rb

@@ -1,3 +1,5 @@
+require 'concurrent'
+
 module Elasticsearch
   module Transport
     module Transport
@@ -44,6 +46,10 @@ def select(options={})
           class RoundRobin
             include Base
 
+            def initialize
+              @current_index
+            end
+
             # Returns the next connection from the collection, rotating them in round-robin fashion.
             #
             # @return [Connections::Connection]
@@ -53,6 +59,7 @@ def select(options={})
               @current = !defined?(@current) || @current.nil? ? 0 : @current+1
               @current = 0 if @current >= connections.size
               connections[@current]
+              require 'pry'; binding.pry
             end
           end
 

elasticsearch-transport/lib/elasticsearch/transport/transport/http/manticore.rb

@@ -47,7 +47,7 @@ class Manticore
 
           def initialize(arguments={}, &block)
             @manticore = build_client(arguments[:options] || {})
-            super()
+            super
           end
 
           # Should just be run once at startup
@@ -115,7 +115,9 @@ def __build_connections
           #
           def __close_connections
             @connections.each {|c| c.dead! }
-            @connections.all.each {|c| c.connection.close }
+            # We don't close the actual manticore client, but rather
+            # just let the pool mechanics close it for us.
+            # It uses the pooling functionality from Apache HC https://hc.apache.org/httpcomponents-client-ga/tutorial/html/connmgmt.html
           end
 
           # Returns an array of implementation specific connection errors.

elasticsearch-transport/lib/elasticsearch/transport/transport/pool.rb

@@ -0,0 +1,139 @@
+module Elasticsearch
+  module Transport
+    module Transport
+      class Pool
+        class RequestFailure < StandardError; end
+
+        def initialize(client_wrapper, resurrect_interval)
+          @state_mutex = Mutex.new
+          @url_info = {}
+          @client = client_wrapper
+          @stopping = false
+          @resurrect_interval = resurrect_interval
+          @resurrectionist = start_resurrectionist
+        end
+
+        def start_resurrectionist
+          Thread.new do
+            slept_for = 0
+            loop do
+              sleep 1
+              slept_for += 1
+               loop do
+                break if @state_mutex.synchronize { @stopping }
+                if slept_for >= @resurrect_interval
+                  slept_for = 0
+                  resurrect_dead!
+                end
+              end
+            end
+          end
+        end
+
+        def resurrect_dead!
+          # Try to keep locking granularity low such that we don't affect IO...
+          @state_mutex.synchronize { @url_info.select {|u,m| m[:dead] } }.each do |url,meta|
+            begin
+              @client.perform_request(url, :get, "/")
+              # If no exception was raised it must have succeeded!
+              @state_mutex.synchronize { m[:dead] = false }
+            rescue RequestFailure => e
+              # NOOP, we'll just try this another time...
+            end
+          end
+        end
+
+        def stop_resurrectionist
+          @state_mutex.synchronize { @stopping = true }
+          @resurrectionist.join # Wait for thread to stop
+        end
+
+        def perform_request(method, path, params={}, body=nil, &block)
+          with_connection do |url|
+            @client.perform_request(url, method, path, params={}, body=nil, &block)
+          end
+        end
+
+        def update_urls(new_urls)
+          @state_mutex.synchronize do
+            # Add new connections
+            new_urls.each do |url|
+              add_url(url)
+            end
+
+            # Delete connections not in the new list
+            @url_info.each do |url,_|
+              remove_url(url)
+            end
+          end
+        end
+
+        def add_url(url)
+          @client.add_url(url)
+          @url_info[connection] ||= empty_url_meta
+        end
+
+        def remove_url(url)
+          @client.remove_url(url)
+          @url_info.delete(url) unless new_connections.include?(url)
+        end
+
+        def empty_url_meta
+          {
+            :open => 0,
+            :dead => false
+          }
+        end
+
+        def with_connection
+          url, url_meta = get_connection
+          yield url
+        rescue RequestFailure => e
+          @state_mutex.synchronize do
+            url_meta[:dead] = true
+            url_meta[:last_error] = e
+          end
+        ensure
+          return_connection(connection)
+        end
+
+        def get_connection
+          @state_mutex.synchronize do
+            # The goal here is to pick a random connection from the least-in-use connections
+            # We want some randomness so that we don't hit the same node over and over, but
+            # we also want more 'fair' behavior in the event of high concurrency
+            eligible_set = nil
+            lowest_value_seen = nil
+            @url_info.each do |url,meta|
+              meta_open = meta[:open]
+              next if meta[:dead]
+
+              if lowest_value_seen.nil? || meta_open < lowest_value_seen
+                lowest_value_seen = meta_open
+                eligible_set = [[url, meta]]
+              elsif lowest_value_seen == meta_open
+                eligible_set << [url, meta]
+              end
+            end
+
+            return nil if eligible_set.nil?
+
+            pick_and_meta = eligible_set.sample
+            pick, pick_meta = pick_and_meta
+            pick_meta[:open] += 1
+
+            pick_and_meta
+          end
+        end
+
+        def return_connection(connection)
+          @state_mutex.synchronize do
+            if @url_info[connection] # Guard against the condition where the connection has already been deleted
+              @url_info[connection][:open] -= 1
+            end
+          end
+        end
+      end
+    end
+  end
+end