From b4ba7433bdb5c222ae0fecd3f6f389f01139a4b7 Mon Sep 17 00:00:00 2001
From: Honza Kral <honza.kral@gmail.com>
Date: Sun, 5 May 2013 00:36:39 +0200
Subject: [PATCH] Make sure dead connections are handled in a thread-safe
 manner

---
 elasticsearch/connection_pool.py           | 28 ++++++++++++++--------
 test_elasticsearch/test_connection_pool.py |  3 ++-
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/elasticsearch/connection_pool.py b/elasticsearch/connection_pool.py
index 536d0d54..0e60f298 100644
--- a/elasticsearch/connection_pool.py
+++ b/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
-
-        # TODO: detect repeated failure and extend the timeout
-        self.dead.append((now + self.dead_timeout, connection))
+        else:
+            # TODO: detect repeated failure and extend the timeout
+            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
-        if not force and self.dead[0][0] > time.time():
+        # retrieve a connection to check
+        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
-        connection = self.dead.pop(0)[1]
+        # either we were forced or the connection is elligible to be retried
         self.connections.append(connection)
 
     def get_connection(self):
diff --git a/test_elasticsearch/test_connection_pool.py b/test_elasticsearch/test_connection_pool.py
index 4fdaf699..da3e0a43 100644
--- a/test_elasticsearch/test_connection_pool.py
+++ b/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)])