├── README ├── main.py └── zkclient.py /README: -------------------------------------------------------------------------------- 1 | A simple example using zkpython 0.4 2 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | import logging 3 | from os.path import basename, join 4 | 5 | from zkclient import ZKClient, zookeeper, watchmethod 6 | 7 | logging.basicConfig( 8 | level = logging.DEBUG, 9 | format = "[%(asctime)s] %(levelname)-8s %(message)s" 10 | ) 11 | 12 | log = logging 13 | 14 | class GJZookeeper(object): 15 | 16 | ZK_HOST = "localhost:2181" 17 | ROOT = "/app" 18 | WORKERS_PATH = join(ROOT, "workers") 19 | MASTERS_NUM = 1 20 | TIMEOUT = 10000 21 | 22 | def __init__(self, verbose = True): 23 | self.VERBOSE = verbose 24 | self.masters = [] 25 | self.is_master = False 26 | self.path = None 27 | 28 | self.zk = ZKClient(self.ZK_HOST, timeout = self.TIMEOUT) 29 | self.say("login ok!") 30 | # init 31 | self.__init_zk() 32 | # register 33 | self.register() 34 | 35 | def __init_zk(self): 36 | """ 37 | create the zookeeper node if not exist 38 | """ 39 | nodes = (self.ROOT, self.WORKERS_PATH) 40 | for node in nodes: 41 | if not self.zk.exists(node): 42 | try: 43 | self.zk.create(node, "") 44 | except: 45 | pass 46 | 47 | @property 48 | def is_slave(self): 49 | return not self.is_master 50 | 51 | def register(self): 52 | """ 53 | register a node for this worker 54 | """ 55 | self.path = self.zk.create(self.WORKERS_PATH + "/worker", "1", flags=zookeeper.EPHEMERAL | zookeeper.SEQUENCE) 56 | self.path = basename(self.path) 57 | self.say("register ok! I'm %s" % self.path) 58 | # check who is the master 59 | self.get_master() 60 | 61 | def get_master(self): 62 | """ 63 | get children, and check who is the smallest child 64 | """ 65 | @watchmethod 66 | def watcher(event): 67 | self.say("child changed, try to get master again.") 68 | self.get_master() 69 | 70 | children = self.zk.get_children(self.WORKERS_PATH, watcher) 71 | children.sort() 72 | self.say("%s's children: %s" % (self.WORKERS_PATH, children)) 73 | 74 | # check if I'm master 75 | self.masters = children[:self.MASTERS_NUM] 76 | if self.path in self.masters: 77 | self.is_master = True 78 | self.say("I've become master!") 79 | else: 80 | self.say("%s is masters, I'm slave" % self.masters) 81 | 82 | 83 | def say(self, msg): 84 | """ 85 | print messages to screen 86 | """ 87 | if self.VERBOSE: 88 | if self.path: 89 | log.info("[ %s(%s) ] %s" % (self.path, "master" if self.is_master else "slave", msg)) 90 | else: 91 | log.info(msg) 92 | 93 | def main(): 94 | gj_zookeeper = GJZookeeper() 95 | 96 | if __name__ == "__main__": 97 | main() 98 | import time 99 | time.sleep(1000) 100 | -------------------------------------------------------------------------------- /zkclient.py: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | # modfied from https://github.com/phunt/zk-smoketest/blob/master/zkclient.py 3 | 4 | import zookeeper, time, threading 5 | from collections import namedtuple 6 | 7 | DEFAULT_TIMEOUT = 30000 8 | VERBOSE = True 9 | 10 | ZOO_OPEN_ACL_UNSAFE = {"perms":0x1f, "scheme":"world", "id" :"anyone"} 11 | 12 | # Mapping of connection state values to human strings. 13 | STATE_NAME_MAPPING = { 14 | zookeeper.ASSOCIATING_STATE: "associating", 15 | zookeeper.AUTH_FAILED_STATE: "auth-failed", 16 | zookeeper.CONNECTED_STATE: "connected", 17 | zookeeper.CONNECTING_STATE: "connecting", 18 | zookeeper.EXPIRED_SESSION_STATE: "expired", 19 | } 20 | 21 | # Mapping of event type to human string. 22 | TYPE_NAME_MAPPING = { 23 | zookeeper.NOTWATCHING_EVENT: "not-watching", 24 | zookeeper.SESSION_EVENT: "session", 25 | zookeeper.CREATED_EVENT: "created", 26 | zookeeper.DELETED_EVENT: "deleted", 27 | zookeeper.CHANGED_EVENT: "changed", 28 | zookeeper.CHILD_EVENT: "child", 29 | } 30 | 31 | class ZKClientError(Exception): 32 | def __init__(self, value): 33 | self.value = value 34 | def __str__(self): 35 | return repr(self.value) 36 | 37 | class ClientEvent(namedtuple("ClientEvent", 'type, connection_state, path')): 38 | """ 39 | A client event is returned when a watch deferred fires. It denotes 40 | some event on the zookeeper client that the watch was requested on. 41 | """ 42 | 43 | @property 44 | def type_name(self): 45 | return TYPE_NAME_MAPPING[self.type] 46 | 47 | @property 48 | def state_name(self): 49 | return STATE_NAME_MAPPING[self.connection_state] 50 | 51 | def __repr__(self): 52 | return "" % ( 53 | self.type_name, self.path, self.state_name) 54 | 55 | 56 | def watchmethod(func): 57 | def decorated(handle, atype, state, path): 58 | event = ClientEvent(atype, state, path) 59 | return func(event) 60 | return decorated 61 | 62 | class ZKClient(object): 63 | def __init__(self, servers, timeout=DEFAULT_TIMEOUT): 64 | self.timeout = timeout 65 | self.connected = False 66 | self.conn_cv = threading.Condition( ) 67 | self.handle = -1 68 | 69 | self.conn_cv.acquire() 70 | if VERBOSE: print("Connecting to %s" % (servers)) 71 | start = time.time() 72 | self.handle = zookeeper.init(servers, self.connection_watcher, timeout) 73 | self.conn_cv.wait(timeout/1000) 74 | self.conn_cv.release() 75 | 76 | if not self.connected: 77 | raise ZKClientError("Unable to connect to %s" % (servers)) 78 | 79 | if VERBOSE: 80 | print("Connected in %d ms, handle is %d" 81 | % (int((time.time() - start) * 1000), self.handle)) 82 | 83 | def connection_watcher(self, h, type, state, path): 84 | self.handle = h 85 | self.conn_cv.acquire() 86 | self.connected = True 87 | self.conn_cv.notifyAll() 88 | self.conn_cv.release() 89 | 90 | def close(self): 91 | return zookeeper.close(self.handle) 92 | 93 | def create(self, path, data="", flags=0, acl=[ZOO_OPEN_ACL_UNSAFE]): 94 | start = time.time() 95 | result = zookeeper.create(self.handle, path, data, acl, flags) 96 | if VERBOSE: 97 | print("Node %s created in %d ms" 98 | % (path, int((time.time() - start) * 1000))) 99 | return result 100 | 101 | def delete(self, path, version=-1): 102 | start = time.time() 103 | result = zookeeper.delete(self.handle, path, version) 104 | if VERBOSE: 105 | print("Node %s deleted in %d ms" 106 | % (path, int((time.time() - start) * 1000))) 107 | return result 108 | 109 | def get(self, path, watcher=None): 110 | return zookeeper.get(self.handle, path, watcher) 111 | 112 | def exists(self, path, watcher=None): 113 | return zookeeper.exists(self.handle, path, watcher) 114 | 115 | def set(self, path, data="", version=-1): 116 | return zookeeper.set(self.handle, path, data, version) 117 | 118 | def set2(self, path, data="", version=-1): 119 | return zookeeper.set2(self.handle, path, data, version) 120 | 121 | 122 | def get_children(self, path, watcher=None): 123 | return zookeeper.get_children(self.handle, path, watcher) 124 | 125 | def async(self, path = "/"): 126 | return zookeeper.async(self.handle, path) 127 | 128 | def acreate(self, path, callback, data="", flags=0, acl=[ZOO_OPEN_ACL_UNSAFE]): 129 | result = zookeeper.acreate(self.handle, path, data, acl, flags, callback) 130 | return result 131 | 132 | def adelete(self, path, callback, version=-1): 133 | return zookeeper.adelete(self.handle, path, version, callback) 134 | 135 | def aget(self, path, callback, watcher=None): 136 | return zookeeper.aget(self.handle, path, watcher, callback) 137 | 138 | def aexists(self, path, callback, watcher=None): 139 | return zookeeper.aexists(self.handle, path, watcher, callback) 140 | 141 | def aset(self, path, callback, data="", version=-1): 142 | return zookeeper.aset(self.handle, path, data, version, callback) 143 | 144 | watch_count = 0 145 | 146 | """Callable watcher that counts the number of notifications""" 147 | class CountingWatcher(object): 148 | def __init__(self): 149 | self.count = 0 150 | global watch_count 151 | self.id = watch_count 152 | watch_count += 1 153 | 154 | def waitForExpected(self, count, maxwait): 155 | """Wait up to maxwait for the specified count, 156 | return the count whether or not maxwait reached. 157 | 158 | Arguments: 159 | - `count`: expected count 160 | - `maxwait`: max milliseconds to wait 161 | """ 162 | waited = 0 163 | while (waited < maxwait): 164 | if self.count >= count: 165 | return self.count 166 | time.sleep(1.0); 167 | waited += 1000 168 | return self.count 169 | 170 | def __call__(self, handle, typ, state, path): 171 | self.count += 1 172 | if VERBOSE: 173 | print("handle %d got watch for %s in watcher %d, count %d" % 174 | (handle, path, self.id, self.count)) 175 | 176 | """Callable watcher that counts the number of notifications 177 | and verifies that the paths are sequential""" 178 | class SequentialCountingWatcher(CountingWatcher): 179 | def __init__(self, child_path): 180 | CountingWatcher.__init__(self) 181 | self.child_path = child_path 182 | 183 | def __call__(self, handle, typ, state, path): 184 | if not self.child_path(self.count) == path: 185 | raise ZKClientError("handle %d invalid path order %s" % (handle, path)) 186 | CountingWatcher.__call__(self, handle, typ, state, path) 187 | 188 | class Callback(object): 189 | def __init__(self): 190 | self.cv = threading.Condition() 191 | self.callback_flag = False 192 | self.rc = -1 193 | 194 | def callback(self, handle, rc, handler): 195 | self.cv.acquire() 196 | self.callback_flag = True 197 | self.handle = handle 198 | self.rc = rc 199 | handler() 200 | self.cv.notify() 201 | self.cv.release() 202 | 203 | def waitForSuccess(self): 204 | while not self.callback_flag: 205 | self.cv.wait() 206 | self.cv.release() 207 | 208 | if not self.callback_flag == True: 209 | raise ZKClientError("asynchronous operation timed out on handle %d" % 210 | (self.handle)) 211 | if not self.rc == zookeeper.OK: 212 | raise ZKClientError( 213 | "asynchronous operation failed on handle %d with rc %d" % 214 | (self.handle, self.rc)) 215 | 216 | 217 | class GetCallback(Callback): 218 | def __init__(self): 219 | Callback.__init__(self) 220 | 221 | def __call__(self, handle, rc, value, stat): 222 | def handler(): 223 | self.value = value 224 | self.stat = stat 225 | self.callback(handle, rc, handler) 226 | 227 | class SetCallback(Callback): 228 | def __init__(self): 229 | Callback.__init__(self) 230 | 231 | def __call__(self, handle, rc, stat): 232 | def handler(): 233 | self.stat = stat 234 | self.callback(handle, rc, handler) 235 | 236 | class ExistsCallback(SetCallback): 237 | pass 238 | 239 | class CreateCallback(Callback): 240 | def __init__(self): 241 | Callback.__init__(self) 242 | 243 | def __call__(self, handle, rc, path): 244 | def handler(): 245 | self.path = path 246 | self.callback(handle, rc, handler) 247 | 248 | class DeleteCallback(Callback): 249 | def __init__(self): 250 | Callback.__init__(self) 251 | 252 | def __call__(self, handle, rc): 253 | def handler(): 254 | pass 255 | self.callback(handle, rc, handler) 256 | --------------------------------------------------------------------------------