# ***** BEGIN LICENSE BLOCK *****
# Version: MPL 1.1/GPL 2.0/LGPL 2.1
#
# The contents of this file are subject to the Mozilla Public License Version
# 1.1 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
# http://www.mozilla.org/MPL/
#
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
#
# The Original Code is Mozilla zktools.
#
# The Initial Developer of the Original Code is Mozilla Foundation.
# Portions created by the Initial Developer are Copyright (C) 2011
# the Initial Developer. All Rights Reserved.
#
# Contributor(s):
# Ben Bangert (bbangert@mozilla.com)
#
# Alternatively, the contents of this file may be used under the terms of
# either the GNU General Public License Version 2 or later (the "GPL"), or
# the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
# in which case the provisions of the GPL or the LGPL are applicable instead
# of those above. If you wish to allow use of your version of this file only
# under the terms of either the GPL or the LGPL, and not to allow others to
# use your version of this file under the terms of the MPL, indicate your
# decision by deleting the provisions above and replace them with the notice
# and other provisions required by the GPL or the LGPL. If you do not delete
# the provisions above, a recipient may use your version of this file under
# the terms of any one of the MPL, the GPL or the LGPL.
#
# ***** END LICENSE BLOCK *****
"""Zookeeper Locking
This module provides a :class:`ZkLock`, which should look familiar to anyone
that has used Python's ``threading.Lock`` class. In addition to normal locking
behavior, revokable shared read/write locks with are also supported. All of the
locks can be revoked as desired. This requires the current lock holder(s) to
release their lock(s).
**Shared Read/Write Locks**
Also known in the Zookeeper Recipes as ``Revocable Shared Locks with Freaking
Laser Beams``, :class:`ZkReadLock` and :class:`ZkWriteLock` locks have been
implemented. A read lock can be acquired as long as no write locks are active,
while a write-lock can only be acquired when there are no other read or write
locks active.
"""
import logging
import threading
import time
import zookeeper
ZOO_OPEN_ACL_UNSAFE = {"perms": 0x1f, "scheme": "world", "id": "anyone"}
IMMEDIATE = object()
log = logging.getLogger(__name__)
class CtxManager(object):
"""A lock context manager"""
def __init__(self, lock_object):
self._lock_object = lock_object
def __enter__(self):
return None # No useful object will be supplied for 'as'
def __exit__(self, exc_type, exc_value, traceback):
self._lock_object.release()
return None # Return a non-true value to propagate exc
[docs]class _LockBase(object):
"""Base lock implementation for subclasses"""
[docs] def __init__(self, connection, lock_name, lock_root='/ZktoolsLocks',
logfile=None):
"""Create a Zookeeper lock object
:param connection: zookeeper connection object
:type connection: ZkConnection instance
:param lock_root: Path to the root lock node to create the locks
under
:type lock_root: string
:param logfile: Path to a file to log the zookeeper stream to
:type logfile: string
"""
self._zk = connection
self._cv = threading.Condition()
self._lock_root = lock_root
self._locks = threading.local()
self._locks.revoked = []
self._log_debug = logging.DEBUG >= log.getEffectiveLevel()
if logfile:
zookeeper.set_log_stream(open(logfile))
else:
zookeeper.set_log_stream(open("/dev/null"))
self._locknode = '%s/%s' % (self._lock_root, lock_name)
self._ensure_lock_dir()
def _ensure_lock_dir(self):
# Ensure our lock dir exists
if self._zk.exists(self._locknode):
return
try:
self._zk.create(self._lock_root, "zktools ZLock dir",
[ZOO_OPEN_ACL_UNSAFE], 0)
except zookeeper.NodeExistsException:
if self._log_debug:
log.debug("Lock node in zookeeper already created")
# Try and create our locking node
try:
self._zk.create(self._locknode, "lock", [ZOO_OPEN_ACL_UNSAFE],
0)
except zookeeper.NodeExistsException:
# Ok if this exists already
pass
[docs] def _acquire_lock(self, node_name, timeout=None, revoke=False):
"""Acquire a lock
Internal function used by read/write lock
:param node_name: Name of the node to use for the lock
:type node_name: str
:param timeout: How long to wait to acquire the lock, set to 0 to
get non-blocking behavior.
:type timeout: int
:param revoke: Whether prior locks should be revoked. Can be set to
True to request and wait for prior locks to release
their lock, or :obj:`IMMEDIATE` to destroy the blocking
read/write locks and attempt to acquire a write lock.
:type revoke: bool or :obj:``IMMEDIATE``
:returns: True if the lock was acquired, False otherwise
:rtype: bool
"""
# First clear out any prior revokation warnings
self._locks.revoked = []
self._locks.removed = []
revoke_lock = self._locks.revoked
# Create a lock node
znode = self._zk.create(self._locknode + node_name, "0",
[ZOO_OPEN_ACL_UNSAFE],
zookeeper.EPHEMERAL | zookeeper.SEQUENCE)
def revoke_watcher(handle, type, state, path):
# This method must be in closure scope to ensure that
# it can append to the thread acquire is called from
# to indicate if this particular thread's lock was
# revoked or removed
if type == zookeeper.CHANGED_EVENT:
data = self._zk.get(path, revoke_watcher)[0]
if data == 'unlock':
revoke_lock.append(True)
elif type == zookeeper.DELETED_EVENT or \
state == zookeeper.EXPIRED_SESSION_STATE:
# Trigger if node was deleted
revoke_lock.append(True)
data = self._zk.get(znode, revoke_watcher)[0]
if data == 'unlock':
revoke_lock.append(True)
keyname = znode[znode.rfind('/') + 1:]
acquired = False
cv = threading.Event()
def lock_watcher(handle, type, state, path):
cv.set()
lock_start = time.time()
first_run = True
while not acquired:
cv.clear()
# Have we been at this longer than the timeout?
if not first_run:
if timeout is not None and time.time() - lock_start > timeout:
try:
self._zk.delete(znode)
except zookeeper.NoNodeException:
pass
return False
# Get all the children of the node
children = self._zk.get_children(self._locknode)
children.sort(key=lambda val: val[val.rfind('-') + 1:])
if len(children) == 0 or not keyname in children:
# Disconnects or other errors can cause this
znode = self._zk.create(
self._locknode + node_name, "0", [ZOO_OPEN_ACL_UNSAFE],
zookeeper.EPHEMERAL | zookeeper.SEQUENCE)
keyname = znode[znode.rfind('/') + 1:]
data = self._zk.get(znode, revoke_watcher)[0]
if data == 'unlock':
revoke_lock.append(True)
continue
acquired, blocking_nodes = self._locks.has_lock(keyname, children)
if acquired:
break
if revoke == IMMEDIATE:
# Remove all prior nodes
for node in blocking_nodes:
try:
self._zk.delete(self._locknode + '/' + node)
except zookeeper.NoNodeException:
pass
continue # Now try again
elif revoke:
# Ask all prior blocking nodes to release
for node in blocking_nodes:
try:
self._zk.set(self._locknode + '/' + node, "unlock")
except zookeeper.NoNodeException:
pass
prior_blocking_node = self._locknode + '/' + blocking_nodes[-1]
exists = self._zk.exists(prior_blocking_node, lock_watcher)
if not exists:
# The node disappeared? Rinse and repeat.
continue
# Wait for a notification from get_children, no longer
# than the timeout
wait_for = None
if timeout is not None:
time_spent = time.time() - lock_start
wait_for = timeout - time_spent
cv.wait(wait_for)
first_run = False
self._locks.lock_node = znode
return CtxManager(self)
[docs] def release(self):
"""Release a lock
:returns: True if the lock was released, or False if it is no
longer valid.
:rtype: bool
"""
self._locks.revoked = []
try:
self._zk.delete(self._locks.lock_node)
del self._locks.lock_node
return True
except (zookeeper.NoNodeException, AttributeError):
return False
[docs] def revoked(self):
"""Indicate if this shared lock has been revoked
:returns: True if the lock has been revoked, False otherwise.
:rtype: bool
"""
return bool(self._locks.revoked)
[docs] def has_lock(self):
"""Check with Zookeeper to see if the lock is acquired
:returns: Whether the lock is acquired or not
:rtype: bool
"""
if not hasattr(self._locks, 'lock_node'):
# So we can check it even if we released
return False
znode = self._locks.lock_node
keyname = znode[znode.rfind('/') + 1:]
# Get all the children of the node
children = self._zk.get_children(self._locknode)
children.sort(key=lambda val: val[val.rfind('-') + 1:])
if keyname not in children:
return False
acquired = self._locks.has_lock(keyname, children)[0]
return bool(acquired)
[docs] def clear(self):
"""Clear out a lock
.. warning::
You must be sure this is a dead lock, as clearing it will
forcably release it.
:returns: True if the lock was cleared, or False if it
is no longer valid.
:rtype: bool
"""
children = self._zk.get_children(self._locknode)
for child in children:
try:
self._zk.delete(self._locknode + '/' + child)
except zookeeper.NoNodeException:
pass
@property
def connected(self):
"""Indicate whether a connection to Zookeeper exists"""
return self._zk.connected
[docs]class ZkLock(_LockBase):
"""Zookeeper Lock
Implements a Zookeeper based lock optionally with lock revocation
should locks be idle for more than a specific set of time.
Example::
from zktools.connection import ZkConnection
from zktools.locking import ZkLock
# Create a connection and a lock
conn = ZkConnection()
my_lock = ZkLock(conn, "my_lock_name")
my_lock.acquire() # wait to acquire lock
# do something with the lock
my_lock.release() # release our lock
"""
[docs] def acquire(self, timeout=None, revoke=False):
"""Acquire a lock
:param timeout: How long to wait to acquire the lock, set to 0 to
get non-blocking behavior.
:type timeout: int
:param revoke: Whether prior locks should be revoked. Can be set to
True to request and wait for prior locks to release
their lock, or :obj:`IMMEDIATE` to destroy the blocking
read/write locks and attempt to acquire a write lock.
:type revoke: bool or :obj:`IMMEDIATE`
:returns: True if the lock was acquired, False otherwise
:rtype: bool
"""
node_name = '/lock-'
self._locks.has_lock = has_write_lock
return self._acquire_lock(node_name, timeout, revoke)
[docs]class ZkReadLock(_LockBase):
"""Shared Zookeeper Read Lock
A read-lock is considered succesful if there are no active write
locks.
This class takes the same initialization parameters as
:class:`ZkLock`.
"""
[docs] def acquire(self, timeout=None, revoke=False):
"""Acquire a shared read lock
:param timeout: How long to wait to acquire the lock, set to 0 to
get non-blocking behavior.
:type timeout: int
:param revoke: Whether prior locks should be revoked. Can be set to
True to request and wait for prior locks to release
their lock, or :obj:`IMMEDIATE` to destroy the blocking
write locks and attempt to acquire a read lock.
:type revoke: bool or :obj:`IMMEDIATE`
:returns: True if the lock was acquired, False otherwise
:rtype: bool
"""
node_name = '/read-'
self._locks.has_lock = has_read_lock
return self._acquire_lock(node_name, timeout, revoke)
[docs]class ZkWriteLock(_LockBase):
"""Shared Zookeeper Write Lock
A write-lock is only succesful if there are no read or write locks
active.
This class takes the same initialization parameters as
:class:`ZkLock`.
"""
[docs] def acquire(self, timeout=None, revoke=False):
"""Acquire a shared write lock
:param timeout: How long to wait to acquire the lock, set to 0 to
get non-blocking behavior.
:type timeout: int
:param revoke: Whether prior locks should be revoked. Can be set to
True to request and wait for prior locks to release
their lock, or :obj:`IMMEDIATE` to destroy the blocking
read/write locks and attempt to acquire a write lock.
:type revoke: bool or :obj:`IMMEDIATE`
:returns: True if the lock was acquired, False otherwise
:rtype: bool
"""
node_name = '/write-'
self._locks.has_lock = has_write_lock
return self._acquire_lock(node_name, timeout, revoke)
[docs]def has_read_lock(keyname, children):
"""Determines if this keyname has a valid read lock
:param keyname: The keyname without full path prefix of the current node
being examined
:type keyname: str
:param children: List of the children nodes at this lock point
:type children: list
"""
prior_nodes = children[:children.index(keyname)]
prior_write_nodes = [x for x in prior_nodes if \
x.startswith('write-')]
if not prior_write_nodes:
return True, None
else:
return False, prior_write_nodes
[docs]def has_write_lock(keyname, children):
"""Determines if this keyname has a valid write lock
:param keyname: The keyname without full path prefix of the current node
being examined
:type keyname: str
:param children: List of the children nodes at this lock point
:type children: list
"""
if keyname == children[0]:
return True, None
return False, children[:children.index(keyname)]