diff --git a/pottery/redlock.py b/pottery/redlock.py index fb62671a..1b4c8641 100644 --- a/pottery/redlock.py +++ b/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) diff --git a/tests/test_redlock.py b/tests/test_redlock.py index 1932c142..4f6ce257 100644 --- a/tests/test_redlock.py +++ b/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))