From e3c76ecf01e7c72582327426fc6467e08f988b8b Mon Sep 17 00:00:00 2001
From: Honza Kral <honza.kral@gmail.com>
Date: Thu, 2 May 2013 17:54:26 +0200
Subject: [PATCH] ConnetionPool to hold all the connections

It provides load balancing and keeps track of the status of individual
connections (dead/alive) and penalizes dead connections with a timeout.
---
 elasticsearch/connection_pool.py           | 74 ++++++++++++++++++++++
 test_elasticsearch/test_connection_pool.py | 64 +++++++++++++++++++
 2 files changed, 138 insertions(+)
 create mode 100644 elasticsearch/connection_pool.py
 create mode 100644 test_elasticsearch/test_connection_pool.py

diff --git a/elasticsearch/connection_pool.py b/elasticsearch/connection_pool.py
new file mode 100644
index 00000000..59784a67
--- /dev/null
+++ b/elasticsearch/connection_pool.py
@@ -0,0 +1,74 @@
+import time
+import random
+
+class ConnectionSelector(object):
+    " Base class for Selectors. "
+    def __init__(self, opts):
+        self.connection_opts = opts
+
+
+class RoundRobinSelector(ConnectionSelector):
+    " Default selector using round-robin. "
+    def __init__(self, opts):
+        super(RoundRobinSelector, self).__init__(opts)
+        self.rr = -1
+
+    def select(self, connections):
+        self.rr += 1
+        self.rr %= len(connections)
+        return connections[self.rr]
+
+
+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 = []
+
+        if randomize_hosts:
+            # randomize the connection list to avoid all clients hitting same node
+            # after startup/restart
+            random.shuffle(self.connections)
+
+        # default timeout after which to try resurrecting a connection
+        self.dead_timeout = dead_timeout
+
+        self.selector = selector_class(dict(connections))
+
+    def mark_dead(self, connection, now=None):
+        # allow inject for testing purposes
+        now = now if now else time.time()
+        try:
+            self.connections.remove(connection)
+        except ValueError:
+            # connection not alive, ignore
+            return
+
+        # TODO: detect repeated failure and extend the timeout
+        self.dead.append((now + self.dead_timeout, connection))
+
+    def resurrect(self, force=False):
+        # no dead connections
+        if not self.dead:
+            return
+
+        # no elligible connections to retry
+        if not force and self.dead[0][0] > time.time():
+            return
+
+        # either we were forced or the node is elligible to be retried
+        connection = self.dead.pop(0)[1]
+        self.connections.append(connection)
+        return connection
+
+    def get_connection(self):
+        connection = self.resurrect()
+        if connection:
+            return connection
+
+        # no live nodes, resurrect one by force
+        if not self.connections:
+            return self.resurrect(True)
+
+        return self.selector.select(self.connections)
+
+
diff --git a/test_elasticsearch/test_connection_pool.py b/test_elasticsearch/test_connection_pool.py
new file mode 100644
index 00000000..a1861057
--- /dev/null
+++ b/test_elasticsearch/test_connection_pool.py
@@ -0,0 +1,64 @@
+import time
+from unittest import TestCase
+
+from elasticsearch.connection_pool import ConnectionPool, RoundRobinSelector
+
+class TestConnectionPool(TestCase):
+    def test_default_round_robin(self):
+        pool = ConnectionPool([(x, {}) for x in range(100)])
+
+        connections = set()
+        for _ in range(100):
+            connections.add(pool.get_connection())
+        self.assertEquals(connections, set(range(100)))
+
+    def test_disable_shuffling(self):
+        pool = ConnectionPool([(x, {}) for x in range(100)], randomize_hosts=False)
+
+        connections = []
+        for _ in range(100):
+            connections.append(pool.get_connection())
+        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)
+
+        connections = []
+        for _ in range(100):
+            connections.append(pool.get_connection())
+        self.assertEquals(connections, [{"number": 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)
+        self.assertEquals(99, len(pool.connections))
+        self.assertEquals([(now + 60, 42)], pool.dead)
+
+    def test_connection_is_skipped_when_dead(self):
+        pool = ConnectionPool([(x, {}) for x in range(2)])
+        pool.mark_dead(0)
+
+        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)
+        pool.mark_dead(1)
+
+        self.assertEquals([], pool.connections)
+        self.assertEquals(0, pool.get_connection())
+        self.assertEquals([0,], 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)
+        self.assertEquals(42, pool.get_connection())
+        self.assertEquals(100, len(pool.connections))
+