fix: make PodLockManager.release_lock atomic compare-and-delete (re-land #21226) (#24466)

* 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 <noreply@anthropic.com>

* 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 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

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(

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