diff --git a/governor.py b/governor.py index b752feb3..1b376270 100755 --- a/governor.py +++ b/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__': diff --git a/helpers/dcs.py b/helpers/dcs.py new file mode 100644 index 00000000..a5e43488 --- /dev/null +++ b/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) diff --git a/helpers/errors.py b/helpers/errors.py deleted file mode 100644 index 3a56a1e0..00000000 --- a/helpers/errors.py +++ /dev/null @@ -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 diff --git a/helpers/etcd.py b/helpers/etcd.py index 3385fbe6..1e971f30 100644 --- a/helpers/etcd.py +++ b/helpers/etcd.py @@ -2,42 +2,26 @@ 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')): +class CurrentLeaderError(EtcdError): + pass - def is_unlocked(self): - return not (self.leader and self.leader.name) + +class EtcdConnectionFailed(EtcdError): + pass class Client: @@ -204,12 +188,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 +207,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 +235,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 +250,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 +269,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 +279,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 +308,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) diff --git a/helpers/ha.py b/helpers/ha.py index 4f0e61bc..3d48505c 100644 --- a/helpers/ha.py +++ b/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') diff --git a/helpers/postgresql.py b/helpers/postgresql.py index 12262456..ffd5dba7 100644 --- a/helpers/postgresql.py +++ b/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()) diff --git a/helpers/zookeeper.py b/helpers/zookeeper.py new file mode 100644 index 00000000..4e7a22c6 --- /dev/null +++ b/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 diff --git a/postgres0.yml b/postgres0.yml index d2a140e4..0cb694a1 100644 --- a/postgres0.yml +++ b/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 diff --git a/postgres1.yml b/postgres1.yml index 44fee37f..98b1cb09 100644 --- a/postgres1.yml +++ b/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 + listen: 127.0.0.1:8010 + connect_address: 127.0.0.1:8010 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 diff --git a/requirements-py2.txt b/requirements-py2.txt index 3e77d016..012d87a1 100644 --- a/requirements-py2.txt +++ b/requirements-py2.txt @@ -3,3 +3,4 @@ dnspython psycopg2 PyYAML requests +kazoo>=2.2.1 diff --git a/requirements-py3.txt b/requirements-py3.txt index 68821d9d..bf481736 100644 --- a/requirements-py3.txt +++ b/requirements-py3.txt @@ -3,3 +3,4 @@ dnspython3 psycopg2 PyYAML requests +kazoo>=2.2.1 diff --git a/tests/test_etcd.py b/tests/test_etcd.py index 9f87b96d..4bb6f510 100644 --- a/tests/test_etcd.py +++ b/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('')) diff --git a/tests/test_governor.py b/tests/test_governor.py index 940fa968..780431c1 100644 --- a/tests/test_governor.py +++ b/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 diff --git a/tests/test_ha.py b/tests/test_ha.py index 574f949f..fcb010c4 100644 --- a/tests/test_ha.py +++ b/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') diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index fc6123a6..d6c01f5b 100644 --- a/tests/test_postgresql.py +++ b/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:rep-pass@127.0.0.1:5435/postgres', None, None, 28) - self.other = Member('test1', 'postgres://replicator:rep-pass@127.0.0.1:5433/postgres', None, None, 28) - self.me = Member('test0', 'postgres://replicator:rep-pass@127.0.0.1:5434/postgres', None, None, 28) + self.leader = Member(0, 'leader', 'postgres://replicator:rep-pass@127.0.0.1:5435/postgres', None, None, 28) + self.other = Member(0, 'test1', 'postgres://replicator:rep-pass@127.0.0.1:5433/postgres', None, None, 28) + self.me = Member(0, 'test0', 'postgres://replicator:rep-pass@127.0.0.1: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') diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py new file mode 100644 index 00000000..3b0bb4a2 --- /dev/null +++ b/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)