From 4fb225a61d66c263daa4e43194118d61e6d745f8 Mon Sep 17 00:00:00 2001 From: Charan Rathore Date: Mon, 28 Sep 2026 10:22:22 +0530 Subject: [PATCH] fix: never release pod lock without atomic compare --- .../db_transaction_queue/pod_lock_manager.py | 28 +++-- .../test_pod_lock_manager.py | 100 ++++++------------ 2 files changed, 47 insertions(+), 81 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 bc67617e444..5dfb57fb060 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -121,8 +121,8 @@ end Release the lock if the current pod holds it. 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. + Without Lua compare-and-delete, leaves the lock to expire rather than + risking deletion of a replacement owner's lock. """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping release_lock") @@ -158,8 +158,7 @@ end """ 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. + If scripting is unavailable, leave the lock in place until its TTL expires. """ script_register: Final = getattr(self.redis_cache, "async_register_script", None) if callable(script_register): @@ -171,23 +170,20 @@ end # the Lua equality check matches and the lock is released result = await self._release_lock_script(keys=[lock_key], args=[json.dumps(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. + except Exception as exc: self._release_lock_script = None verbose_proxy_logger.warning( - "Lua compare-and-delete failed for lock_key=%s, falling back to GET+DEL", + "Lua compare-and-delete failed for lock_key=%s; leaving lock to expire: %s", lock_key, + exc, ) + return 0 - current_value = await self.redis_cache.async_get_cache(lock_key) - 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) - return int(result or 0) + verbose_proxy_logger.warning( + "Lua compare-and-delete unavailable for lock_key=%s; leaving lock to expire", + lock_key, + ) + return 0 @staticmethod def _emit_acquired_lock_event(cronjob_id: str, pod_id: str): 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 4684c3213d6..047a43e33d9 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 @@ -112,62 +112,41 @@ async def test_acquire_lock_expired(pod_lock_manager, mock_redis): @pytest.mark.asyncio -async def test_release_lock_success(pod_lock_manager, mock_redis): - """ - Test that the release lock works when the current pod holds the lock - """ - # Mock get_cache to return this pod's ID +async def test_release_lock_without_script_leaves_own_lock_to_expire(pod_lock_manager, mock_redis): mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id - # Mock successful deletion - mock_redis.async_delete_cache.return_value = 1 - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) + await pod_lock_manager.release_lock(cronjob_id="test_job") - # Verify get_cache was called - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - # Verify delete_cache was called - mock_redis.async_delete_cache.assert_called_once_with(lock_key) - - -@pytest.mark.asyncio -async def test_release_lock_different_pod(pod_lock_manager, mock_redis): - """ - Test that the release lock doesn't delete when a different pod holds the lock - """ - # Mock get_cache to return a different pod's ID - mock_redis.async_get_cache.return_value = "different_pod_id" - - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) - - # Verify get_cache was called - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - # Verify delete_cache was NOT called + mock_redis.async_get_cache.assert_not_called() mock_redis.async_delete_cache.assert_not_called() @pytest.mark.asyncio -async def test_release_lock_no_lock(pod_lock_manager, mock_redis): - """ - Test release lock behavior when no lock exists - """ - # Mock get_cache to return None (no lock) - mock_redis.async_get_cache.return_value = None +async def test_release_lock_without_script_preserves_new_owner(pod_lock_manager, mock_redis): + class ReplacedLockCache: + def __init__(self): + self.owner = pod_lock_manager.pod_id + self.reads = 0 + self.deletes = 0 - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) + async def async_get_cache(self, key): + self.reads += 1 + old_owner = self.owner + self.owner = "new-pod" + return old_owner - # Verify get_cache was called - lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_get_cache.assert_called_once_with(lock_key) - # Verify delete_cache was NOT called - mock_redis.async_delete_cache.assert_not_called() + async def async_delete_cache(self, key): + self.deletes += 1 + self.owner = None + return 1 + + cache = ReplacedLockCache() + pod_lock_manager.redis_cache = cache + await pod_lock_manager.release_lock(cronjob_id="test_job") + + assert cache.owner == pod_lock_manager.pod_id + assert cache.reads == 0 + assert cache.deletes == 0 @pytest.mark.asyncio @@ -248,14 +227,11 @@ async def test_bytes_handling(pod_lock_manager, mock_redis): ) assert result == True - # Reset for release test - mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id.encode("utf-8") - mock_redis.async_delete_cache.return_value = 1 - - await pod_lock_manager.release_lock( - cronjob_id="test_job", - ) - mock_redis.async_delete_cache.assert_called_once() + # Release without Lua cannot safely use decoded bytes and a plain delete. + mock_redis.async_get_cache.reset_mock() + await pod_lock_manager.release_lock(cronjob_id="test_job") + mock_redis.async_get_cache.assert_not_called() + mock_redis.async_delete_cache.assert_not_called() @pytest.mark.asyncio @@ -442,23 +418,17 @@ async def test_release_lock_preserves_lock_held_by_other_pod(): @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). - """ +async def test_release_lock_lua_failure_leaves_lock_to_expire(pod_lock_manager, mock_redis): 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 + script_callable.assert_awaited_once_with(keys=[lock_key], args=[json.dumps(pod_lock_manager.pod_id)]) + mock_redis.async_get_cache.assert_not_called() + mock_redis.async_delete_cache.assert_not_called() assert pod_lock_manager._release_lock_script is None