python zookeeper

zoo.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
from kazoo.client import KazooClient
from kazoo.protocol.states import KazooState
from kazoo.exceptions import ZookeeperError, NoNodeError


class ZkHelper:
    def __init__(self, zookeeper_hosts):
        self._zookeeper_hosts = zookeeper_hosts # example: h1:p1,h2:p2,h3:p3
        self._zoo = None

    def _watch_session(self, state):
        if KazooState.LOST == state:
            logging.error("Zookeeper session lost")
            self.reset()
        elif KazooState.SUSPENDED == state:
            logging.warning("Zookeeper session suspended")
        elif KazooState.CONNECTED == state:
            logging.debug("Zookeeper session connected")

    def _start(self):
        if self._zoo is not None:
            return True

        try:
            self._zoo = KazooClient(hosts=self._zookeeper_hosts, timeout=10)
            self._zoo.add_listener(self._watch_session)
            self._zoo.start(timeout=5)
            return True
        except Exception as e:
            self.reset()
            logging.error("Start zookeeper client failed. exception:[%r]" % (e,))
            return False

    def _createMasterNode(self, master_path, node_ip):
        try:
            self._zoo.create(master_path, node_ip, ephemeral=True)
            return True
        except Exception as e:
            logging.warning("Create master node failed. exception:[%r]" % (e,))
            return False

    def reset(self):
        try:
            if self._zoo is not None:
                self._zoo.remove_listener(self._watch_session)
                self._zoo.stop()
                self._zoo.close()
        except Exception as e:
            logging.error("Reset failed. exception:[%r]" % (e,))
        finally:
            self._zoo = None

    def isMasterNode(self, root_path, node_ip):
        try:
            if not self._start():
                logging.error("Judge master node failed, start error")
                return False

            self._zoo.ensure_path(root_path)
            master_path = root_path + "/master"

            try:
                data, stat = self._zoo.get(master_path)
                return data == node_ip
            except NoNodeError:
                return self._createMasterNode(master_path, node_ip)
            except ZookeeperError as e:
                logging.error("Judge master node failed. exception:[%r]" % (e,))
                return False
        except Exception as e:
            logging.error("Judge master node failed. exception:[%r]" % (e,))
            return False

    def getChildNodes(self, path):
        try:
            if not self._start():
                logging.error("Get child nodes failed, start error")
                return False, None

            child_list = self._zoo.get_children(path)
            return True, child_list
        except Exception as e:
            logging.error("Get child nodes failed. exception:[%r]" % (e,))
            return False, None

    def getNodeVal(self, path):
        try:
            if not self._start():
                logging.error("Get node value failed, start error")
                return False, None

            val, stat = self._zoo.get(path)
            return True, val
        except Exception as e:
            logging.error("Get node value failed. exception:[%r]" % (e,))
            return False, None

if __name__ == '__main__':
    zk = ZkHelper("localhost:2181,localhost:2182,localhost:2183")
    import time
    while True:
        print "childs:",zk.getChildNodes("/zoo")
        print "isMaster:", zk.isMasterNode("/zoo", "3.3.3.3")
        time.sleep(3)