Add exhibitor support

List of ZooKeeper nodes could be periodically updated from Exhibitor
Since we know that each Exhibitor accompanies one ZooKeeper node, list
of Exhibitor nodes also maintained. Exhibitor assumes that all ZooKeeper
nodes are using the same client port, 2181. The same assumption is valid
for Exhibitor, it should always listen on the same port on all nodes.

Original list of Exhibitor nodes is cached and used as a fallback when
it failed ito query information with using maintained list.

helpers/zookeeper.py

@@ -1,8 +1,13 @@
 import logging
+import random
+import requests
+import time
 
 from helpers.dcs import AbstractDCS, Cluster, DCSError, Member, parse_connection_string
+from helpers.utils import sleep
 from kazoo.client import KazooClient, KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
+from requests.exceptions import RequestException
 
 logger = logging.getLogger(__name__)
 
@@ -11,15 +16,73 @@ class ZooKeeperError(DCSError):
     pass
 
 
+class ExhibitorEnsembleProvider:
+
+    TIMEOUT = 3.1
+
+    def __init__(self, hosts, port, uri_path='/exhibitor/v1/cluster/list', poll_interval=300):
+        self._exhibitor_port = port
+        self._uri_path = uri_path
+        self._poll_interval = poll_interval
+        self._exhibitors = hosts
+        self._master_exhibitors = hosts
+        self._zookeeper_hosts = ''
+        self._next_poll = None
+        while not self.poll():
+            logger.info('waiting on exhibitor')
+            sleep(5)
+
+    def poll(self):
+        if self._next_poll and self._next_poll > time.time():
+            return False
+
+        json = self._query_exhibitors(self._exhibitors)
+        if not json:
+            json = self._query_exhibitors(self._master_exhibitors)
+
+        if isinstance(json, dict) and 'servers' in json and 'port' in json:
+            self._next_poll = time.time() + self._poll_interval
+            zookeeper_hosts = ','.join([h + ':' + str(json['port']) for h in sorted(json['servers'])])
+            if self._zookeeper_hosts != zookeeper_hosts:
+                logger.info('ZooKeeper connection string has changed: %s => %s', self._zookeeper_hosts, zookeeper_hosts)
+                self._zookeeper_hosts = zookeeper_hosts
+                self._exhibitors = json['servers']
+                return True
+        return False
+
+    def _query_exhibitors(self, exhibitors):
+        random.shuffle(exhibitors)
+        for host in exhibitors:
+            uri = 'http://{}:{}{}'.format(host, self._exhibitor_port, self._uri_path)
+            try:
+                response = requests.get(uri, timeout=self.TIMEOUT)
+                return response.json()
+            except RequestException:
+                pass
+        return None
+
+    @property
+    def zookeeper_hosts(self):
+        return self._zookeeper_hosts
+
+
 class ZooKeeper(AbstractDCS):
 
     def __init__(self, name, config):
         super(ZooKeeper, self).__init__(name, config)
-        self.fetch_cluster = True
-        self.members = []
-        self.leader = None
-        self.last_leader_operation = 0
-        self.client = KazooClient(hosts=config['hosts'],
+
+        hosts = config.get('hosts', [])
+        if isinstance(hosts, list):
+            hosts = ','.join(hosts)
+
+        self.exhibitor = None
+        if 'exhibitor' in config:
+            exhibitor = config['exhibitor']
+            interval = exhibitor.get('poll_interval', 300)
+            self.exhibitor = ExhibitorEnsembleProvider(exhibitor['hosts'], exhibitor['port'], poll_interval=interval)
+            hosts = self.exhibitor.zookeeper_hosts
+
+        self.client = KazooClient(hosts=hosts,
                                   timeout=(config.get('session_timeout', None) or 30),
                                   command_retry={
                                       'deadline': (config.get('reconnect_timeout', None) or 10),
@@ -28,6 +91,12 @@ class ZooKeeper(AbstractDCS):
                                   connection_retry={'max_delay': 1, 'max_tries': -1})
         self.client.add_listener(self.session_listener)
         self.cluster_event = self.client.handler.event_object()
+
+        self.fetch_cluster = True
+        self.members = []
+        self.leader = None
+        self.last_leader_operation = 0
+
         self.client.start(None)
 
     def session_listener(self, state):
@@ -87,6 +156,9 @@ class ZooKeeper(AbstractDCS):
                 self.last_leader_operation = int(last_leader_operation[0])
 
     def get_cluster(self):
+        if self.exhibitor and self.exhibitor.poll():
+            self.client.set_hosts(self.exhibitor.zookeeper_hosts)
+
         if self.fetch_cluster:
             try:
                 self.client.retry(self._inner_load_cluster)

postgres0.yml

@@ -12,7 +12,16 @@ etcd:
 #  scope: batman
 #  session_timeout: *ttl
 #  reconnect_timeout: *loop_wait
-#  hosts: 127.0.0.1:2181
+#  hosts:
+#    - 127.0.0.1:2181
+#    - 127.0.0.2:2181
+#  exhibitor:
+#    poll_interval: 300
+#    port: 8181
+#    hosts:
+#      - host1
+#      - host2
+#      - host3
 postgresql:
   name: postgresql0
   listen: 127.0.0.1:5432

postgres1.yml

@@ -12,7 +12,16 @@ etcd:
 #  scope: batman
 #  session_timeout: *ttl
 #  reconnect_timeout: *loop_wait
-#  hosts: 127.0.0.1:2181
+#  hosts:
+#    - 127.0.0.1:2181
+#    - 127.0.0.2:2181
+#  exhibitor:
+#    poll_interval: 300
+#    port: 8181
+#    hosts:
+#      - host1
+#      - host2
+#      - host3
 postgresql:
   name: postgresql1
   listen: 127.0.0.1:5433

tests/test_etcd.py

@@ -43,6 +43,8 @@ def requests_get(url, **kwargs):
             response.content = members
     elif url.endswith('/bad_response'):
         response.content = '{'
+    elif url.startswith('http://exhibitor'):
+        response.content = '{"servers":["127.0.0.1","127.0.0.2","127.0.0.3"],"port":2181}'
     elif url.startswith('http://local'):
         raise requests.exceptions.RequestException()
     elif url.startswith('http://remote') or url.startswith('http://127.0.0.1') or url.startswith('http://error'):

tests/test_zookeeper.py

@@ -1,11 +1,12 @@
 import helpers.zookeeper
 import unittest
+import requests
 
-from helpers.zookeeper import ZooKeeper, ZooKeeperError
+from helpers.zookeeper import ExhibitorEnsembleProvider, ZooKeeper, ZooKeeperError
 from kazoo.client import KazooState
 from kazoo.exceptions import NoNodeError, NodeExistsError
 from kazoo.protocol.states import ZnodeStat
-from test_etcd import MockPostgresql
+from test_etcd import MockPostgresql, requests_get
 
 
 class MockEvent:
@@ -88,6 +89,27 @@ class MockKazooClient:
             self.leader = True
             raise Exception
 
+    def set_hosts(self, hosts, randomize_hosts=None):
+        pass
+
+
+def exhibitor_sleep(_):
+    raise Exception
+
+
+class TestExhibitorEnsembleProvider(unittest.TestCase):
+
+    def __init__(self, method_name='runTest'):
+        self.setUp = self.set_up
+        super(TestExhibitorEnsembleProvider, self).__init__(method_name)
+
+    def set_up(self):
+        requests.get = requests_get
+        helpers.zookeeper.sleep = exhibitor_sleep
+
+    def test_init(self):
+        self.assertRaises(Exception, ExhibitorEnsembleProvider, ['localhost'], 8181)
+
 
 class TestZooKeeper(unittest.TestCase):
 
@@ -96,8 +118,9 @@ class TestZooKeeper(unittest.TestCase):
         super(TestZooKeeper, self).__init__(method_name)
 
     def set_up(self):
+        requests.get = requests_get
         helpers.zookeeper.KazooClient = MockKazooClient
-        self.zk = ZooKeeper('foo', {'hosts': 'localhost:2181', 'scope': 'test'})
+        self.zk = ZooKeeper('foo', {'exhibitor': {'hosts': ['localhost', 'exhibitor'], 'port': 8181}, 'scope': 'test'})
 
     def test_session_listener(self):
         self.zk.session_listener(KazooState.SUSPENDED)
@@ -112,6 +135,7 @@ class TestZooKeeper(unittest.TestCase):
 
     def test_get_cluster(self):
         self.assertRaises(ZooKeeperError, self.zk.get_cluster)
+        self.zk.exhibitor.poll = lambda: True
         self.zk.get_cluster()
         self.zk.touch_member('foo')
         self.zk.delete_leader()