Merge pull request #28 from zalando/feature/multidcs

Feature/multidcs

governor.py

@@ -6,26 +6,34 @@ import time
 import yaml
 
 from helpers.api import RestApiServer
-from helpers.etcd import Etcd
-from helpers.postgresql import Postgresql
-from helpers.ha import Ha
-from helpers.utils import setup_signal_handlers, sleep
 from helpers.aws import AWSConnection
+from helpers.etcd import Etcd
+from helpers.ha import Ha
+from helpers.postgresql import Postgresql
+from helpers.utils import setup_signal_handlers, sleep
+from helpers.zookeeper import ZooKeeper
 
 
 class Governor:
 
     def __init__(self, config):
         self.nap_time = config['loop_wait']
-        self.etcd = Etcd(config['etcd'])
         self.aws = AWSConnection(config)
-        self.postgresql = Postgresql(config['postgresql'], self.aws.on_role_change)
-        self.ha = Ha(self.postgresql, self.etcd)
+        self.postgresql = Postgresql(config['postgresql'])
+        self.ha = Ha(self.postgresql, self.get_dcs(self.postgresql.name, config))
         host, port = config['restapi']['listen'].split(':')
         self.api = RestApiServer(self, config['restapi'])
         self.next_run = time.time()
         self.shutdown_member_ttl = 300
 
+    @staticmethod
+    def get_dcs(name, config):
+        if 'etcd' in config:
+            return Etcd(name, config['etcd'])
+        if 'zookeeper' in config:
+            return ZooKeeper(name, config['zookeeper'])
+        raise Exception('Can not find sutable configuration of distributed configuration store')
+
     def touch_member(self, ttl=None):
         connection_string = self.postgresql.connection_string + '?application_name=' + self.api.connection_string
         if self.ha.cluster:
@@ -33,26 +41,26 @@ class Governor:
                 # Do not update member TTL when it is far from being expired
                 if m.name == self.postgresql.name and m.real_ttl() > self.shutdown_member_ttl:
                     return True
-        return self.etcd.touch_member(self.postgresql.name, connection_string, ttl)
+        return self.ha.dcs.touch_member(connection_string, ttl)
 
     def initialize(self):
         # wait for etcd to be available
         while not self.touch_member():
-            logging.info('waiting on etcd')
+            logging.info('waiting on DCS')
             sleep(5)
 
         # is data directory empty?
         if self.postgresql.data_directory_empty():
             # racing to initialize
-            if self.etcd.race('/initialize', self.postgresql.name):
+            if self.ha.dcs.race('/initialize'):
                 self.postgresql.initialize()
-                self.etcd.take_leader(self.postgresql.name)
+                self.ha.dcs.take_leader()
                 self.postgresql.start()
                 self.postgresql.create_replication_user()
                 self.postgresql.create_connection_users()
             else:
                 while True:
-                    leader = self.etcd.current_leader()
+                    leader = self.ha.dcs.current_leader()
                     if leader and self.postgresql.sync_from_leader(leader):
                         self.postgresql.write_recovery_conf(leader)
                         self.postgresql.start()
@@ -68,7 +76,7 @@ class Governor:
         if nap_time <= 0:
             self.next_run = current_time
         else:
-            sleep(nap_time)
+            self.ha.dcs.sleep(nap_time)
 
     def run(self):
         self.api.start()
@@ -108,7 +116,7 @@ def main():
     finally:
         governor.touch_member(governor.shutdown_member_ttl)  # schedule member removal
         governor.postgresql.stop()
-        governor.etcd.delete_leader(governor.postgresql.name)
+        governor.ha.dcs.delete_leader()
 
 
 if __name__ == '__main__':

helpers/dcs.py

@@ -0,0 +1,92 @@
+import abc
+import sys
+
+from collections import namedtuple
+from helpers.utils import calculate_ttl, sleep
+
+if sys.hexversion >= 0x03000000:
+    from urllib.parse import urlparse, urlunparse, parse_qsl
+else:
+    from urlparse import urlparse, urlunparse, parse_qsl
+
+
+def parse_connection_string(value):
+    scheme, netloc, path, params, query, fragment = urlparse(value)
+    conn_url = urlunparse((scheme, netloc, path, params, '', fragment))
+    api_url = ([v for n, v in parse_qsl(query) if n == 'application_name'] or [None])[0]
+    return conn_url, api_url
+
+
+class DCSError(Exception):
+
+    def __init__(self, value):
+        self.value = value
+
+    def __str__(self):
+        """
+        >>> str(DCSError('foo'))
+        "'foo'"
+        """
+        return repr(self.value)
+
+
+class Member(namedtuple('Member', 'index,name,conn_url,api_url,expiration,ttl')):
+
+    def real_ttl(self):
+        return calculate_ttl(self.expiration) or -1
+
+
+class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')):
+
+    def is_unlocked(self):
+        return not (self.leader and self.leader.name)
+
+
+class AbstractDCS:
+
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, name, config):
+        self._name = name
+        self._base_path = '/service/' + config['scope']
+
+    def client_path(self, path):
+        return self._base_path + path
+
+    @abc.abstractmethod
+    def get_cluster(self):
+        """get_cluster"""
+
+    @abc.abstractmethod
+    def update_leader(self, state_handler):
+        """update_leader"""
+
+    @abc.abstractmethod
+    def attempt_to_acquire_leader(self):
+        """attempt_to_acquire_leader"""
+
+    def current_leader(self):
+        try:
+            cluster = self.get_cluster()
+            return None if cluster.is_unlocked() else cluster.leader
+        except DCSError:
+            return None
+
+    @abc.abstractmethod
+    def touch_member(self, connection_string, ttl=None):
+        """touch_member"""
+
+    @abc.abstractmethod
+    def take_leader(self):
+        """take_leader"""
+
+    @abc.abstractmethod
+    def race(self, path):
+        """race"""
+
+    @abc.abstractmethod
+    def delete_leader(self):
+        """delete_leader"""
+
+    def sleep(self, timeout):
+        sleep(timeout)

helpers/errors.py

@@ -1,15 +0,0 @@
-class EtcdError(Exception):
-
-    def __init__(self, value):
-        self.value = value
-
-    def __str__(self):
-        return repr(self.value)
-
-
-class CurrentLeaderError(EtcdError):
-    pass
-
-
-class EtcdConnectionFailed(EtcdError):
-    pass

helpers/etcd.py

@@ -2,42 +2,22 @@ import logging
 import random
 import requests
 import socket
-import sys
 
-from collections import namedtuple
 from dns.exception import DNSException
 from dns import resolver
-from helpers.errors import CurrentLeaderError, EtcdError, EtcdConnectionFailed
-from helpers.utils import calculate_ttl, sleep
+from helpers.dcs import AbstractDCS, Cluster, DCSError, Member, parse_connection_string
+from helpers.utils import sleep
 from requests.exceptions import RequestException
 
-if sys.hexversion >= 0x03000000:
-    from urllib.parse import urlparse, urlunparse, parse_qsl
-else:
-    from urlparse import urlparse, urlunparse, parse_qsl
-
 logger = logging.getLogger(__name__)
 
 
-class Member(namedtuple('Member', 'name,conn_url,api_url,expiration,ttl')):
-
-    @staticmethod
-    def fromNode(node):
-        scheme, netloc, path, params, query, fragment = urlparse(node['value'])
-        conn_url = urlunparse((scheme, netloc, path, params, '', fragment))
-        api_url = ([v for n, v in parse_qsl(query) if n == 'application_name'] or [None])[0]
-        expiration = node.get('expiration', None)
-        ttl = node.get('ttl', None)
-        return Member(node['key'].split('/')[-1], conn_url, api_url, expiration, ttl)
-
-    def real_ttl(self):
-        return calculate_ttl(self.expiration) or -1
+class EtcdError(DCSError):
+    pass
 
 
-class Cluster(namedtuple('Cluster', 'initialize,leader,last_leader_operation,members')):
-
-    def is_unlocked(self):
-        return not (self.leader and self.leader.name)
+class EtcdConnectionFailed(EtcdError):
+    pass
 
 
 class Client:
@@ -204,12 +184,12 @@ class Client:
         return False
 
 
-class Etcd:
+class Etcd(AbstractDCS):
 
-    def __init__(self, config):
+    def __init__(self, name, config):
+        super(Etcd, self).__init__(name, config)
         self.ttl = config['ttl']
         self.member_ttl = config.get('member_ttl', 3600)
-        self._base_path = '/keys/service/' + config['scope']
         self.client = self.get_etcd_client(config)
 
     def get_etcd_client(self, config):
@@ -223,7 +203,7 @@ class Etcd:
         return client
 
     def client_path(self, path):
-        return self._base_path + path
+        return '/keys' + super(Etcd, self).client_path(path)
 
     def get_client_path(self, path):
         return self.client.get(self.client_path(path))
@@ -251,6 +231,13 @@ class Etcd:
                 return n
         return None
 
+    @staticmethod
+    def member(node):
+        conn_url, api_url = parse_connection_string(node['value'])
+        expiration = node.get('expiration', None)
+        ttl = node.get('ttl', None)
+        return Member(node['modifiedIndex'], node['key'].split('/')[-1], conn_url, api_url, expiration, ttl)
+
     def get_cluster(self):
         try:
             response, status_code = self.get_client_path('?recursive=true')
@@ -259,7 +246,7 @@ class Etcd:
                 initialize = True if node else False
                 # get list of members
                 node = self.find_node(response['node'], '/members') or {'nodes': []}
-                members = [Member.fromNode(n) for n in node['nodes']]
+                members = [self.member(n) for n in node['nodes']]
 
                 # get last leader operation
                 last_leader_operation = 0
@@ -278,7 +265,7 @@ class Etcd:
                             leader = m
                             break
                     if not leader:
-                        leader = Member(node['value'], None, None, None, None)
+                        leader = Member(-1, node['value'], None, None, None, None)
 
                 return Cluster(initialize, leader, last_leader_operation, members)
             elif status_code == 404:
@@ -288,28 +275,21 @@ class Etcd:
 
         raise EtcdError('Etcd is not responding properly')
 
-    def current_leader(self):
+    def touch_member(self, connection_string, ttl=None):
         try:
-            cluster = self.get_cluster()
-            return None if cluster.is_unlocked() else cluster.leader
-        except EtcdError:
-            raise CurrentLeaderError('Etcd is not responding properly')
-
-    def touch_member(self, member, connection_string, ttl=None):
-        try:
-            return self.put_client_path('/members/' + member, value=connection_string, ttl=ttl or self.member_ttl)
+            return self.put_client_path('/members/' + self._name, value=connection_string, ttl=ttl or self.member_ttl)
         except EtcdError:
             return False
 
-    def take_leader(self, value):
+    def take_leader(self):
         try:
-            return self.put_client_path('/leader', value=value, ttl=self.ttl)
+            return self.put_client_path('/leader', value=self._name, ttl=self.ttl)
         except EtcdError:
             return False
 
-    def attempt_to_acquire_leader(self, value):
+    def attempt_to_acquire_leader(self):
         try:
-            ret = self.put_client_path('/leader', value=value, ttl=self.ttl, prevExist=False)
+            ret = self.put_client_path('/leader', value=self._name, ttl=self.ttl, prevExist=False)
             ret or logger.info('Could not take out TTL lock')
             return ret
         except EtcdError:
@@ -324,14 +304,11 @@ class Etcd:
             return True
         return False
 
-    def race(self, path, value):
+    def race(self, path):
         try:
-            return self.put_client_path(path, value=value, prevExist=False)
+            return self.put_client_path(path, value=self._name, prevExist=False)
         except EtcdError:
             return False
 
-    def delete_member(self, member):
-        return self.delete_client_path('/members/' + member)
-
-    def delete_leader(self, value):
-        return self.delete_client_path('/leader?prevValue=' + value)
+    def delete_leader(self):
+        return self.delete_client_path('/leader?prevValue=' + self._name)

helpers/ha.py

@@ -1,6 +1,6 @@
 import logging
 
-from helpers.errors import EtcdError
+from helpers.dcs import DCSError
 from psycopg2 import InterfaceError, OperationalError
 
 logger = logging.getLogger(__name__)
@@ -10,17 +10,17 @@ class Ha:
 
     def __init__(self, state_handler, etcd):
         self.state_handler = state_handler
-        self.etcd = etcd
+        self.dcs = etcd
         self.cluster = None
 
-    def load_cluster_from_etcd(self):
-        self.cluster = self.etcd.get_cluster()
+    def load_cluster_from_dcs(self):
+        self.cluster = self.dcs.get_cluster()
 
     def acquire_lock(self):
-        return self.etcd.attempt_to_acquire_leader(self.state_handler.name)
+        return self.dcs.attempt_to_acquire_leader()
 
     def update_lock(self):
-        return self.etcd.update_leader(self.state_handler)
+        return self.dcs.update_leader(self.state_handler)
 
     def has_lock(self):
         lock_owner = self.cluster.leader and self.cluster.leader.name
@@ -35,7 +35,7 @@ class Ha:
 
     def run_cycle(self):
         try:
-            self.load_cluster_from_etcd()
+            self.load_cluster_from_dcs()
             if not self.state_handler.is_healthy():
                 has_lock = self.has_lock()
                 self.state_handler.write_recovery_conf(None if has_lock else self.cluster.leader)
@@ -43,7 +43,7 @@ class Ha:
                 if not has_lock:
                     return 'started as a secondary'
                 logger.info('started as readonly because i had the session lock')
-                self.load_cluster_from_etcd()
+                self.load_cluster_from_dcs()
 
             if self.cluster.is_unlocked():
                 if self.state_handler.is_healthiest_node(self.cluster):
@@ -54,7 +54,7 @@ class Ha:
                             self.state_handler.promote()
                             return 'promoted self to leader by acquiring session lock'
                     else:
-                        self.load_cluster_from_etcd()
+                        self.load_cluster_from_dcs()
                         if self.state_handler.is_leader():
                             self.demote()
                             return 'demoted self due after trying and failing to obtain lock'
@@ -62,7 +62,7 @@ class Ha:
                             self.follow_the_leader()
                             return 'following new leader after trying and failing to obtain lock'
                 else:
-                    self.load_cluster_from_etcd()
+                    self.load_cluster_from_dcs()
                     if self.state_handler.is_leader():
                         self.demote()
                         return 'demoting self because i am not the healthiest node'
@@ -84,10 +84,10 @@ class Ha:
                     else:
                         self.follow_the_leader()
                         return 'no action.  i am a secondary and i am following a leader'
-        except EtcdError:
-            logger.error('Error communicating with Etcd')
+        except DCSError:
+            logger.error('Error communicating with DCS')
             if self.state_handler.is_leader():
                 self.state_handler.demote(None)
-                return 'demoted self because etcd is not accessible and i was a leader'
+                return 'demoted self because DCS is not accessible and i was a leader'
         except (InterfaceError, OperationalError):
             logger.error('Error communicating with Postgresql.  Will try again')

helpers/postgresql.py

@@ -422,4 +422,4 @@ primary_conninfo = '{}'
         self.sync_replication_slots([])
 
     def last_operation(self):
-        return self.xlog_position()
+        return str(self.xlog_position())

helpers/zookeeper.py

@@ -0,0 +1,157 @@
+import logging
+
+from helpers.dcs import AbstractDCS, Cluster, DCSError, Member, parse_connection_string
+from kazoo.client import KazooClient, KazooState
+from kazoo.exceptions import NoNodeError, NodeExistsError
+
+logger = logging.getLogger(__name__)
+
+
+class ZooKeeperError(DCSError):
+    pass
+
+
+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'],
+                                  timeout=(config.get('session_timeout', None) or 30),
+                                  command_retry={
+                                      'deadline': (config.get('reconnect_timeout', None) or 10),
+                                      'max_delay': 1,
+                                      'max_tries': -1},
+                                  connection_retry={'max_delay': 1, 'max_tries': -1})
+        self.client.add_listener(self.session_listener)
+        self.cluster_event = self.client.handler.event_object()
+        self.client.start(None)
+
+    def session_listener(self, state):
+        if state in [KazooState.SUSPENDED, KazooState.LOST]:
+            self.cluster_watcher(None)
+
+    def cluster_watcher(self, event):
+        self.fetch_cluster = True
+        self.cluster_event.set()
+
+    def get_node(self, name, watch=None):
+        try:
+            return self.client.get(self.client_path(name), watch)
+        except NoNodeError:
+            pass
+        except:
+            logger.exception('get_node')
+        return None
+
+    @staticmethod
+    def member(name, value, znode):
+        conn_url, api_url = parse_connection_string(value)
+        return Member(znode.mzxid, name, conn_url, api_url, None, None)
+
+    def load_members(self):
+        members = []
+        for member in self.client.get_children(self.client_path('/members'), self.cluster_watcher):
+            data = self.get_node('/members/' + member)
+            if data is not None:
+                members.append(self.member(member, *data))
+        return members
+
+    def _inner_load_cluster(self):
+        self.cluster_event.clear()
+        leader = self.get_node('/leader', self.cluster_watcher)
+        self.members = self.load_members()
+        if leader:
+            if leader[0] == self._name:
+                client_id = self.client.client_id
+                if client_id is not None and client_id[0] != leader[1].ephemeralOwner:
+                    logger.info('I am leader but not owner of the session. Removing leader node')
+                    self.client.delete(self.client_path('/leader'))
+                    leader = None
+
+            if leader:
+                for member in self.members:
+                    if member.name == leader[0]:
+                        leader = member
+                        self.fetch_cluster = False
+                        break
+            if not isinstance(leader, Member):
+                leader = Member(-1, leader, None, None, None, None)
+        self.leader = leader
+        if self.fetch_cluster:
+            last_leader_operation = self.get_node('/optime/leader')
+            if last_leader_operation:
+                self.last_leader_operation = int(last_leader_operation[0])
+
+    def get_cluster(self):
+        if self.fetch_cluster:
+            try:
+                self.client.retry(self._inner_load_cluster)
+            except:
+                logger.exception('get_cluster')
+                self.session_listener(KazooState.LOST)
+                raise ZooKeeperError('ZooKeeper in not responding properly')
+        return Cluster(True, self.leader, self.last_leader_operation, self.members)
+
+    def _create(self, path, value, **kwargs):
+        try:
+            self.client.retry(self.client.create, self.client_path(path), value, **kwargs)
+            return True
+        except:
+            return False
+
+    def attempt_to_acquire_leader(self):
+        ret = self._create('/leader', self._name, makepath=True, ephemeral=True)
+        ret or logger.info('Could not take out TTL lock')
+        return ret
+
+    def race(self, path):
+        return self._create(path, self._name, makepath=True)
+
+    def touch_member(self, connection_string, ttl=None):
+        for m in self.members:
+            if m.name == self._name:
+                return True
+        path = self.client_path('/members/' + self._name)
+        try:
+            self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+            return True
+        except NodeExistsError:
+            try:
+                self.client.retry(self.client.delete, path)
+                self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True)
+                return True
+            except:
+                logger.exception('touch_member')
+        return False
+
+    def take_leader(self):
+        return self.attempt_to_acquire_leader()
+
+    def update_leader(self, state_handler):
+        last_operation = state_handler.last_operation()
+        if last_operation != self.last_leader_operation:
+            self.last_leader_operation = last_operation
+            path = self.client_path('/optime/leader')
+            try:
+                self.client.retry(self.client.set, path, last_operation)
+            except NoNodeError:
+                try:
+                    self.client.retry(self.client.create, path, last_operation, makepath=True)
+                except:
+                    logger.exception('Failed to create %s', path)
+            except:
+                logger.exception('Failed to update %s', path)
+        return True
+
+    def delete_leader(self):
+        if isinstance(self.leader, Member) and self.leader.name == self._name:
+            self.client.delete(self.client_path('/leader'))
+
+    def sleep(self, timeout):
+        self.cluster_event.wait(timeout)
+        if self.cluster_event.isSet():
+            self.fetch_cluster = True

postgres0.yml

@@ -1,12 +1,18 @@
-loop_wait: 10
+ttl: &ttl 30
+loop_wait: &loop_wait 10
 restapi:
   listen: 127.0.0.1:8008
   connect_address: 127.0.0.1:8008
 etcd:
   scope: batman
-  ttl: 30
+  ttl: *ttl
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#zookeeper:
+#  scope: batman
+#  session_timeout: *ttl
+#  reconnect_timeout: *loop_wait
+#  hosts: 127.0.0.1:2181
 postgresql:
   name: postgresql0
   listen: 127.0.0.1:5432

postgres1.yml

@@ -1,12 +1,18 @@
-loop_wait: 10
+ttl: &ttl 30
+loop_wait: &loop_wait 10
 restapi:
   listen: 127.0.0.1:8009
   connect_address: 127.0.0.1:8009
 etcd:
   scope: batman
-  ttl: 30
+  ttl: *ttl
   host: 127.0.0.1:4001
   #discovery_srv: my-etcd.domain
+#zookeeper:
+#  scope: batman
+#  session_timeout: *ttl
+#  reconnect_timeout: *loop_wait
+#  hosts: 127.0.0.1:2181
 postgresql:
   name: postgresql1
   listen: 127.0.0.1:5433

requirements-py2.txt

@@ -3,3 +3,4 @@ dnspython
 psycopg2
 PyYAML
 requests
+kazoo>=2.2.1

requirements-py3.txt

@@ -3,3 +3,4 @@ dnspython3
 psycopg2
 PyYAML
 requests
+kazoo>=2.2.1

tests/test_etcd.py

@@ -7,8 +7,8 @@ import time
 import unittest
 
 from dns.exception import DNSException
-from helpers.errors import EtcdError, CurrentLeaderError, EtcdConnectionFailed
-from helpers.etcd import Client, Cluster, Etcd, Member
+from helpers.dcs import Cluster, Member
+from helpers.etcd import Client, Etcd, EtcdConnectionFailed, EtcdError
 
 
 class MockResponse:
@@ -109,9 +109,9 @@ class TestMember(unittest.TestCase):
 
     def test_real_ttl(self):
         now = datetime.datetime.utcnow()
-        member = Member('a', 'b', 'c', (now + datetime.timedelta(seconds=2)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
+        member = Member(0, 'a', 'b', 'c', (now + datetime.timedelta(seconds=2)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
         self.assertLess(member.real_ttl(), 2)
-        self.assertEquals(Member('a', 'b', 'c', '', None).real_ttl(), -1)
+        self.assertEquals(Member(0, 'a', 'b', 'c', '', None).real_ttl(), -1)
 
 
 class TestClient(unittest.TestCase):
@@ -180,7 +180,7 @@ class TestEtcd(unittest.TestCase):
         requests.put = requests_put
         requests.delete = requests_delete
         time.sleep = time_sleep
-        self.etcd = Etcd({'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
+        self.etcd = Etcd('foo', {'ttl': 30, 'host': 'localhost:2379', 'scope': 'test'})
 
     def test_get_etcd_client(self):
         time.sleep = time_sleep_exception
@@ -206,16 +206,16 @@ class TestEtcd(unittest.TestCase):
         self.etcd.get_cluster()
 
     def test_current_leader(self):
-        self.assertRaises(CurrentLeaderError, self.etcd.current_leader)
+        self.assertIsNone(self.etcd.current_leader())
 
     def test_touch_member(self):
         self.assertFalse(self.etcd.touch_member('', ''))
 
     def test_take_leader(self):
-        self.assertFalse(self.etcd.take_leader(''))
+        self.assertFalse(self.etcd.take_leader())
 
     def test_attempt_to_acquire_leader(self):
-        self.assertFalse(self.etcd.attempt_to_acquire_leader(''))
+        self.assertFalse(self.etcd.attempt_to_acquire_leader())
 
     def test_update_leader(self):
         url = self.etcd.client._base_uri = self.etcd.client._base_uri.replace('local', 'remote')
@@ -224,7 +224,4 @@ class TestEtcd(unittest.TestCase):
         self.assertFalse(self.etcd.update_leader(MockPostgresql()))
 
     def test_race(self):
-        self.assertFalse(self.etcd.race('', ''))
-
-    def test_delete_member(self):
-        self.assertFalse(self.etcd.delete_member(''))
+        self.assertFalse(self.etcd.race(''))

tests/test_governor.py

@@ -1,4 +1,5 @@
 import datetime
+import helpers.zookeeper
 import psycopg2
 import requests
 import subprocess
@@ -8,10 +9,12 @@ import unittest
 import yaml
 
 from governor import Governor, main
-from helpers.etcd import Cluster, Member
+from helpers.dcs import Cluster, Member
+from helpers.zookeeper import ZooKeeper
+from test_etcd import requests_get, requests_put, requests_delete
 from test_ha import true, false
 from test_postgresql import Postgresql, subprocess_call, psycopg2_connect
-from test_etcd import requests_get, requests_put, requests_delete
+from test_zookeeper import MockKazooClient
 
 if sys.hexversion >= 0x03000000:
     import http.server as BaseHTTPServer
@@ -57,6 +60,11 @@ class TestGovernor(unittest.TestCase):
         Postgresql.write_pg_hba = self.write_pg_hba
         Postgresql.write_recovery_conf = self.write_recovery_conf
 
+    def test_get_dcs(self):
+        helpers.zookeeper.KazooClient = MockKazooClient
+        self.assertIsInstance(self.g.get_dcs('', {'zookeeper': {'scope': '', 'hosts': ''}}), ZooKeeper)
+        self.assertRaises(Exception, self.g.get_dcs, '', {})
+
     def test_governor_main(self):
         main()
         sys.argv = ['governor.py', 'postgres0.yml']
@@ -77,18 +85,18 @@ class TestGovernor(unittest.TestCase):
 
     def test_touch_member(self):
         now = datetime.datetime.utcnow()
-        member = Member(self.g.postgresql.name, 'b', 'c', (now + datetime.timedelta(
+        member = Member(0, self.g.postgresql.name, 'b', 'c', (now + datetime.timedelta(
             seconds=self.g.shutdown_member_ttl + 10)).strftime('%Y-%m-%dT%H:%M:%S.%fZ'), None)
         self.g.ha.cluster = Cluster(True, member, 0, [member])
         self.g.touch_member()
 
     def test_governor_initialize(self):
         self.g.postgresql.should_use_s3_to_create_replica = false
-        self.g.etcd.client._base_uri = 'http://remote'
+        self.g.ha.dcs.client._base_uri = 'http://remote'
         self.g.postgresql.data_directory_empty = true
-        self.g.etcd.race = true
+        self.g.ha.dcs.race = true
         self.g.initialize()
-        self.g.etcd.race = false
+        self.g.ha.dcs.race = false
         self.g.initialize()
         self.g.postgresql.data_directory_empty = false
         self.g.touch_member = self.touch_member

tests/test_ha.py

@@ -1,7 +1,7 @@
 import unittest
 import requests
 
-from helpers.errors import EtcdError
+from helpers.dcs import DCSError
 from helpers.etcd import Cluster, Etcd
 from helpers.ha import Ha
 from test_etcd import requests_get, requests_put, requests_delete
@@ -57,7 +57,7 @@ def nop(*args, **kwargs):
 
 
 def dead_etcd():
-    raise EtcdError('Etcd is not responding properly')
+    raise DCSError('Etcd is not responding properly')
 
 
 class TestHa(unittest.TestCase):
@@ -71,11 +71,11 @@ class TestHa(unittest.TestCase):
         requests.put = requests_put
         requests.delete = requests_delete
         self.p = MockPostgresql()
-        self.e = Etcd({'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
+        self.e = Etcd('foo', {'ttl': 30, 'host': 'remotehost:2379', 'scope': 'test'})
         self.ha = Ha(self.p, self.e)
-        self.ha.load_cluster_from_etcd()
+        self.ha.load_cluster_from_dcs()
         self.ha.cluster = Cluster(False, None, None, [])
-        self.ha.load_cluster_from_etcd = nop
+        self.ha.load_cluster_from_dcs = nop
 
     def test_start_as_slave(self):
         self.p.is_healthy = false
@@ -133,5 +133,5 @@ class TestHa(unittest.TestCase):
         self.assertEquals(self.ha.run_cycle(), 'no action.  i am a secondary and i am following a leader')
 
     def test_no_etcd_connection_master_demote(self):
-        self.ha.load_cluster_from_etcd = dead_etcd
-        self.assertEquals(self.ha.run_cycle(), 'demoted self because etcd is not accessible and i was a leader')
+        self.ha.load_cluster_from_dcs = dead_etcd
+        self.assertEquals(self.ha.run_cycle(), 'demoted self because DCS is not accessible and i was a leader')

tests/test_postgresql.py

@@ -7,7 +7,7 @@ import shutil
 import subprocess
 import unittest
 
-from helpers.etcd import Cluster, Member
+from helpers.dcs import Cluster, Member
 from helpers.postgresql import Postgresql
 
 
@@ -122,9 +122,9 @@ class TestPostgresql(unittest.TestCase):
         psycopg2.connect = psycopg2_connect
         if not os.path.exists(self.p.data_dir):
             os.makedirs(self.p.data_dir)
-        self.leader = Member('leader', 'postgres://replicator:[email protected]:5435/postgres', None, None, 28)
-        self.other = Member('test1', 'postgres://replicator:[email protected]:5433/postgres', None, None, 28)
-        self.me = Member('test0', 'postgres://replicator:[email protected]:5434/postgres', None, None, 28)
+        self.leader = Member(0, 'leader', 'postgres://replicator:[email protected]:5435/postgres', None, None, 28)
+        self.other = Member(0, 'test1', 'postgres://replicator:[email protected]:5433/postgres', None, None, 28)
+        self.me = Member(0, 'test0', 'postgres://replicator:[email protected]:5434/postgres', None, None, 28)
 
     def tear_down(self):
         shutil.rmtree('data')
@@ -195,6 +195,4 @@ class TestPostgresql(unittest.TestCase):
         self.assertTrue(self.p.promote())
 
     def test_last_operation(self):
-        self.assertEquals(self.p.last_operation(), 0)
-
-
+        self.assertEquals(self.p.last_operation(), '0')

tests/test_zookeeper.py

@@ -0,0 +1,138 @@
+import helpers.zookeeper
+import unittest
+
+from helpers.zookeeper import ZooKeeper, ZooKeeperError
+from kazoo.client import KazooState
+from kazoo.exceptions import NoNodeError, NodeExistsError
+from kazoo.protocol.states import ZnodeStat
+from test_etcd import MockPostgresql
+
+
+class MockEvent:
+
+    def clear(self):
+        pass
+
+    def set(self):
+        pass
+
+    def wait(self, timeout):
+        pass
+
+    def isSet(self):
+        return True
+
+
+class MockEventHandler:
+
+    def event_object(self):
+        return MockEvent()
+
+
+class MockKazooClient:
+
+    def __init__(self, **kwargs):
+        self.handler = MockEventHandler()
+        self.leader = False
+        self.exists = True
+
+    def start(self, timeout):
+        pass
+
+    @property
+    def client_id(self):
+        return (-1, '')
+
+    def add_listener(self, cb):
+        pass
+
+    def retry(self, func, *args, **kwargs):
+        func(*args, **kwargs)
+
+    def get(self, path, watch=None):
+        if path == '/service/test/no_node':
+            raise NoNodeError
+        elif path == '/service/test/other_exception':
+            raise Exception()
+        elif '/members/' in path:
+            return (
+                'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/governor',
+                ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)
+            )
+        elif path.endswith('/optime/leader'):
+            return '1'
+        elif path.endswith('/leader'):
+            if self.leader:
+                return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, -1, 0, 0, 0))
+            return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0))
+
+    def get_children(self, path, watch=None, include_data=False):
+        return ['foo', 'bar', 'buzz']
+
+    def create(self, path, value="", acl=None, ephemeral=False, sequence=False, makepath=False):
+        if path.endswith('/initialize') or path == '/service/test/optime/leader':
+            raise Exception
+        elif value == 'retry' or (value == 'exists' and self.exists):
+            raise NodeExistsError
+
+    def set(self, path, value, version=-1):
+        if path == '/service/bla/optime/leader':
+            raise Exception
+        raise NoNodeError
+
+    def delete(self, path, version=-1, recursive=False):
+        self.exists = False
+        if path == '/service/test/leader':
+            if self.leader:
+                return
+            self.leader = True
+            raise Exception
+
+
+class TestZooKeeper(unittest.TestCase):
+
+    def __init__(self, method_name='runTest'):
+        self.setUp = self.set_up
+        super(TestZooKeeper, self).__init__(method_name)
+
+    def set_up(self):
+        helpers.zookeeper.KazooClient = MockKazooClient
+        self.zk = ZooKeeper('foo', {'hosts': 'localhost:2181', 'scope': 'test'})
+
+    def test_session_listener(self):
+        self.zk.session_listener(KazooState.SUSPENDED)
+
+    def test_get_node(self):
+        self.assertIsNone(self.zk.get_node('/no_node'))
+        self.assertIsNone(self.zk.get_node('/other_exception'))
+
+    def test__inner_load_cluster(self):
+        self.zk._base_path = self.zk._base_path.replace('test', 'bla')
+        self.zk._inner_load_cluster()
+
+    def test_get_cluster(self):
+        self.assertRaises(ZooKeeperError, self.zk.get_cluster)
+        self.zk.get_cluster()
+        self.zk.touch_member('foo')
+        self.zk.delete_leader()
+
+    def test_race(self):
+        self.assertFalse(self.zk.race('/initialize'))
+
+    def test_touch_member(self):
+        self.zk.touch_member('new')
+        self.zk.touch_member('exists')
+        self.zk.touch_member('retry')
+
+    def test_take_leader(self):
+        self.zk.take_leader()
+
+    def test_update_leader(self):
+        self.zk.last_leader_operation = -1
+        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+        self.zk._base_path = self.zk._base_path.replace('test', 'bla')
+        self.zk.last_leader_operation = -1
+        self.assertTrue(self.zk.update_leader(MockPostgresql()))
+
+    def test_sleep(self):
+        self.zk.sleep(0)