Implement locks for RedisCluster

redis/client.py

@@ -1082,7 +1082,10 @@ def lock(
         continue trying forever. ``blocking_timeout`` can be specified as a
         float or integer, both representing the number of seconds to wait.
 
-        ``lock_class`` forces the specified lock implementation.
+        ``lock_class`` forces the specified lock implementation. Note that as
+        of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
+        a Lua-based lock). So, it's unlikely you'll need this parameter, unless
+        you have created your own custom lock class.
 
         ``thread_local`` indicates whether the lock token is placed in
         thread-local storage. By default, the token is placed in thread local

redis/cluster.py

@@ -28,6 +28,7 @@
     TimeoutError,
     TryAgainError,
 )
+from redis.lock import Lock
 from redis.utils import (
     dict_merge,
     list_keys_to_dict,
@@ -742,6 +743,72 @@ def pipeline(self, transaction=None, shard_hint=None):
             reinitialize_steps=self.reinitialize_steps,
         )
 
+    def lock(
+        self,
+        name,
+        timeout=None,
+        sleep=0.1,
+        blocking_timeout=None,
+        lock_class=None,
+        thread_local=True,
+    ):
+        """
+        Return a new Lock object using key ``name`` that mimics
+        the behavior of threading.Lock.
+
+        If specified, ``timeout`` indicates a maximum life for the lock.
+        By default, it will remain locked until release() is called.
+
+        ``sleep`` indicates the amount of time to sleep per loop iteration
+        when the lock is in blocking mode and another client is currently
+        holding the lock.
+
+        ``blocking_timeout`` indicates the maximum amount of time in seconds to
+        spend trying to acquire the lock. A value of ``None`` indicates
+        continue trying forever. ``blocking_timeout`` can be specified as a
+        float or integer, both representing the number of seconds to wait.
+
+        ``lock_class`` forces the specified lock implementation. Note that as
+        of redis-py 3.0, the only lock class we implement is ``Lock`` (which is
+        a Lua-based lock). So, it's unlikely you'll need this parameter, unless
+        you have created your own custom lock class.
+
+        ``thread_local`` indicates whether the lock token is placed in
+        thread-local storage. By default, the token is placed in thread local
+        storage so that a thread only sees its token, not a token set by
+        another thread. Consider the following timeline:
+
+            time: 0, thread-1 acquires `my-lock`, with a timeout of 5 seconds.
+                     thread-1 sets the token to "abc"
+            time: 1, thread-2 blocks trying to acquire `my-lock` using the
+                     Lock instance.
+            time: 5, thread-1 has not yet completed. redis expires the lock
+                     key.
+            time: 5, thread-2 acquired `my-lock` now that it's available.
+                     thread-2 sets the token to "xyz"
+            time: 6, thread-1 finishes its work and calls release(). if the
+                     token is *not* stored in thread local storage, then
+                     thread-1 would see the token value as "xyz" and would be
+                     able to successfully release the thread-2's lock.
+
+        In some use cases it's necessary to disable thread local storage. For
+        example, if you have code where one thread acquires a lock and passes
+        that lock instance to a worker thread to release later. If thread
+        local storage isn't disabled in this case, the worker thread won't see
+        the token set by the thread that acquired the lock. Our assumption
+        is that these cases aren't common and as such default to using
+        thread local storage."""
+        if lock_class is None:
+            lock_class = Lock
+        return lock_class(
+            self,
+            name,
+            timeout=timeout,
+            sleep=sleep,
+            blocking_timeout=blocking_timeout,
+            thread_local=thread_local,
+        )
+
     def _determine_nodes(self, *args, **kwargs):
         command = args[0]
         nodes_flag = kwargs.pop("nodes_flag", None)

redis/lock.py

@@ -180,7 +180,7 @@ def acquire(self, blocking=None, blocking_timeout=None, token=None):
         if token is None:
             token = uuid.uuid1().hex.encode()
         else:
-            encoder = self.redis.connection_pool.get_encoder()
+            encoder = self.redis.get_encoder()
             token = encoder.encode(token)
         if blocking is None:
             blocking = self.blocking
@@ -224,7 +224,7 @@ def owned(self):
         # need to always compare bytes to bytes
         # TODO: this can be simplified when the context manager is finished
         if stored_token and not isinstance(stored_token, bytes):
-            encoder = self.redis.connection_pool.get_encoder()
+            encoder = self.redis.get_encoder()
             stored_token = encoder.encode(stored_token)
         return self.local.token is not None and stored_token == self.local.token
 

tests/test_lock.py

@@ -9,7 +9,6 @@
 from .conftest import _get_client
 
 
-@pytest.mark.onlynoncluster
 class TestLock:
     @pytest.fixture()
     def r_decoded(self, request):
@@ -223,7 +222,6 @@ def test_reacquiring_lock_no_longer_owned_raises_error(self, r):
             lock.reacquire()
 
 
-@pytest.mark.onlynoncluster
 class TestLockClassSelection:
     def test_lock_class_argument(self, r):
         class MyLock: