diff --git a/elasticsearch/connection_pool.py b/elasticsearch/connection_pool.py index 0e60f298..041de9b4 100644 --- a/elasticsearch/connection_pool.py +++ b/elasticsearch/connection_pool.py @@ -28,6 +28,7 @@ class ConnectionPool(object): def __init__(self, connections, dead_timeout=60, selector_class=RoundRobinSelector, randomize_hosts=True, **kwargs): self.connections = [c for (c, opts) in connections] self.dead = PriorityQueue(len(self.connections)) + self.dead_count = {} if randomize_hosts: # randomize the connection list to avoid all clients hitting same node @@ -39,7 +40,7 @@ class ConnectionPool(object): self.selector = selector_class(dict(connections)) - def mark_dead(self, connection, now=None): + def mark_dead(self, connection, dead_count, now=None): # allow inject for testing purposes now = now if now else time.time() try: @@ -49,7 +50,15 @@ class ConnectionPool(object): return else: # TODO: detect repeated failure and extend the timeout - self.dead.put((now + self.dead_timeout, connection)) + self.dead_count[connection] = dead_count + self.dead.put((now + self.dead_timeout * 2 ** (dead_count - 1), connection)) + + def mark_live(self, connection): + try: + del self.dead_count[connection] + except KeyError: + # race condition, safe to ignore + pass def resurrect(self, force=False): # no dead connections @@ -74,6 +83,8 @@ class ConnectionPool(object): if not self.connections: self.resurrect(True) - return self.selector.select(self.connections) + connection = self.selector.select(self.connections) + + return connection, self.dead_count.get(connection, 0) diff --git a/elasticsearch/transport.py b/elasticsearch/transport.py index 66797314..2e93f8e5 100644 --- a/elasticsearch/transport.py +++ b/elasticsearch/transport.py @@ -92,26 +92,29 @@ class Transport(object): self.sniffs_due_to_failure = 0 self.sniff_after_requests = self.sniff_after_requests_original - def mark_dead(self, connection, sniffing=False): + def mark_dead(self, connection, dead_count, sniffing=False): if not sniffing and self.sniff_on_connection_fail: self.sniff_hosts(True) else: - self.connection_pool.mark_dead(connection) + self.connection_pool.mark_dead(connection, dead_count) def perform_request(self, method, url, params=None, body=None, sniffing=False): for attempt in range(self.max_retries): - connection = self.get_connection(sniffing) + connection, dead_count = self.get_connection(sniffing) if body: body = self.serializer.dumps(body) try: status, raw_data = connection.perform_request(method, url, params, body) except TransportError: - self.mark_dead(connection, sniffing) + self.mark_dead(connection, dead_count + 1, sniffing) # raise exception on last retry if attempt + 1 == self.max_retries: raise else: + # resurrected connection didn't fail, confirm it's live status + if dead_count: + self.connection_pool.mark_live(connection) return status, self.serializer.loads(raw_data) diff --git a/test_elasticsearch/test_connection_pool.py b/test_elasticsearch/test_connection_pool.py index da3e0a43..3be4d59e 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()) + connections.add(pool.get_connection()[0]) self.assertEquals(connections, set(range(100))) def test_disable_shuffling(self): @@ -17,50 +17,67 @@ class TestConnectionPool(TestCase): connections = [] for _ in range(100): - connections.append(pool.get_connection()) + connections.append(pool.get_connection()[0]) self.assertEquals(connections, list(range(100))) def test_selectors_have_access_to_connection_opts(self): class MySelector(RoundRobinSelector): def select(self, connections): - return self.connection_opts[super(MySelector, self).select(connections)] - pool = ConnectionPool([(x, {"number": x}) for x in range(100)], selector_class=MySelector, randomize_hosts=False) + return self.connection_opts[super(MySelector, self).select(connections)]["actual"] + pool = ConnectionPool([(x, {"actual": x*x}) for x in range(100)], selector_class=MySelector, randomize_hosts=False) connections = [] for _ in range(100): - connections.append(pool.get_connection()) - self.assertEquals(connections, [{"number": x} for x in range(100)]) + connections.append(pool.get_connection()[0]) + 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, now=now) + pool.mark_dead(42, 1, 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) + pool.mark_dead(0, 1) - self.assertEquals([1, 1, 1], [pool.get_connection(), pool.get_connection(), pool.get_connection(), ]) + self.assertEquals([(1, 0), (1, 0), (1, 0)], [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) - pool.mark_dead(1) + pool.mark_dead(0, 2) # failed twice, longer timeout + pool.mark_dead(1, 1) # failed the first time, first to be resurrected self.assertEquals([], pool.connections) - self.assertEquals(0, pool.get_connection()) - self.assertEquals([0,], pool.connections) + self.assertEquals((1, 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, now=now-61) + pool.mark_dead(42, 1, now=now-61) pool.get_connection() self.assertEquals(42, pool.connections[-1]) self.assertEquals(100, len(pool.connections)) + 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) + + self.assertEquals(3, pool.dead_count[42]) + self.assertEquals((now + 4*60, 42), pool.dead.get()) + + 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) + + self.assertEquals(3, pool.dead_count[42]) + pool.mark_live(42) + self.assertNotIn(42, pool.dead_count) + diff --git a/test_elasticsearch/test_transport.py b/test_elasticsearch/test_transport.py index 066641db..41ee9ad2 100644 --- a/test_elasticsearch/test_transport.py +++ b/test_elasticsearch/test_transport.py @@ -61,7 +61,7 @@ class TestTransport(TestCase): t = Transport([{'exception': TransportError('abandon ship')}], connection_class=DummyConnection) self.assertRaises(TransportError, t.perform_request, 'GET', '/') - self.assertEquals(3, len(t.get_connection().calls)) + self.assertEquals(3, len(t.get_connection()[0].calls)) def test_failed_connection_will_be_marked_as_dead(self): t = Transport([{'exception': TransportError('abandon ship')}], connection_class=DummyConnection) @@ -69,10 +69,19 @@ class TestTransport(TestCase): self.assertRaises(TransportError, t.perform_request, 'GET', '/') self.assertEquals(0, len(t.connection_pool.connections)) + 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) + + t.perform_request('GET', '/') + self.assertEquals(1, len(t.connection_pool.connections)) + self.assertEquals(0, len(t.connection_pool.dead_count)) + 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().host) + self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) def test_sniff_on_fail_triggers_sniffing_on_fail(self): t = Transport([{'exception': TransportError('abandon ship')}, {"data": CLUSTER_NODES}], @@ -80,7 +89,7 @@ class TestTransport(TestCase): self.assertRaises(TransportError, t.perform_request, 'GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection().host) + self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) def test_sniff_after_n_requests(self): t = Transport([{"data": CLUSTER_NODES}], @@ -89,11 +98,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(), DummyConnection) + self.assertIsInstance(t.get_connection()[0], DummyConnection) t.perform_request('GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection().host) + self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) def test_sniff_on_failure_shortens_sniff_after_n_requests(self): t = Transport([{'exception': TransportError('abandon ship')}, {"data": CLUSTER_NODES}], @@ -102,6 +111,6 @@ class TestTransport(TestCase): self.assertRaises(TransportError, t.perform_request, 'GET', '/') self.assertEquals(1, len(t.connection_pool.connections)) - self.assertEquals('http://1.1.1.1:123', t.get_connection().host) + self.assertEquals('http://1.1.1.1:123', t.get_connection()[0].host) self.assertEquals(3, t.sniff_after_requests)