diff --git a/patroni/zookeeper.py b/patroni/zookeeper.py index 7c01f278..708cbafa 100644 --- a/patroni/zookeeper.py +++ b/patroni/zookeeper.py @@ -4,7 +4,7 @@ import requests import time from kazoo.client import KazooClient, KazooState -from kazoo.exceptions import NoNodeError, NodeExistsError, KazooException +from kazoo.exceptions import NoNodeError, NodeExistsError from patroni.dcs import AbstractDCS, Cluster, DCSError, Leader, Member, parse_connection_string from patroni.utils import sleep from requests.exceptions import RequestException @@ -92,9 +92,8 @@ class ZooKeeper(AbstractDCS): self.client.add_listener(self.session_listener) self.cluster_event = self.client.handler.event_object() + self.cluster = None self.fetch_cluster = True - self.members = [] - self.leader = None self.last_leader_operation = 0 self.client.start(None) @@ -111,28 +110,39 @@ class ZooKeeper(AbstractDCS): try: return self.client.get(key, watch) except NoNodeError: - pass - except: - logger.exception('get_node') - return None + 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 get_children(self, key, watch=None): + try: + return self.client.get_children(key, watch) + except NoNodeError: + return [] + def load_members(self): members = [] - for member in self.client.get_children(self.members_path, self.cluster_watcher): - data = self.get_node(self.member_path) + for member in self.get_children(self.members_path, self.cluster_watcher): + data = self.get_node(self.members_path + 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(self.leader_path, self.cluster_watcher) - self.members = self.load_members() + nodes = set(self.get_children(self.client_path(''))) + + # get initialize flag + initialize = self._INITIALIZE in nodes + + # get list of members + members = self.load_members() if self._MEMBERS[:-1] in nodes else [] + + # get leader + leader = self.get_node(self.leader_path, self.cluster_watcher) if self._LEADER in nodes else None if leader: client_id = self.client.client_id if leader[0] == self._name and client_id is not None and client_id[0] != leader[1].ephemeralOwner: @@ -142,15 +152,14 @@ class ZooKeeper(AbstractDCS): if leader: member = Member(-1, leader[0], None, None, None, None) - member = ([m for m in self.members if m.name == leader[0]] or [member])[0] + member = ([m for m in members if m.name == leader[0]] or [member])[0] leader = Leader(leader[1].mzxid, None, None, member) self.fetch_cluster = member.index == -1 - self.leader = leader - if self.fetch_cluster: - last_leader_operation = self.get_node(self.leader_optime_path) - if last_leader_operation: - self.last_leader_operation = int(last_leader_operation[0]) + # get last leader operation + self.last_leader_operation = self.get_node(self.leader_optime_path) if self.fetch_cluster else None + self.last_leader_operation = 0 if self.last_leader_operation is None else int(self.last_leader_operation[0]) + self.cluster = Cluster(initialize, leader, self.last_leader_operation, members) def get_cluster(self): if self.exhibitor and self.exhibitor.poll(): @@ -163,7 +172,7 @@ class ZooKeeper(AbstractDCS): 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) + return self.cluster def _create(self, path, value, **kwargs): try: @@ -181,9 +190,8 @@ class ZooKeeper(AbstractDCS): return self._create(self.initialize_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 + if self.cluster and any(m.name == self._name for m in self.cluster.members): + return True path = self.member_path try: self.client.retry(self.client.create, path, connection_string, makepath=True, ephemeral=True) @@ -217,16 +225,19 @@ class ZooKeeper(AbstractDCS): return True def delete_leader(self): - if isinstance(self.leader, Leader) and self.leader.name == self._name: - self.client.delete(self.leader_path) + if isinstance(self.cluster, Cluster) and self.cluster.leader.name == self._name: + self.client.delete(self.leader_path, version=self.cluster.leader.index) - def cancel_initialization(self): + def _cancel_initialization(self): node = self.get_node(self.initialize_path) if node and node[0] == self._name: - try: - self.client.retry(self.client.delete, self.initialize_path, version=node[1].mzxid) - except KazooException: - logger.exception("Unable to delete initialize key") + self.client.delete(self.initialize_path, version=node[1].mzxid) + + def cancel_initialization(self): + try: + self.client.retry(self._cancel_initialization) + except: + logger.exception("Unable to delete initialize key") def watch(self, timeout): self.cluster_event.wait(timeout) diff --git a/tests/test_zookeeper.py b/tests/test_zookeeper.py index 0a95d0bc..4b31bad2 100644 --- a/tests/test_zookeeper.py +++ b/tests/test_zookeeper.py @@ -58,8 +58,6 @@ class MockKazooClient: def get(self, path, watch=None): if path == '/no_node': raise NoNodeError - elif path == '/other_exception': - raise Exception() elif '/members/' in path: return ( 'postgres://repuser:rep-pass@localhost:5434/postgres?application_name=http://127.0.0.1:8009/patroni', @@ -75,6 +73,10 @@ class MockKazooClient: return ('foo', ZnodeStat(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)) def get_children(self, path, watch=None, include_data=False): + if path == '/no_node': + raise NoNodeError + elif path in ['/service/bla/', '/service/test/']: + return ['initialize', 'leader', 'members', 'optime'] return ['foo', 'bar', 'buzz'] def create(self, path, value="", acl=None, ephemeral=False, sequence=False, makepath=False): @@ -136,7 +138,9 @@ class TestZooKeeper(unittest.TestCase): def test_get_node(self): self.assertIsNone(self.zk.get_node('/no_node')) - self.assertIsNone(self.zk.get_node('/other_exception')) + + def test_get_children(self): + self.assertListEqual(self.zk.get_children('/no_node'), []) def test__inner_load_cluster(self): self.zk._base_path = self.zk._base_path.replace('test', 'bla')