From f92490c308030bc69a1345540e1c3ffce8e44d5d Mon Sep 17 00:00:00 2001 From: Joe Reyna Date: Wed, 15 Apr 2026 17:33:21 -0700 Subject: [PATCH] fix: make PodLockManager.release_lock atomic compare-and-delete (re-land #21226) (#24466) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: make PodLockManager.release_lock atomic compare-and-delete Re-lands #21226 (reverted in #21469). release_lock() previously did GET + compare + DEL in separate calls, leaving a window where another pod could reacquire the lock between the GET and DEL, causing a stale owner to delete a live lock. Fix: use a Redis Lua script for atomic compare-and-delete. Script registration is cached per PodLockManager instance. Falls back to the old GET+DEL path for cache backends that don't expose async_register_script. Original revert was due to e2e tests running in CI without Redis. Those tests now carry @pytest.mark.skip(reason="Requires Redis connection.") so this re-land is safe. Co-Authored-By: Claude Sonnet 4.6 * fix: add Lua fallback on execution error + test coverage gaps Address Greptile review feedback on #24466: 1. Wrap Lua script execution in try/except — if Redis clears loaded scripts (restart) or scripting is disabled, fall back to GET+DEL rather than letting the exception propagate and leave the lock held until TTL. Reset cached script handle so the next call re-registers. 2. Add test_release_lock_lua_path_emits_released_event — verifies _emit_released_lock_event is called when Lua path returns 1. 3. Add test_release_lock_falls_back_to_get_del_when_lua_execution_fails — verifies the fallback path is taken and script handle is reset. Co-Authored-By: Claude Sonnet 4.6 --------- Co-authored-by: Claude Sonnet 4.6 --- .../db_transaction_queue/pod_lock_manager.py | 95 ++++++++++++------- .../test_pod_lock_manager.py | 80 ++++++++++++++++ 2 files changed, 141 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index 6435498ae03..5e0ddef9eaa 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -21,9 +21,18 @@ class PodLockManager: Ensures that only one pod can run a cron job at a time. """ + _COMPARE_AND_DELETE_LOCK_SCRIPT = """ +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +else + return 0 +end +""" + def __init__(self, redis_cache: Optional[RedisCache] = None): self.pod_id = str(uuid.uuid4()) self.redis_cache = redis_cache + self._release_lock_script: Optional[Any] = None @staticmethod def get_redis_lock_key(cronjob_id: str) -> str: @@ -107,53 +116,35 @@ class PodLockManager: ): """ Release the lock if the current pod holds it. - Uses get and delete commands to ensure that only the owner can release the lock. + Uses an atomic Lua compare-and-delete to prevent TOCTOU races where a + stale owner could delete a newly reacquired lock. + Falls back to GET + DEL for cache implementations that don't support + script registration. """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping release_lock") return try: - cronjob_id = cronjob_id verbose_proxy_logger.debug( "Pod %s attempting to release Redis lock for cronjob_id=%s", self.pod_id, cronjob_id, ) lock_key = PodLockManager.get_redis_lock_key(cronjob_id) - - current_value = await self.redis_cache.async_get_cache(lock_key) - if current_value is not None: - if isinstance(current_value, bytes): - current_value = current_value.decode("utf-8") - if current_value == self.pod_id: - result = await self.redis_cache.async_delete_cache(lock_key) - if result == 1: - verbose_proxy_logger.info( - "Pod %s successfully released Redis lock for cronjob_id=%s", - self.pod_id, - cronjob_id, - ) - self._emit_released_lock_event( - cronjob_id=cronjob_id, - pod_id=self.pod_id, - ) - else: - verbose_proxy_logger.warning( - "Pod %s failed to release Redis lock for cronjob_id=%s. " - "Lock will expire after its TTL.", - self.pod_id, - cronjob_id, - ) - else: - verbose_proxy_logger.debug( - "Pod %s cannot release Redis lock for cronjob_id=%s because it is held by pod %s", - self.pod_id, - cronjob_id, - current_value, - ) + result = await self._compare_and_delete_lock(lock_key=lock_key) + if result == 1: + verbose_proxy_logger.info( + "Pod %s successfully released Redis lock for cronjob_id=%s", + self.pod_id, + cronjob_id, + ) + self._emit_released_lock_event( + cronjob_id=cronjob_id, + pod_id=self.pod_id, + ) else: verbose_proxy_logger.debug( - "Pod %s attempted to release Redis lock for cronjob_id=%s, but no lock was found", + "Pod %s failed to release Redis lock for cronjob_id=%s (lock missing or held by another pod)", self.pod_id, cronjob_id, ) @@ -162,6 +153,42 @@ class PodLockManager: f"Error releasing Redis lock for {cronjob_id}: {e}" ) + async def _compare_and_delete_lock(self, lock_key: str) -> int: + """ + Atomically delete lock key only if current pod owns it. + + Falls back to get/delete for non-RedisCache implementations that do not + expose Lua script registration. + """ + script_register = getattr(self.redis_cache, "async_register_script", None) + if callable(script_register): + try: + if self._release_lock_script is None: + self._release_lock_script = script_register( + self._COMPARE_AND_DELETE_LOCK_SCRIPT + ) + result = await self._release_lock_script( + keys=[lock_key], args=[self.pod_id] + ) + return int(result or 0) + except Exception: + # Lua execution failed (e.g. Redis restart cleared loaded scripts, + # or scripting is disabled). Reset cached script handle and fall + # through to the GET + DEL fallback so the lock is still released. + self._release_lock_script = None + verbose_proxy_logger.warning( + "Lua compare-and-delete failed for lock_key=%s, falling back to GET+DEL", + lock_key, + ) + + current_value = await self.redis_cache.async_get_cache(lock_key) # type: ignore + if isinstance(current_value, bytes): + current_value = current_value.decode("utf-8") + if current_value != self.pod_id: + return 0 + result = await self.redis_cache.async_delete_cache(lock_key) # type: ignore + return int(result or 0) + @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): asyncio.create_task( diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index e83fd75c3a0..eb795fbc012 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -307,3 +307,83 @@ async def test_lock_takeover_race_condition(mock_redis): cronjob_id="test_job", ) assert result2 == False + + +@pytest.mark.asyncio +async def test_release_lock_uses_atomic_compare_delete_script_when_available( + pod_lock_manager, mock_redis +): + """ + Test that release_lock prefers atomic compare-and-delete Lua script when + redis cache exposes script registration. + """ + script_callable = AsyncMock(return_value=1) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + await pod_lock_manager.release_lock(cronjob_id="test_job") + + lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") + mock_redis.async_register_script.assert_called_once_with( + PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT + ) + script_callable.assert_called_once_with( + keys=[lock_key], args=[pod_lock_manager.pod_id] + ) + mock_redis.async_get_cache.assert_not_called() + mock_redis.async_delete_cache.assert_not_called() + + +@pytest.mark.asyncio +async def test_release_lock_reuses_registered_script(pod_lock_manager, mock_redis): + """ + Test script registration is cached on manager instance and reused. + """ + script_callable = AsyncMock(return_value=0) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + await pod_lock_manager.release_lock(cronjob_id="test_job") + await pod_lock_manager.release_lock(cronjob_id="test_job") + + assert mock_redis.async_register_script.call_count == 1 + + +@pytest.mark.asyncio +async def test_release_lock_lua_path_emits_released_event( + pod_lock_manager, mock_redis +): + """ + Test that _emit_released_lock_event is called when the Lua path returns 1 + (successful release). + """ + script_callable = AsyncMock(return_value=1) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + + with patch.object(pod_lock_manager, "_emit_released_lock_event") as mock_emit: + await pod_lock_manager.release_lock(cronjob_id="test_job") + + mock_emit.assert_called_once_with( + cronjob_id="test_job", pod_id=pod_lock_manager.pod_id + ) + + +@pytest.mark.asyncio +async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails( + pod_lock_manager, mock_redis +): + """ + Test that release_lock falls back to GET+DEL when Lua script execution + raises (e.g. Redis restart cleared loaded scripts). + """ + script_callable = AsyncMock(side_effect=Exception("NOSCRIPT")) + mock_redis.async_register_script = MagicMock(return_value=script_callable) + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id + mock_redis.async_delete_cache.return_value = 1 + + await pod_lock_manager.release_lock(cronjob_id="test_job") + + # Lua failed — should have fallen back to GET+DEL + lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") + mock_redis.async_get_cache.assert_called_once_with(lock_key) + mock_redis.async_delete_cache.assert_called_once_with(lock_key) + # Cached script handle should be reset so next call re-registers + assert pod_lock_manager._release_lock_script is None