Merge pull request #2 from zalando/feature/exhibitor

Feature/exhibitor

.travis.yml

@@ -9,5 +9,6 @@ install:
   - pip install coveralls
 script:
   - python setup.py test
+  - python setup.py flake8
 after_success:
   - coveralls

helpers/aws.py

@@ -1,5 +1,4 @@
 import logging
-import re
 import requests
 from requests.exceptions import RequestException
 import yaml
@@ -9,6 +8,7 @@ logger = logging.getLogger(__name__)
 
 
 class AWSConnection:
+
     def __init__(self, config):
         self.available = False
         self.config = config
@@ -43,7 +43,7 @@ class AWSConnection:
         if not self.available:
             return False
 
-        tags = {'Name': 'spilo_'+self.cluster_name, 'Role': role, 'Instance': self.instance_id}
+        tags = {'Name': 'spilo_' + self.cluster_name, 'Role': role, 'Instance': self.instance_id}
         try:
             conn = boto.ec2.connect_to_region(self.region)
             volumes = conn.get_all_volumes(filters={'attachment.instance-id': self.instance_id})

helpers/postgresql.py

@@ -9,6 +9,7 @@ from helpers.utils import sleep
 
 if sys.hexversion >= 0x03000000:
     from urllib.parse import urlparse
+    long = int
 else:
     from urlparse import urlparse
 
@@ -169,7 +170,8 @@ class Postgresql:
 
         try:
             latest_backup = subprocess.check_output(self.wal_e_path.split() + ['backup-list', '--detail', 'LATEST'])
-            # name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start wal_segment_backup_stop wal_segment_offset_backup_stop
+            # name    last_modified   expanded_size_bytes wal_segment_backup_start    wal_segment_offset_backup_start
+            #                                                   wal_segment_backup_stop wal_segment_offset_backup_stop
             # base_00000001000000000000007F_00000040  2015-05-18T10:13:25.000Z
             # 20310671    00000001000000000000007F    00000040
             # 00000001000000000000007F    00000240

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_postgresql.py

@@ -118,7 +118,7 @@ class TestPostgresql(unittest.TestCase):
                                              'password': 'rep-pass',
                                              'network': '127.0.0.1/32'},
                              'parameters': {'foo': 'bar'}, 'recovery_conf': {'foo': 'bar'}},
-                             on_change_callback=lambda state: True)
+                            on_change_callback=lambda state: True)
         psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)

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()