From 43b12af3a7ea5307b06273a26bbd0fd8cc712e59 Mon Sep 17 00:00:00 2001 From: Alexander Kukushkin Date: Tue, 7 Jul 2015 12:45:14 +0200 Subject: [PATCH] Implement possibility to work against ZooKeeper This implementation is using the same interface (AbstractDCS) as Etcd class. It means that there should be no problem to implement another plugin to work agains Consul for example. --- governor.py | 5 +- helpers/zookeeper.py | 157 ++++++++++++++++++++++++++++++++++++++++ postgres0.yml | 10 ++- postgres1.yml | 14 +++- requirements-py2.txt | 1 + requirements-py3.txt | 1 + tests/test_governor.py | 10 ++- tests/test_zookeeper.py | 138 +++++++++++++++++++++++++++++++++++ 8 files changed, 328 insertions(+), 8 deletions(-) create mode 100644 helpers/zookeeper.py create mode 100644 tests/test_zookeeper.py diff --git a/governor.py b/governor.py index ec4b39ad..b3c8a48c 100755 --- a/governor.py +++ b/governor.py @@ -7,9 +7,10 @@ import yaml from helpers.api import RestApiServer from helpers.etcd import Etcd -from helpers.postgresql import Postgresql 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: @@ -27,6 +28,8 @@ class Governor: 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): 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 0b5910f9..7dc453b4 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 5d5687d3..1f85542d 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 436f88ab..f2703a77 100644 --- a/requirements-py2.txt +++ b/requirements-py2.txt @@ -2,3 +2,4 @@ dnspython psycopg2 PyYAML requests +kazoo>=2.2.1 diff --git a/requirements-py3.txt b/requirements-py3.txt index 13f53d0e..93428c08 100644 --- a/requirements-py3.txt +++ b/requirements-py3.txt @@ -2,3 +2,4 @@ dnspython3 psycopg2 PyYAML requests +kazoo>=2.2.1 diff --git a/tests/test_governor.py b/tests/test_governor.py index daced98b..ed9d6a51 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 @@ -9,9 +10,11 @@ import yaml from governor import Governor, main 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'] 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)