Added PoolingConnection for non thread-safe connection classes

elasticsearch/connection/memcached.py

@@ -6,9 +6,9 @@ except ImportError:
     from urllib.parse import urlencode
 
 from ..exceptions import TransportError, ConnectionError
-from .base import Connection
+from .pooling import PoolingConnection
 
-class MemcachedConnection(Connection):
+class MemcachedConnection(PoolingConnection):
     transport_schema = 'memcached'
 
     method_map = {
@@ -25,9 +25,10 @@ class MemcachedConnection(Connection):
         except ImportError:
             raise ImproperlyConfigured("You need to install pylibmc to use the MemcachedConnection class.")
         super(MemcachedConnection, self).__init__(host=host, port=port, **kwargs)
-        self.mc = pylibmc.Client(['%s:%s' % (host, port)],behaviors={"tcp_nodelay": True})
+        self._make_connection = lambda: pylibmc.Client(['%s:%s' % (host, port)], behaviors={"tcp_nodelay": True})
 
     def perform_request(self, method, url, params=None, body=None, timeout=None):
+        mc = self._get_connection()
         url = self.url_prefix + url
         if params:
             url = '%s?%s' % (url, urlencode(params or {}))
@@ -41,10 +42,10 @@ class MemcachedConnection(Connection):
             if mc_method == 'set':
                 # no response from set commands
                 response = ''
-                if not json.dumps(self.mc.set(url, body)):
+                if not json.dumps(mc.set(url, body)):
                     status = 500
             else:
-                response = self.mc.get(url)
+                response = mc.get(url)
 
             duration = time.time() - start
             if response:
@@ -52,6 +53,8 @@ class MemcachedConnection(Connection):
         except Exception as e:
             self.log_request_fail(method, full_url, time.time() - start, exception=e)
             raise ConnectionError('N/A', str(e), e)
+        finally:
+            self._release_connection(mc)
 
         # try not to load the json every time
         if response and response[0] == '{' and ('"status"' in response or '"error"' in response):

elasticsearch/connection/pooling.py

@@ -0,0 +1,22 @@
+from .base import Connection
+
+class PoolingConnection(Connection):
+    def __init__(self, *args, **kwargs):
+        self._max_pool_size = kwargs.pop('max_connection_pool_size', 50)
+        self._free_connections = []
+        self._in_use_connections = set()
+        super(PoolingConnection, self).__init__(*args, **kwargs)
+
+    def _get_connection(self):
+        try:
+            con = self._free_connections.pop()
+        except IndexError:
+            con = self._make_connection()
+
+        self._in_use_connections.add(con)
+        return con
+
+    def _release_connection(self, con):
+        self._in_use_connections.remove(con)
+        self._free_connections.append(con)
+

elasticsearch/connection/thrift.py

@@ -12,9 +12,9 @@ except ImportError:
     THRIFT_AVAILABLE = False
 
 from ..exceptions import ConnectionError
-from .base import Connection
+from .pooling import PoolingConnection
 
-class ThriftConnection(Connection):
+class ThriftConnection(PoolingConnection):
     transport_schema = 'thrift'
 
     def __init__(self, host='localhost', port=9500, framed_transport=False, **kwargs):
@@ -22,9 +22,13 @@ class ThriftConnection(Connection):
             raise ImproperlyConfigured("Thrift is not available.")
 
         super(ThriftConnection, self).__init__(host=host, port=port, **kwargs)
-        socket = TSocket.TSocket(host, port)
+        self._framed_transport = framed_transport
+        self._tsocket_args = (host, port)
+
+    def _make_connection(self):
+        socket = TSocket.TSocket(*self._tsocket_args)
         socket.setTimeout(self.timeout * 1000.0)
-        if framed_transport:
+        if self._framed_transport:
             transport = TTransport.TFramedTransport(socket)
         else:
             transport = TTransport.TBufferedTransport(socket)
@@ -32,20 +36,22 @@ class ThriftConnection(Connection):
         protocol = TBinaryProtocol.TBinaryProtocolAccelerated(transport)
         client = Rest.Client(protocol)
         transport.open()
-        self.tclient = client
-        self.ttransport = transport
+        return client
 
     def perform_request(self, method, url, params=None, body=None, timeout=None):
         request = RestRequest(method=Method._NAMES_TO_VALUES[method.upper()], uri=url,
                     parameters=params, body=body)
 
         start = time.time()
+        tclient = self._get_connection()
         try:
-            response = self.tclient.execute(request)
+            response = tclient.execute(request)
             duration = time.time() - start
         except TException as e:
             self.log_request_fail(method, url, time.time() - start, exception=e)
             raise ConnectionError('N/A', str(e), e)
+        finally:
+            self._release_connection(tclient)
 
         if not (200 <= response.status < 300):
             self.log_request_fail(method, url, duration, response.status)