diff --git a/elasticsearch/connection_pool.py b/elasticsearch/connection_pool.py index 70158013..c68a7a3b 100644 --- a/elasticsearch/connection_pool.py +++ b/elasticsearch/connection_pool.py @@ -110,13 +110,12 @@ class ConnectionPool(object): self.selector = selector_class(dict(connections)) - def mark_dead(self, connection, dead_count, now=None): + def mark_dead(self, connection, now=None): """ Mark the connection as dead (failed). Remove it from the live pool and put it on a timeout. :arg connection: the failed instance - :arg dead_count: number of consecutive failures """ # allow inject for testing purposes now = now if now else time.time() @@ -126,6 +125,7 @@ class ConnectionPool(object): # connection not alive or another thread marked it already, ignore return else: + dead_count = self.dead_count.get(connection, 0) + 1 self.dead_count[connection] = dead_count self.dead.put((now + self.dead_timeout * 2 ** (dead_count - 1), connection)) @@ -190,6 +190,6 @@ class ConnectionPool(object): connection = self.selector.select(self.connections) - return connection, self.dead_count.get(connection, 0) + return connection diff --git a/elasticsearch/transport.py b/elasticsearch/transport.py index e94fd127..7d099429 100644 --- a/elasticsearch/transport.py +++ b/elasticsearch/transport.py @@ -176,16 +176,15 @@ class Transport(object): self.set_connections(hosts) - def mark_dead(self, connection, dead_count): + def mark_dead(self, connection): """ Mark a connection as dead (failed) in the connection pool. If sniffing on failure is enabled this will initiate the sniffing process. :arg connection: instance of :class:`~elasticsearch.Connection` that failed - :arg dead_count: number of successive failures for this connection """ # mark as dead even when sniffing to avoid hitting this host during the sniff process - self.connection_pool.mark_dead(connection, dead_count) + self.connection_pool.mark_dead(connection) if self.sniff_on_connection_fail: self.sniff_hosts() @@ -212,20 +211,19 @@ class Transport(object): body = self.serializer.dumps(body) for attempt in range(self.max_retries + 1): - connection, dead_count = self.get_connection() + connection = self.get_connection() try: status, raw_data = connection.perform_request(method, url, params, body) except ConnectionError: - self.mark_dead(connection, dead_count + 1) + self.mark_dead(connection) # raise exception on last retry if attempt == self.max_retries: raise else: # resurrected connection didn't fail, confirm it's live status - if dead_count: - self.connection_pool.mark_live(connection) + self.connection_pool.mark_live(connection) data = None if raw_data: data = self.serializer.loads(raw_data) diff --git a/test_elasticsearch/test_connection_pool.py b/test_elasticsearch/test_connection_pool.py index 3be4d59e..51b98667 100644 --- a/test_elasticsearch/test_connection_pool.py +++ b/test_elasticsearch/test_connection_pool.py @@ -9,7 +9,7 @@ class TestConnectionPool(TestCase): connections = set() for _ in range(100): - connections.add(pool.get_connection()[0]) + connections.add(pool.get_connection()) self.assertEquals(connections, set(range(100))) def test_disable_shuffling(self): @@ -17,7 +17,7 @@ class TestConnectionPool(TestCase): connections = [] for _ in range(100): - connections.append(pool.get_connection()[0]) + connections.append(pool.get_connection()) self.assertEquals(connections, list(range(100))) def test_selectors_have_access_to_connection_opts(self): @@ -28,38 +28,39 @@ class TestConnectionPool(TestCase): connections = [] for _ in range(100): - connections.append(pool.get_connection()[0]) + connections.append(pool.get_connection()) self.assertEquals(connections, [x*x for x in range(100)]) def test_dead_nodes_are_removed_from_active_connections(self): pool = ConnectionPool([(x, {}) for x in range(100)]) now = time.time() - pool.mark_dead(42, 1, now=now) + pool.mark_dead(42, now=now) self.assertEquals(99, len(pool.connections)) self.assertEquals(1, pool.dead.qsize()) self.assertEquals((now + 60, 42), pool.dead.get()) def test_connection_is_skipped_when_dead(self): pool = ConnectionPool([(x, {}) for x in range(2)]) - pool.mark_dead(0, 1) + pool.mark_dead(0) - self.assertEquals([(1, 0), (1, 0), (1, 0)], [pool.get_connection(), pool.get_connection(), pool.get_connection(), ]) + self.assertEquals([1, 1, 1], [pool.get_connection(), pool.get_connection(), pool.get_connection(), ]) def test_connection_is_forcibly_resurrected_when_no_live_ones_are_availible(self): pool = ConnectionPool([(x, {}) for x in range(2)]) - pool.mark_dead(0, 2) # failed twice, longer timeout - pool.mark_dead(1, 1) # failed the first time, first to be resurrected + pool.dead_count[0] = 1 + pool.mark_dead(0) # failed twice, longer timeout + pool.mark_dead(1) # failed the first time, first to be resurrected self.assertEquals([], pool.connections) - self.assertEquals((1, 1), pool.get_connection()) + self.assertEquals(1, pool.get_connection()) self.assertEquals([1,], pool.connections) def test_connection_is_resurrected_after_its_timeout(self): pool = ConnectionPool([(x, {}) for x in range(100)]) now = time.time() - pool.mark_dead(42, 1, now=now-61) + pool.mark_dead(42, now=now-61) pool.get_connection() self.assertEquals(42, pool.connections[-1]) self.assertEquals(100, len(pool.connections)) @@ -67,7 +68,8 @@ class TestConnectionPool(TestCase): def test_already_failed_connection_has_longer_timeout(self): pool = ConnectionPool([(x, {}) for x in range(100)]) now = time.time() - pool.mark_dead(42, 3, now=now) + pool.dead_count[42] = 2 + pool.mark_dead(42, now=now) self.assertEquals(3, pool.dead_count[42]) self.assertEquals((now + 4*60, 42), pool.dead.get()) @@ -75,7 +77,8 @@ class TestConnectionPool(TestCase): def test_dead_count_is_wiped_clean_for_connection_if_marked_live(self): pool = ConnectionPool([(x, {}) for x in range(100)]) now = time.time() - pool.mark_dead(42, 3, now=now) + pool.dead_count[42] = 2 + pool.mark_dead(42, now=now) self.assertEquals(3, pool.dead_count[42]) pool.mark_live(42) diff --git a/test_elasticsearch/test_transport.py b/test_elasticsearch/test_transport.py index 249c7582..0080dc38 100644 --- a/test_elasticsearch/test_transport.py +++ b/test_elasticsearch/test_transport.py @@ -62,7 +62,7 @@ class TestTransport(TestCase): t = Transport([{'exception': ConnectionError('abandon ship')}], connection_class=DummyConnection) self.assertRaises(ConnectionError, t.perform_request, 'GET', '/') - self.assertEquals(4, len(t.get_connection()[0].calls)) + self.assertEquals(4, len(t.get_connection().calls)) def test_failed_connection_will_be_marked_as_dead(self): t = Transport([{'exception': ConnectionError('abandon ship')}], connection_class=DummyConnection) @@ -72,8 +72,8 @@ class TestTransport(TestCase): def test_resurrected_connection_will_be_marked_as_live_on_success(self): t = Transport([{}], connection_class=DummyConnection) - con, failed_count = t.connection_pool.get_connection() - t.connection_pool.mark_dead(con, 2) + con = t.connection_pool.get_connection() + t.connection_pool.mark_dead(con) t.perform_request('GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) @@ -85,12 +85,12 @@ class TestTransport(TestCase): t.sniff_hosts() self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) + self.assertEquals('http://1.1.1.1:123', t.get_connection().host) def test_sniff_on_start_fetches_and_uses_nodes_list(self): t = Transport([{'data': CLUSTER_NODES}], connection_class=DummyConnection, sniff_on_start=True) self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) + self.assertEquals('http://1.1.1.1:123', t.get_connection().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) @@ -98,7 +98,7 @@ class TestTransport(TestCase): t.sniff_hosts() self.assertEquals(1, len(t.connection_pool.connections)) - self.assertIs(connection, t.get_connection()[0]) + self.assertIs(connection, t.get_connection()) def test_sniff_on_fail_triggers_sniffing_on_fail(self): t = Transport([{'exception': ConnectionError('abandon ship')}, {"data": CLUSTER_NODES}], @@ -106,7 +106,7 @@ class TestTransport(TestCase): self.assertRaises(ConnectionError, t.perform_request, 'GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) + self.assertEquals('http://1.1.1.1:123', t.get_connection().host) def test_sniff_after_n_seconds(self): t = Transport([{"data": CLUSTER_NODES}], @@ -115,11 +115,11 @@ class TestTransport(TestCase): for _ in range(4): t.perform_request('GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) - self.assertIsInstance(t.get_connection()[0], DummyConnection) + self.assertIsInstance(t.get_connection(), DummyConnection) t.last_sniff = time.time() - 5.1 t.perform_request('GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) + self.assertEquals('http://1.1.1.1:123', t.get_connection().host) self.assertTrue(time.time() - 1 < t.last_sniff < time.time() + 0.01 )