From 667160e9dc421f5f8721f385dbc17d9ede3e01c0 Mon Sep 17 00:00:00 2001 From: Honza Kral Date: Sun, 25 Aug 2013 17:00:46 +0200 Subject: [PATCH] Added PoolingConnection for non thread-safe connection classes --- elasticsearch/connection/memcached.py | 13 ++++++++----- elasticsearch/connection/pooling.py | 22 ++++++++++++++++++++++ elasticsearch/connection/thrift.py | 20 +++++++++++++------- 3 files changed, 43 insertions(+), 12 deletions(-) create mode 100644 elasticsearch/connection/pooling.py diff --git a/elasticsearch/connection/memcached.py b/elasticsearch/connection/memcached.py index 17094f0a..153a2505 100644 --- a/elasticsearch/connection/memcached.py +++ b/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): diff --git a/elasticsearch/connection/pooling.py b/elasticsearch/connection/pooling.py new file mode 100644 index 00000000..58641101 --- /dev/null +++ b/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) + diff --git a/elasticsearch/connection/thrift.py b/elasticsearch/connection/thrift.py index 70396cac..7fe96805 100644 --- a/elasticsearch/connection/thrift.py +++ b/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)