From dd19553ba709616272e0255a8d720b5561486ce5 Mon Sep 17 00:00:00 2001 From: Honza Kral Date: Fri, 24 May 2013 01:29:31 +0200 Subject: [PATCH] Have sniff_hosts reuse existing connection to not waste open sockets --- elasticsearch/connection_pool.py | 1 + elasticsearch/transport.py | 17 ++++++++++++++--- test_elasticsearch/test_transport.py | 8 ++++++++ 3 files changed, 23 insertions(+), 3 deletions(-) diff --git a/elasticsearch/connection_pool.py b/elasticsearch/connection_pool.py index 477bde39..70158013 100644 --- a/elasticsearch/connection_pool.py +++ b/elasticsearch/connection_pool.py @@ -94,6 +94,7 @@ class ConnectionPool(object): :arg randomize_hosts: shuffle the list of connections upon arrival to avoid dog piling effect across processes """ + self.connection_opts = connections self.connections = [c for (c, opts) in connections] # PriorityQueue for thread safety and ease of timeout management self.dead = PriorityQueue(len(self.connections)) diff --git a/elasticsearch/transport.py b/elasticsearch/transport.py index 3aa1fc57..41b4fd3a 100644 --- a/elasticsearch/transport.py +++ b/elasticsearch/transport.py @@ -95,19 +95,30 @@ class Transport(object): def set_connections(self, hosts): """ - Instantiate all the connections and crate new connection pool to hold them. + Instantiate all the connections and crate new connection pool to hold + them. Tries to identify unchanged hosts and re-use existing + :class:`~elasticsearch.Connection` instances. :arg hosts: same as `__init__` """ # construct the connections def _create_connection(host): + # if this is not the initial setup look at the existing connection + # options and identify connections that haven't changed and can be + # kept around. + if hasattr(self, 'connection_pool'): + for (connection, old_host) in self.connection_pool.connection_opts: + if old_host == host: + return connection + + # previously unseen params, create new connection kwargs = self.kwargs.copy() kwargs.update(host) return self.connection_class(**kwargs) - connections = list(map(_create_connection, hosts)) + connections = map(_create_connection, hosts) # pass the hosts dicts to the connection pool to optionally extract parameters from - self.connection_pool = self.connection_pool_class(zip(connections, hosts), **self.kwargs) + self.connection_pool = self.connection_pool_class(list(zip(connections, hosts)), **self.kwargs) def get_connection(self, sniffing=False): """ diff --git a/test_elasticsearch/test_transport.py b/test_elasticsearch/test_transport.py index 917128fc..9d53b676 100644 --- a/test_elasticsearch/test_transport.py +++ b/test_elasticsearch/test_transport.py @@ -84,6 +84,14 @@ class TestTransport(TestCase): self.assertEquals(1, len(t.connection_pool.connections)) self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) + def test_sniff_reuses_connection_instances_if_possible(self): + t = Transport([{'data': CLUSTER_NODES}, {"host": "1.1.1.1", "port": 123}], connection_class=DummyConnection, randomize_hosts=False) + connection = t.connection_pool.connections[1] + + t.sniff_hosts() + self.assertEquals(1, len(t.connection_pool.connections)) + self.assertIs(connection, t.get_connection()[0]) + def test_sniff_on_fail_triggers_sniffing_on_fail(self): t = Transport([{'exception': ConnectionError('abandon ship')}, {"data": CLUSTER_NODES}], connection_class=DummyConnection, sniff_on_connection_fail=True, max_retries=1, randomize_hosts=False)