Implement read-write (shared) lock from ZooKeeper recipe.

Reference: http://zookeeper.apache.org/doc/trunk/recipes.html#Shared+Locks

This is a writer-preference shared lock.

Most of the lock protocol is the same as for an exclusive lock, except that it needs to be parameterized on:
1) Node name - lock "kind" - a write lock (__lock__) or read lock (__rlock__).
2) Which node names block which other node names (read locks block only writers, writers block both).
3) Finding the predecessor node according to the rule in #2.

These changes are integrated to the main Lock class. The changes are backward compatible with existing users of Lock. Users should normally never pass the Lock node_name and exclude_names kwargs directly, they should use WLock and RLock.

1 files changed, 83 insertions, 10 deletions

diff --git a/kazoo/recipe/lock.py b/kazoo/recipe/lock.py
index d351539..f8ccd85 100644
--- a/kazoo/recipe/lock.py
+++ b/kazoo/recipe/lock.py
@@ -69,9 +69,9 @@ class Lock(object):
     acquired will block.
 
     """
-    _NODE_NAME = '__lock__'
 
-    def __init__(self, client, path, identifier=None):
+    def __init__(self, client, path, identifier=None, node_name=None,
+                 exclude_names=None):
         """Create a Kazoo lock.
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
@@ -79,7 +79,13 @@ class Lock(object):
         :param identifier: Name to use for this lock contender. This
                            can be useful for querying to see who the
                            current lock contenders are.
-
+        :param node_name: Node name, after the contender UUID, before the
+                          sequence number. Involved in read/write locks. For a
+                          normal (exclusive) lock, leave unset.
+        :param exclude_names: Node names which exclude this contender when
+                              present at a lower sequence number. Involved in
+                              read/write locks. For a normal (exclusive) lock,
+                              leave unset.
         """
         self.client = client
         self.path = path
@@ -91,11 +97,19 @@ class Lock(object):
 
         self.wake_event = client.handler.event_object()
 
+        if node_name is None:
+            node_name = "__lock__"
+        self.node_name = node_name
+
+        if exclude_names is None:
+            exclude_names = [self.node_name]
+        self.exclude_names = exclude_names
+
         # props to Netflix Curator for this trick. It is possible for our
         # create request to succeed on the server, but for a failure to
         # prevent us from getting back the full path name. We prefix our
         # lock name with a uuid and can check for its presence on retry.
-        self.prefix = uuid.uuid4().hex + self._NODE_NAME
+        self.prefix = uuid.uuid4().hex + self.node_name
         self.create_path = self.path + "/" + self.prefix
 
         self.create_tried = False
@@ -224,14 +238,15 @@ class Lock(object):
                 # node was removed
                 raise ForceRetryError()
 
-            if self.acquired_lock(children, our_index):
+            predecessor = self.predecessor(children, our_index)
+            if not predecessor:
                 return True
 
             if not blocking:
                 return False
 
             # otherwise we are in the mix. watch predecessor and bide our time
-            predecessor = self.path + "/" + children[our_index - 1]
+            predecessor = self.path + "/" + predecessor
             self.client.add_listener(self._watch_session)
             try:
                 if self.client.exists(predecessor, self._watch_predecessor):
@@ -242,8 +257,11 @@ class Lock(object):
             finally:
                 self.client.remove_listener(self._watch_session)
 
-    def acquired_lock(self, children, index):
-        return index == 0
+    def predecessor(self, children, index):
+        for c in children[:index]:
+            if any(n in c for n in self.exclude_names):
+                return c
+        return None
 
     def _watch_predecessor(self, event):
         self.wake_event.set()
@@ -252,8 +270,13 @@ class Lock(object):
         children = self.client.get_children(self.path)
 
         # can't just sort directly: the node names are prefixed by uuids
-        lockname = self._NODE_NAME
-        children.sort(key=lambda c: c[c.find(lockname) + len(lockname):])
+        def _seq(c):
+            for name in ["__lock__", "__rlock__"]:
+                idx = c.find(name)
+                if idx != -1:
+                    return c[idx + len(name):]
+            raise ValueError("Unknown node type: %s" % c)
+        children.sort(key=_seq)
         return children
 
     def _find_node(self):
@@ -323,6 +346,56 @@ class Lock(object):
         self.release()
 
 
+def WLock(*args, **kwargs):
+    """Kazoo read-write lock (writer side).
+
+    Example usage:
+
+    .. code-block:: python
+
+        zk = KazooClient()
+        lock = WLock("/lockpath", "my-identifier")
+        with lock:  # blocks waiting for any outstanding readers and writers
+            # do something with the lock
+
+    This is an implementation of the ZooKeeper "Shared Locks" recipe:
+    http://zookeeper.apache.org/doc/trunk/recipes.html#Shared+Locks
+
+    This is a writer-preference shared lock.
+
+    The lock path passed to WLock and RLock must match for them to communicate.
+    """
+    kwargs["node_name"] = "__lock__"
+    # Both write and read locks exclude new writers.
+    kwargs["exclude_names"] = ["__lock__", "__rlock__"]
+    return Lock(*args, **kwargs)
+
+
+def RLock(*args, **kwargs):
+    """Kazoo read-write lock (reader side).
+
+    Example usage:
+
+    .. code-block:: python
+
+        zk = KazooClient()
+        lock = RLock("/lockpath", "my-identifier")
+        with lock:  # blocks waiting for any outstanding writers
+            # do something with the lock
+
+    This is an implementation of the ZooKeeper "Shared Locks" recipe:
+    http://zookeeper.apache.org/doc/trunk/recipes.html#Shared+Locks
+
+    This is a writer-preference shared lock.
+
+    The lock path passed to WLock and RLock must match for them to communicate.
+    """
+    kwargs["node_name"] = "__rlock__"
+    # Only write locks exclude new readers.
+    kwargs["exclude_names"] = ["__lock__"]
+    return Lock(*args, **kwargs)
+
+
 class Semaphore(object):
     """A Zookeeper-based Semaphore