Make sure dead connections are handled in a thread-safe manner

elasticsearch/connection_pool.py

@@ -1,6 +1,11 @@
 import time
 import random
 
+try:
+    from Queue import PriorityQueue
+except ImportError:
+    from queue import PriorityQueue
+
 class ConnectionSelector(object):
     " Base class for Selectors. "
     def __init__(self, opts):
@@ -22,7 +27,7 @@ class RoundRobinSelector(ConnectionSelector):
 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 = []
+        self.dead = PriorityQueue(len(self.connections))
 
         if randomize_hosts:
             # randomize the connection list to avoid all clients hitting same node
@@ -40,23 +45,26 @@ class ConnectionPool(object):
         try:
             self.connections.remove(connection)
         except ValueError:
-            # connection not alive, ignore
+            # connection not alive or another thread marked it already, ignore
             return
-
+        else:
-        # TODO: detect repeated failure and extend the timeout
+            # TODO: detect repeated failure and extend the timeout
-        self.dead.append((now + self.dead_timeout, connection))
+            self.dead.put((now + self.dead_timeout, connection))
 
     def resurrect(self, force=False):
         # no dead connections
-        if not self.dead:
+        if self.dead.empty():
             return
 
-        # no elligible connections to retry
+        # retrieve a connection to check
-        if not force and self.dead[0][0] > time.time():
+        timeout, connection = self.dead.get()
+
+        if not force and timeout > time.time():
+            # return it back if not eligible and not forced
+            self.dead.put((timeout, connection))
             return
 
-        # either we were forced or the node is elligible to be retried
+        # either we were forced or the connection is elligible to be retried
-        connection = self.dead.pop(0)[1]
         self.connections.append(connection)
 
     def get_connection(self):

test_elasticsearch/test_connection_pool.py

@@ -37,7 +37,8 @@ class TestConnectionPool(TestCase):
         now = time.time()
         pool.mark_dead(42, now=now)
         self.assertEquals(99, len(pool.connections))
-        self.assertEquals([(now + 60, 42)], pool.dead)
+        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)])