scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix: make PodLockManager.release_lock atomic compare-and-delete

Browse source 
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>
This commit is contained in:
joereyna 2026-03-23 20:54:54 -07:00
parent 72a461ba4a
commit cd004398a7
No known key found for this signature in database
GPG key ID: 37E09E2BDB5920E5
2 changed files with 89 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

85
litellm/proxy/db/db_transaction_queue/pod_lock_manager.py
View file

@ -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,32 @@ 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):
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)
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(

38
tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py
View file

@ -307,3 +307,41 @@ 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