diff --git a/helpers/zookeeper.py b/helpers/zookeeper.py index 4e7a22c6..cb2918cd 100644 --- a/helpers/zookeeper.py +++ b/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) diff --git a/postgres0.yml b/postgres0.yml index 0cb694a1..7b8527dc 100644 --- a/postgres0.yml +++ b/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 diff --git a/postgres1.yml b/postgres1.yml index 680091f1..5be3fc25 100644 --- a/postgres1.yml +++ b/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 diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 609be245..599b4bb2 100644 --- a/tests/test_etcd.py +++ b/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'): diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py index 0799413b..172a8c5f 100644 --- a/tests/test_zookeeper.py +++ b/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()