updated for unlimited redlock extensions

pottery/redlock.py

@@ -253,7 +253,7 @@ def __init__(self,
                 automatically release this Redlock, unless it's already been
                 released
             num_extensions -- the number of times that this Redlock's lease can
-                be extended
+                be extended.  Set to ``-1`` for unlimited extensions.
             context_manager_blocking -- when using this Redlock as a context
                 manager, whether to block when acquiring
             context_manager_timeout -- if context_manager_blocking, how long to
@@ -520,7 +520,10 @@ def extend(self, *, raise_on_redis_errors: bool | None = None) -> None:
             True
             >>> printer_lock.release()
         '''
-        if self._extension_num >= self.num_extensions:
+        if (
+            self.num_extensions != -1
+            and self._extension_num >= self.num_extensions
+        ):
             raise TooManyExtensions(self.key, self.masters)
 
         with BailOutExecutor() as executor:
@@ -542,7 +545,9 @@ def extend(self, *, raise_on_redis_errors: bool | None = None) -> None:
                     )
                 else:
                     if num_masters_extended > len(self.masters) // 2:
-                        self._extension_num += 1
+                        if self.num_extensions != -1:
+                            self._extension_num += 1
+
                         return
 
         self._check_enough_masters_up(raise_on_redis_errors, redis_errors)

tests/test_redlock.py

@@ -118,6 +118,20 @@ def test_extend(redlock: Redlock) -> None:
         with pytest.raises(TooManyExtensions):
             redlock.extend()
 
+    @staticmethod
+    def test_extend_unlimited(redis: Redis) -> None:
+        redlock = Redlock(masters={redis}, key='printer', auto_release_time=.2, num_extensions=-1)
+        assert not redis.exists(redlock.key)
+        with pytest.raises(ExtendUnlockedLock):
+            redlock.extend()
+        assert redlock.acquire()
+        for extension_num in range(Redlock._NUM_EXTENSIONS):
+            redlock.extend()
+        # Extend the lock once more and an exception should not be raised
+        redlock.extend()
+        assert redlock._extension_num == 0
+        redlock.release()
+
     @staticmethod
     def test_acquire_then_release(redlock: Redlock) -> None:
         redis = next(iter(redlock.masters))