Threadsafe connection pooling

Author: andrewvc

elasticsearch-transport/elasticsearch-transport.gemspec

@@ -29,6 +29,7 @@ Gem::Specification.new do |s|
 
   s.add_development_dependency "bundler", "> 1"
   s.add_development_dependency "rake"
+  s.add_dependency "concurrent-ruby"
 
   if defined?(RUBY_VERSION) && RUBY_VERSION > '1.9'
     s.add_development_dependency "elasticsearch-extensions"

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

@@ -99,8 +99,7 @@ def __rebuild_connections(arguments={})
           @state_mutex.synchronize do
             @hosts       = arguments[:hosts]    || []
             @options     = arguments[:options]  || {}
-            __close_connections
-            @connections = __build_connections
+            @connections.update(__build_connections)
           end
         end
 

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

@@ -1,3 +1,6 @@
+require 'concurrent'
+require 'thread'
+
 module Elasticsearch
   module Transport
     module Transport
@@ -22,9 +25,28 @@ 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))
+            @state_mutex = Mutex.new
+          end
+
+          # Adds new connections to the internal pool, and closing connections not in the list
+          #
+          # @option arguments [Array] :new_connections An Array of {Connection} objects
+          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
+            end
+          end
+
+          def add_connection(connection)
+            @connections.put_if_absent(connection, true)
           end
+          alias :<< :add_connection
 
           # Returns an Array of hosts information in this collection as Hashes.
           #
@@ -39,7 +61,7 @@ def hosts
           # @return [Array]
           #
           def connections
-            @connections.reject { |c| c.dead? }
+            @connections.keys.reject { |c| c.dead? }
           end
           alias :alive :connections
 
@@ -48,15 +70,15 @@ def connections
           # @return [Array]
           #
           def dead
-            @connections.select { |c| c.dead? }
+            @connections.keys.select { |c| c.dead? }
           end
 
           # Returns an Array of all connections, both dead and alive
           #
           # @return [Array]
           #
           def all
-            @connections
+            @connections.keys
           end
 
           # Returns a connection.
@@ -67,7 +89,7 @@ def all
           # @return [Connection]
           #
           def get_connection(options={})
-            if connections.empty? && dead_connection = dead.sort { |a,b| a.failures <=> b.failures }.first
+            if connections.empty? && (dead_connection = dead.sort { |a,b| a.failures <=> b.failures }.first)
               dead_connection.alive!
             end
             selector.select(options)
@@ -77,10 +99,9 @@ def each(&block)
             connections.each(&block)
           end
 
-          def slice(*args)
-            connections.slice(*args)
+          def [](key)
+            @connections[key]
           end
-          alias :[] :slice
 
           def size
             connections.size

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

@@ -118,6 +118,19 @@ def resurrectable?
             }
           end
 
+          def close
+            self.connection.close
+          end
+
+          def eql?(other)
+            other.instance_of?(self.class) &&  self.host.to_s == other.host.to_s
+          end
+          alias :== :eql?
+
+          def hash
+            self.host.hash
+          end
+
           # @return [String]
           #
           def to_s

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

@@ -97,20 +97,15 @@ def __build_connections
               @request_options[:headers] = options[:headers]
             end
 
-            Connections::Collection.new \
-              :connections => hosts.map { |host|
-                host[:protocol]   = host[:scheme] || DEFAULT_PROTOCOL
-                host[:port]     ||= DEFAULT_PORT
+            hosts.map do |host|
+              host[:protocol]   = host[:scheme] || DEFAULT_PROTOCOL
+              host[:port]     ||= DEFAULT_PORT
 
-                host.delete(:user)     # auth is not supported here.
-                host.delete(:password) # use the headers
+              host.delete(:user)     # auth is not supported here.
+              host.delete(:password) # use the headers
 
-                Connections::Connection.new \
-                  :host => host,
-                  :connection => @manticore
-              },
-              :selector_class => options[:selector_class],
-              :selector => options[:selector]
+             Connections::Connection.new(:host => host, :connection => @manticore)
+            end
           end
 
           # Closes all connections by marking them as dead