def mutex(self, args, kwargs, delete=False):
'''Creates the mutex locks and yields the mutex status.'''
conf = self.app.conf
global_timeout = getattr(conf, 'MUTEX_TIMEOUT', None)
items = getcallargs(self.run, *args, **kwargs)
timeout = items.get('mutex_timeout') or global_timeout or 3600
success = False
try:
hosts = getattr(conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
if client.exists(lock_node):
if time.time() - client.get(lock_node)[1].created > timeout:
client.delete(lock_node)
success = True
else:
success = True
except kazoo.exceptions.KazooException:
yield False
else:
if success:
client.create(lock_node, makepath=True)
yield True
if delete:
client.delete(lock_node)
else:
yield False
finally:
try:
client.stop()
client.close()
except kazoo.exceptions.KazooException:
pass
def get_or_create_node(self,
path,
value='',
acl=None,
ephemeral=False,
sequence=False,
makepath=False):
client = None
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
try:
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
if not client.exists(path):
client.create(path,
value=value,
acl=acl,
ephemeral=ephemeral,
sequence=sequence,
makepath=makepath)
return client.get(path)
except Exception as ex:
raise ex
finally:
if client:
client.stop()
client.close()
def mutex(self, args, kwargs, delete=False):
'''Creates the mutex locks and yields the mutex status.'''
conf = self.app.conf
global_timeout = getattr(conf, 'MUTEX_TIMEOUT', None)
items = getcallargs(self.run, *args, **kwargs)
timeout = items.get('mutex_timeout') or global_timeout or 3600
success = False
try:
hosts = getattr(conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
if client.exists(lock_node):
if time.time() - client.get(lock_node)[1].created > timeout:
client.delete(lock_node)
success = True
else:
success = True
except kazoo.exceptions.KazooException:
yield False
else:
if success:
client.create(lock_node, makepath=True)
yield True
if delete:
client.delete(lock_node)
else:
yield False
finally:
try:
client.stop()
client.close()
except kazoo.exceptions.KazooException:
pass
def delete(node):
'''Delete a node (useful for clearning a lock/semaphore)'''
client = kazoo.client.KazooClient('localhost:2181')
client.start()
client.delete(node, recursive=True)
client.stop()
client.close()
def delete(node):
'''Delete a node (useful for clearning a lock/semaphore)'''
client = kazoo.client.KazooClient('localhost:2181')
client.start()
client.delete(node, recursive=True)
client.stop()
client.close()
def get_or_create_node(self,
path,
value='',
acl=None,
ephemeral=False,
sequence=False,
makepath=False):
client = None
try:
client = self.get_kazoo_client()
client.start()
if not client.exists(path):
client.create(path,
value=value,
acl=acl,
ephemeral=ephemeral,
sequence=sequence,
makepath=makepath)
return client.get(path)
except Exception as ex:
raise ex
finally:
if client:
client.stop()
client.close()
def test_bad_create_args(self):
# We need a non-namespaced client for this test
client = self._get_nonchroot_client()
try:
client.start()
self.assertRaises(ValueError, client.create, "1/2", "val1")
finally:
client.stop()
def now(self):
'''Get the catch-up time or current time'''
client = kazoo.client.KazooClient(hosts=self.hosts)
client.start(timeout=1)
if client.exists(self.now_node):
ret = cPickle.loads(client.get(self.now_node)[0])
else:
ret = celery.current_app._get_current_object().now()
client.stop()
client.close()
return ret
def set_value(self, path, value):
client = None
try:
client = self.get_kazoo_client()
client.start()
client.set(path, value)
except Exception as ex:
raise ex
finally:
if client:
client.stop()
client.close()
def set_value(self, path, value):
client = None
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
try:
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
client.set(path, value)
except Exception as ex:
raise ex
finally:
if client:
client.stop()
client.close()
def _remove_lock(self, args, kwargs):
'''Remove mutex for given args and kwargs.'''
client = None
try:
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
if client.exists(lock_node):
client.delete(lock_node)
finally:
if hasattr(client, 'stop'):
client.stop()
client.close()
def _remove_lock(self, args, kwargs):
'''Remove mutex for given args and kwargs.'''
client = None
try:
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
if client.exists(lock_node):
client.delete(lock_node)
finally:
if hasattr(client, 'stop'):
client.stop()
client.close()
def after_return(self, status, retval, task_id, args, kwargs, einfo):
'''Delete lock node of task, regardles of status.'''
if not (self.request.called_directly or self.request.is_eager):
## Only remove the lock if the job was not called locally
client = None
try:
hosts = self._get_app().conf.ZOOKEEPER_HOSTS
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
if client.exists(lock_node):
client.delete(lock_node)
finally:
if hasattr(client, 'stop'):
client.stop()
client.close()
def _lock(self, args, kwargs, use_lock=False):
'''Actual method for creating lock/semaphore.'''
max_leases = kwargs.pop('max_leases', 1)
identifier = kwargs.pop('identifier', None)
blocking = kwargs.pop('blocking', False)
retry = False if blocking else kwargs.pop('retry', False)
if not identifier:
identifier = '%s->%s' % (self.request.id, self.request.hostname)
node_name = 'locks' if use_lock else 'semaphores'
node_path = u'/zkcelery/%s/%s' % (node_name, self.name)
for value in (unicode(x) for x in args):
# This replace here converts a slash into a fraction-slash.
# They look the same but ZooKeeper uses slashes to denote a
# new node and since a value could contain a slash (e.g. a
# uri) we want to make it into a non-reserved character.
node_path += u'.%s' % (value.replace('/', u'\u2044'))
client = lock = None
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
try:
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
if use_lock:
lock = client.Lock(node_path, identifier=identifier)
else:
lock = client.Semaphore(node_path,
identifier=identifier,
max_leases=max_leases)
success = lock.acquire(blocking=blocking)
if retry:
if success:
yield
else:
self.retry(**kwargs)
else:
yield success
except kazoo.exceptions.KazooException:
if retry:
self.retry(**kwargs)
else:
yield False
finally:
if lock:
lock.release()
if client:
client.stop()
client.close()
def list_controls():
'''View current locks/semaphores/mutexes'''
client = kazoo.client.KazooClient('localhost:2181')
client.start()
# Locks
lock_node = u'/zkcelery/locks'
lock_header = False
if client.exists(lock_node):
for lock_name in client.get_children(lock_node):
full_path = u'%s/%s' % (lock_node, lock_name)
lock = client.Lock(full_path)
contenders = lock.contenders()
if contenders:
if not lock_header:
print 'Current Locks'
lock_header = True
print '\t%s - %s' % (full_path, contenders)
# Semaphores
semaphore_node = u'/zkcelery/semaphores'
semaphore_header = False
if client.exists(semaphore_node):
for semaphore_name in client.get_children(semaphore_node):
if '__lock__' not in semaphore_name:
full_path = u'%s/%s' % (semaphore_node, semaphore_name)
semaphore = client.Semaphore(full_path)
lease_holders = semaphore.lease_holders()
if lease_holders:
if not semaphore_header:
print 'Current Semaphores'
semaphore_header = True
print '\t%s - %s' % (full_path, lease_holders)
# Mutex
mutex_node = u'/zkcelery/mutexes'
mutex_header = False
if client.exists(mutex_node):
for mutex_name in client.get_children(mutex_node):
full_path = u'%s/%s' % (mutex_node, mutex_name)
lock_time = time.localtime(client.get(full_path)[1].created)
lock_str = time.strftime('Set at %Y-%m-%d %H:%M:%S', lock_time)
if not mutex_header:
print 'Current Mutexes'
mutex_header = True
print '\t%s - %s' % (full_path, lock_str)
client.stop()
client.close()
def _lock(self, args, kwargs, use_lock=False):
'''Actual method for creating lock/semaphore.'''
max_leases = kwargs.pop('max_leases', 1)
identifier = kwargs.pop('identifier', None)
blocking = kwargs.pop('blocking', False)
retry = False if blocking else kwargs.pop('retry', False)
if not identifier:
identifier = '%s->%s' % (self.request.id, self.request.hostname)
node_name = 'locks' if use_lock else 'semaphores'
node_path = u'/zkcelery/%s/%s' % (node_name, self.name)
for value in (unicode(x) for x in args):
# This replace here converts a slash into a fraction-slash.
# They look the same but ZooKeeper uses slashes to denote a
# new node and since a value could contain a slash (e.g. a
# uri) we want to make it into a non-reserved character.
node_path += u'.%s' % (value.replace('/', u'\u2044'))
client = lock = None
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
try:
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
if use_lock:
lock = client.Lock(node_path, identifier=identifier)
else:
lock = client.Semaphore(node_path, identifier=identifier,
max_leases=max_leases)
success = lock.acquire(blocking=blocking)
if retry:
if success:
yield
else:
self.retry(**kwargs)
else:
yield success
except kazoo.exceptions.KazooException:
if retry:
self.retry(**kwargs)
else:
yield False
finally:
if lock:
lock.release()
if client:
client.stop()
client.close()
def list_controls():
'''View current locks/semaphores/mutexes'''
client = kazoo.client.KazooClient('localhost:2181')
client.start()
# Locks
lock_node = u'/zkcelery/locks'
lock_header = False
if client.exists(lock_node):
for lock_name in client.get_children(lock_node):
full_path = u'%s/%s' % (lock_node, lock_name)
lock = client.Lock(full_path)
contenders = lock.contenders()
if contenders:
if not lock_header:
print 'Current Locks'
lock_header = True
print '\t%s - %s' % (full_path, contenders)
# Semaphores
semaphore_node = u'/zkcelery/semaphores'
semaphore_header = False
if client.exists(semaphore_node):
for semaphore_name in client.get_children(semaphore_node):
if '__lock__' not in semaphore_name:
full_path = u'%s/%s' % (semaphore_node, semaphore_name)
semaphore = client.Semaphore(full_path)
lease_holders = semaphore.lease_holders()
if lease_holders:
if not semaphore_header:
print 'Current Semaphores'
semaphore_header = True
print '\t%s - %s' % (full_path, lease_holders)
# Mutex
mutex_node = u'/zkcelery/mutexes'
mutex_header = False
if client.exists(mutex_node):
for mutex_name in client.get_children(mutex_node):
full_path = u'%s/%s' % (mutex_node, mutex_name)
lock_time = time.localtime(client.get(full_path)[1].created)
lock_str = time.strftime('Set at %Y-%m-%d %H:%M:%S', lock_time)
if not mutex_header:
print 'Current Mutexes'
mutex_header = True
print '\t%s - %s' % (full_path, lock_str)
client.stop()
client.close()
def connect(zoo_nodes, timeout, start_timeout, start_retries, randomize_hosts, chroot): # pylint: disable=R0913
hosts = ",".join(zoo_nodes)
client = Client(hosts=hosts, timeout=timeout, randomize_hosts=randomize_hosts)
client.chroot = chroot
while start_retries is None or start_retries > 0:
remaining = ( "inf" if start_retries is None else start_retries )
_logger.info("Trying to connect to ZK, attempts remaining: %s (timeout: %d; hosts: %s)",
remaining, start_timeout, hosts)
try:
client.start(timeout=start_timeout)
break
except Exception:
_logger.exception("Can't connect to ZK in this time")
if start_retries is not None:
start_retries -= 1
_logger.info("Started ZK client on hosts: %s", hosts)
return client
def mutex(self, args, kwargs, delete=False):
'''Creates the mutex locks and yields the mutex status.'''
client = lock = None
try:
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
identifier = '%s->%s' % (self.request.id, self.request.hostname)
lock = client.Lock(lock_node, identifier=identifier)
mutex_acquired = lock.acquire(blocking=False)
except kazoo.exceptions.KazooException:
yield False
else:
yield mutex_acquired
finally:
if lock:
lock.release()
if client:
client.stop()
client.close()
def mutex(self, args, kwargs, delete=False):
'''Creates the mutex locks and yields the mutex status.'''
client = lock = None
try:
hosts = getattr(self.app.conf, 'ZOOKEEPER_HOSTS', '127.0.0.1:2181')
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
identifier = '%s->%s' % (self.request.id, self.request.hostname)
lock = client.Lock(lock_node, identifier=identifier)
mutex_acquired = lock.acquire(blocking=False)
except kazoo.exceptions.KazooException:
yield False
else:
yield mutex_acquired
finally:
if lock:
lock.release()
if client:
client.stop()
client.close()
def connect(zoo_nodes, timeout, start_timeout, start_retries, randomize_hosts,
chroot): # pylint: disable=R0913
hosts = ",".join(zoo_nodes)
client = Client(hosts=hosts,
timeout=timeout,
randomize_hosts=randomize_hosts)
client.chroot = chroot
while start_retries is None or start_retries > 0:
remaining = ("inf" if start_retries is None else start_retries)
_logger.info(
"Trying to connect to ZK, attempts remaining: %s (timeout: %d; hosts: %s)",
remaining, start_timeout, hosts)
try:
client.start(timeout=start_timeout)
break
except Exception:
_logger.exception("Can't connect to ZK in this time")
if start_retries is not None:
start_retries -= 1
_logger.info("Started ZK client on hosts: %s", hosts)
return client
def init_zookeeper_client(host="zoo1", port=2181, max_retries=5):
"""
Tries to initialize the connection to zookeeper.
Args:
host(str, optional): zookeeper host
port(int, optional): zookeeper port
max_retries(int, optional): How many times should the connection
establishment be retried
Returns:
kazoo.client.KazooClient: An initialized zookeeper client
Raises:
EnvironmentError: If the client cannot connect to zookeeper
"""
zk_logger = logging.getLogger("zookeeper")
zk_logger.info("initializing connection to zookeeper")
hosts = "{:s}:{:d}".format(host, port)
client = kazoo.client.KazooClient(hosts=hosts)
backoff_in_seconds = 1
tries = 0
while tries < max_retries:
try:
zk_logger.info("Trying to connect to {:s}".format(hosts))
client.start()
if client.connected:
zk_logger.info("Connected to zookeeper")
return client
except kazoo.handlers.threading.KazooTimeoutError:
tries += 1
backoff_in_seconds *= 2
zk_logger.warn("Failed to connect to {:s}".format(hosts))
zk_logger.warn("Waiting {:d} seconds to reconnect to {:s}".format(
backoff_in_seconds, hosts))
time.sleep(backoff_in_seconds)
error_message = "Could not connect to {:s}".format(hosts)
zk_logger.error(error_message)
raise EnvironmentError(error_message)
def mutex(self, args, kwargs, delete=False):
'''Creates the mutex locks and yields the mutex status.'''
conf = self._get_app().conf
global_timeout = getattr(conf, 'MUTEX_TIMEOUT', None)
items = inspect.getcallargs(self.run, *args, **kwargs)
timeout = items.get('mutex_timeout') or global_timeout or 3600
success = False
try:
hosts = conf.ZOOKEEPER_HOSTS
client = kazoo.client.KazooClient(hosts=hosts)
client.start()
lock_node = self._get_node(args, kwargs)
if client.exists(lock_node):
if time.time() - float(client.get(lock_node)[0]) > timeout:
client.delete(lock_node)
success = True
else:
success = True
except kazoo.exceptions.KazooException as exc:
logger.exception('Error stopping execution')
yield False
else:
if success:
client.create(lock_node, six.text_type(time.time()).encode(), makepath=True)
yield True
if delete:
client.delete(lock_node)
else:
logger.debug('This task has been locked.')
yield False
finally:
try:
client.stop()
client.close()
except kazoo.exceptions.KazooException as exc:
pass
def __init__(
self,
connection_string=None,
session_timeout=None,
wait = False
):
if session_timeout is None:
session_timeout = 10.0
self.session_timeout = session_timeout
if not connection_string:
connection_string = os.environ.get(
"ZC_ZK_CONNECTION_STRING", "127.0.0.1:2181")
if isinstance(connection_string, basestring):
client = kazoo.client.KazooClient(
connection_string, session_timeout)
started = False
else:
client = connection_string
started = True
self.close = lambda : None
self.client = client
for alias in aliases:
setattr(self, alias, getattr(client, alias))
self.ephemeral = {}
self.state = None
def watch_session(state):
restore = False
logger.info("watch_session %s" % state)
if state == kazoo.protocol.states.KazooState.CONNECTED:
restore = self.state == kazoo.protocol.states.KazooState.LOST
logger.info('connected')
self.state = state
if restore:
@zc.thread.Thread
def restore():
for path, data in list(self.ephemeral.items()):
logger.info("restoring ephemeral %s", path)
try:
self.create(
path, data['data'], data['acl'], ephemeral=True)
except kazoo.exceptions.NodeExistsError:
pass # threads? <shrug>
client.add_listener(watch_session)
if started:
watch_session(client.state)
else:
while 1:
try:
client.start()
except Exception:
logger.critical("Can't connect to ZooKeeper at %r",
connection_string)
if wait:
time.sleep(1)
else:
raise FailedConnect(connection_string)
else:
break
from kazoo import client
from vyper import v
client = client.KazooClient()
client.start()
v.set_config_type("json")
v.add_remote_provider("zookeeper", client, "/config.json")
v.read_remote_config()
print("Hello " + v.get("hello"))
def get_zookeeper_client(ensemble, read_only=True):
client = kazoo.client.KazooClient(hosts=external_storage_string(ensemble),
read_only=read_only)
client.start()
return client
def client(host):
client = kazoo.client.KazooClient(hosts=host)
client.start()
return client
