from kazoo.client import KazooClient
from time import sleep
zk = KazooClient()
zk.start()
lock = zk.Lock("/lockpath", "my-identifier")
with lock: # blocks waiting for lock acquisition
# do something with the lock
print('get lock')
sleep(2)
lock = zk.ReadLock("/lockpath", "my-identifier")
with lock: # blocks waiting for outstanding writers
# do something with the lock
print('read lock')
import os
import time
from datetime import datetime
from kazoo.client import KazooClient
zk = KazooClient(hosts='0.0.0.0:2181')
zk.start()
logs = []
rlock = zk.ReadLock('lock1', 'paris')
for i in range(1000):
try:
tic = datetime.now()
done = rlock.acquire()
tac = datetime.now() - tic
logs += ['time to acquire read lock :' + str(tac)]
except Exception as e:
logs += ["ReadLock acquisition failed" + str(e)]
if not done:
logs += ["ReadLock acquisition failed" + str(e)]
time.sleep(0.005)
try:
tic = datetime.now()
done = rlock.release()
tac = datetime.now() - tic
logs += ['time to release read lock :' + str(tac)]
except Exception as e:
logs += ["ReadLock release failed" + str(e)]
if not done:
logs += ["ReadLock release failed" + str(e)]
class ZKHandler(object):
def __init__(self, config, logger=None):
"""
Initialize an instance of the ZKHandler class with config
A zk_conn object will be created but not started
A ZKSchema instance will be created
"""
self.encoding = "utf8"
self.coordinators = config["coordinators"]
self.logger = logger
self.zk_conn = KazooClient(hosts=self.coordinators)
self._schema = ZKSchema()
#
# Class meta-functions
#
def coordinators(self):
return str(self.coordinators)
def log(self, message, state=""):
if self.logger is not None:
self.logger.out(message, state)
else:
print(message)
#
# Properties
#
@property
def schema(self):
return self._schema
#
# State/connection management
#
def listener(self, state):
"""
Listen for KazooState changes and log accordingly.
This function does not do anything except for log the state, and Kazoo handles the rest.
"""
if state == KazooState.CONNECTED:
self.log("Connection to Zookeeper resumed", state="o")
else:
self.log(
"Connection to Zookeeper lost with state {}".format(state),
state="w")
def connect(self, persistent=False):
"""
Start the zk_conn object and connect to the cluster
"""
try:
self.zk_conn.start()
if persistent:
self.log("Connection to Zookeeper started", state="o")
self.zk_conn.add_listener(self.listener)
except Exception as e:
raise ZKConnectionException(self, e)
def disconnect(self, persistent=False):
"""
Stop and close the zk_conn object and disconnect from the cluster
The class instance may be reused later (avoids persistent connections)
"""
self.zk_conn.stop()
self.zk_conn.close()
if persistent:
self.log("Connection to Zookeeper terminated", state="o")
#
# Schema helper actions
#
def get_schema_path(self, key):
"""
Get the Zookeeper path for {key} from the current schema based on its format.
If {key} is a tuple of length 2, it's treated as a path plus an item instance of that path (e.g. a node, a VM, etc.).
If {key} is a tuple of length 4, it is treated as a path plus an item instance, as well as another item instance of the subpath.
If {key} is just a string, it's treated as a lone path (mostly used for the 'base' schema group.
Otherwise, returns None since this is not a valid key.
This function also handles the special case where a string that looks like an existing path (i.e. starts with '/') is passed;
in that case it will silently return the same path back. This was mostly a migration functionality and is deprecated.
"""
if isinstance(key, tuple):
# This is a key tuple with both an ipath and an item
if len(key) == 2:
# 2-length normal tuple
ipath, item = key
elif len(key) == 4:
# 4-length sub-level tuple
ipath, item, sub_ipath, sub_item = key
return self.schema.path(ipath, item=item) + self.schema.path(
sub_ipath, item=sub_item)
else:
# This is an invalid key
return None
elif isinstance(key, str):
# This is a key string with just an ipath
ipath = key
item = None
# This is a raw key path, used by backup/restore functionality
if re.match(r"^/", ipath):
return ipath
else:
# This is an invalid key
return None
return self.schema.path(ipath, item=item)
#
# Key Actions
#
def exists(self, key):
"""
Check if a key exists
"""
path = self.get_schema_path(key)
if path is None:
# This path is invalid, this is likely due to missing schema entries, so return False
return False
stat = self.zk_conn.exists(path)
if stat:
return True
else:
return False
def read(self, key):
"""
Read data from a key
"""
try:
path = self.get_schema_path(key)
if path is None:
# This path is invalid; this is likely due to missing schema entries, so return None
return None
return self.zk_conn.get(path)[0].decode(self.encoding)
except NoNodeError:
return None
def write(self, kvpairs):
"""
Create or update one or more keys' data
"""
if type(kvpairs) is not list:
self.log("ZKHandler error: Key-value sequence is not a list",
state="e")
return False
transaction = self.zk_conn.transaction()
for kvpair in kvpairs:
if type(kvpair) is not tuple:
self.log(
"ZKHandler error: Key-value pair '{}' is not a tuple".
format(kvpair),
state="e",
)
return False
key = kvpair[0]
value = kvpair[1]
path = self.get_schema_path(key)
if path is None:
# This path is invalid; this is likely due to missing schema entries, so continue
continue
if not self.exists(key):
# Creating a new key
transaction.create(path, str(value).encode(self.encoding))
else:
# Updating an existing key
data = self.zk_conn.get(path)
version = data[1].version
# Validate the expected version after the execution
new_version = version + 1
# Update the data
transaction.set_data(path, str(value).encode(self.encoding))
# Check the data
try:
transaction.check(path, new_version)
except TypeError:
self.log(
"ZKHandler error: Key '{}' does not match expected version"
.format(path),
state="e",
)
return False
try:
transaction.commit()
return True
except Exception as e:
self.log(
"ZKHandler error: Failed to commit transaction: {}".format(e),
state="e")
return False
def delete(self, keys, recursive=True):
"""
Delete a key or list of keys (defaults to recursive)
"""
if type(keys) is not list:
keys = [keys]
for key in keys:
if self.exists(key):
try:
path = self.get_schema_path(key)
self.zk_conn.delete(path, recursive=recursive)
except Exception as e:
self.log(
"ZKHandler error: Failed to delete key {}: {}".format(
path, e),
state="e",
)
return False
return True
def children(self, key):
"""
Lists all children of a key
"""
try:
path = self.get_schema_path(key)
if path is None:
# This path is invalid; this is likely due to missing schema entries, so return None
return None
return self.zk_conn.get_children(path)
except NoNodeError:
return None
def rename(self, kkpairs):
"""
Rename one or more keys to a new value
"""
if type(kkpairs) is not list:
self.log("ZKHandler error: Key-key sequence is not a list",
state="e")
return False
transaction = self.zk_conn.transaction()
def rename_element(transaction, source_path, destination_path):
data = self.zk_conn.get(source_path)[0]
transaction.create(destination_path, data)
if self.children(source_path):
for child_path in self.children(source_path):
child_source_path = "{}/{}".format(source_path, child_path)
child_destination_path = "{}/{}".format(
destination_path, child_path)
rename_element(transaction, child_source_path,
child_destination_path)
transaction.delete(source_path)
for kkpair in kkpairs:
if type(kkpair) is not tuple:
self.log(
"ZKHandler error: Key-key pair '{}' is not a tuple".format(
kkpair),
state="e",
)
return False
source_key = kkpair[0]
source_path = self.get_schema_path(source_key)
if source_path is None:
# This path is invalid; this is likely due to missing schema entries, so continue
continue
destination_key = kkpair[1]
destination_path = self.get_schema_path(destination_key)
if destination_path is None:
# This path is invalid; this is likely due to missing schema entries, so continue
continue
if not self.exists(source_key):
self.log(
"ZKHander error: Source key '{}' does not exist".format(
source_path),
state="e",
)
return False
if self.exists(destination_key):
self.log(
"ZKHander error: Destination key '{}' already exists".
format(destination_path),
state="e",
)
return False
rename_element(transaction, source_path, destination_path)
try:
transaction.commit()
return True
except Exception as e:
self.log(
"ZKHandler error: Failed to commit transaction: {}".format(e),
state="e")
return False
#
# Lock actions
#
def readlock(self, key):
"""
Acquires a read lock on a key
"""
count = 1
lock = None
path = self.get_schema_path(key)
while True:
try:
lock_id = str(uuid.uuid1())
lock = self.zk_conn.ReadLock(path, lock_id)
break
except NoNodeError:
self.log(
"ZKHandler warning: Failed to acquire read lock on nonexistent path {}"
.format(path),
state="e",
)
return None
except Exception as e:
if count > 5:
self.log(
"ZKHandler warning: Failed to acquire read lock after 5 tries: {}"
.format(e),
state="e",
)
break
else:
time.sleep(0.5)
count += 1
continue
return lock
def writelock(self, key):
"""
Acquires a write lock on a key
"""
count = 1
lock = None
path = self.get_schema_path(key)
while True:
try:
lock_id = str(uuid.uuid1())
lock = self.zk_conn.WriteLock(path, lock_id)
break
except NoNodeError:
self.log(
"ZKHandler warning: Failed to acquire write lock on nonexistent path {}"
.format(path),
state="e",
)
return None
except Exception as e:
if count > 5:
self.log(
"ZKHandler warning: Failed to acquire write lock after 5 tries: {}"
.format(e),
state="e",
)
break
else:
time.sleep(0.5)
count += 1
continue
return lock
def exclusivelock(self, key):
"""
Acquires an exclusive lock on a key
"""
count = 1
lock = None
path = self.get_schema_path(key)
while True:
try:
lock_id = str(uuid.uuid1())
lock = self.zk_conn.Lock(path, lock_id)
break
except NoNodeError:
self.log(
"ZKHandler warning: Failed to acquire exclusive lock on nonexistent path {}"
.format(path),
state="e",
)
return None
except Exception as e:
if count > 5:
self.log(
"ZKHandler warning: Failed to acquire exclusive lock after 5 tries: {}"
.format(e),
state="e",
)
break
else:
time.sleep(0.5)
count += 1
continue
return lock
