mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
refactor(redis-store): refresh locks take a user and server target, the store owns the key
This commit is contained in:
Yujong Lee
Devin AI
parent
89d9a7649e
commit
b8e43df44b
5 changed files with 83 additions and 73 deletions
|
|
@ -20,7 +20,7 @@ from typing import Final
|
|||
from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import (
|
||||
LockAcquisition,
|
||||
)
|
||||
from litellm.redis.stores.mcp_credential_refresh import McpRefreshLockStore
|
||||
from litellm.redis.stores.mcp_credential_refresh import McpRefreshLockStore, RefreshLockTarget
|
||||
from litellm.redis.transport.contract import RedisFailure
|
||||
|
||||
|
||||
|
|
@ -28,17 +28,17 @@ from litellm.redis.transport.contract import RedisFailure
|
|||
class RedisDistributedLock:
|
||||
store: McpRefreshLockStore
|
||||
|
||||
async def acquire(self, key: str, token: str, ttl_seconds: float) -> LockAcquisition:
|
||||
acquired: Final = await self.store.acquire(key, token, ttl=timedelta(seconds=ttl_seconds))
|
||||
async def acquire(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> LockAcquisition:
|
||||
acquired: Final = await self.store.acquire(target, token, ttl=timedelta(seconds=ttl_seconds))
|
||||
if isinstance(acquired, RedisFailure):
|
||||
return LockAcquisition.ERROR
|
||||
return LockAcquisition.ACQUIRED if acquired else LockAcquisition.HELD
|
||||
|
||||
async def extend(self, key: str, token: str, ttl_seconds: float) -> bool:
|
||||
return await self.store.extend(key, token, ttl=timedelta(seconds=ttl_seconds)) is True
|
||||
async def extend(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> bool:
|
||||
return await self.store.extend(target, token, ttl=timedelta(seconds=ttl_seconds)) is True
|
||||
|
||||
async def release(self, key: str, token: str) -> None:
|
||||
_ = await self.store.release(key, token)
|
||||
async def release(self, target: RefreshLockTarget, token: str) -> None:
|
||||
_ = await self.store.release(target, token)
|
||||
|
||||
async def is_held(self, key: str) -> bool:
|
||||
return await self.store.is_held(key) is True
|
||||
async def is_held(self, target: RefreshLockTarget) -> bool:
|
||||
return await self.store.is_held(target) is True
|
||||
|
|
|
|||
|
|
@ -31,7 +31,7 @@ from typing import Final, Protocol
|
|||
from litellm.proxy.mcp_server.outbound_credentials.oauth_token_store import (
|
||||
OAuthToken,
|
||||
)
|
||||
from litellm.redis.stores.mcp_credential_refresh import refresh_lock_key
|
||||
from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget
|
||||
|
||||
|
||||
class LockAcquisition(Enum):
|
||||
|
|
@ -51,13 +51,13 @@ class DistributedLock(Protocol):
|
|||
after PX-expiry); ``extend`` refreshes the ``PX`` lease only for the owner; ``is_held`` is
|
||||
``EXISTS`` (so a waiter can poll without taking the lock)."""
|
||||
|
||||
async def acquire(self, key: str, token: str, ttl_seconds: float) -> LockAcquisition: ...
|
||||
async def acquire(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> LockAcquisition: ...
|
||||
|
||||
async def extend(self, key: str, token: str, ttl_seconds: float) -> bool: ...
|
||||
async def extend(self, target: RefreshLockTarget, token: str, ttl_seconds: float) -> bool: ...
|
||||
|
||||
async def release(self, key: str, token: str) -> None: ...
|
||||
async def release(self, target: RefreshLockTarget, token: str) -> None: ...
|
||||
|
||||
async def is_held(self, key: str) -> bool: ...
|
||||
async def is_held(self, target: RefreshLockTarget) -> bool: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -86,11 +86,11 @@ class RedisRefreshCoordinator:
|
|||
refresh: Callable[[], Awaitable[OAuthToken | None]],
|
||||
reread: Callable[[], Awaitable[OAuthToken | None]],
|
||||
) -> OAuthToken | None:
|
||||
key: Final = refresh_lock_key(user_id, server_id)
|
||||
target: Final = RefreshLockTarget(user_id, server_id)
|
||||
token: Final = self.new_token()
|
||||
match await self.lock.acquire(key, token, self.lock_ttl_seconds):
|
||||
match await self.lock.acquire(target, token, self.lock_ttl_seconds):
|
||||
case LockAcquisition.ACQUIRED:
|
||||
return await self._refresh_with_lease_renewal(key, token, refresh)
|
||||
return await self._refresh_with_lease_renewal(target, token, refresh)
|
||||
case LockAcquisition.ERROR:
|
||||
# No election happened (lock backend down), so waiting would just re-read the
|
||||
# still-expired token. Refresh anyway; worst case is an extra refresh, not a stale bearer.
|
||||
|
|
@ -100,34 +100,34 @@ class RedisRefreshCoordinator:
|
|||
# the token it persisted - the winner wrote the fresh token to the store, so a plain
|
||||
# re-read sees it without us refreshing again.
|
||||
deadline: Final = self.clock() + self.wait_timeout_seconds
|
||||
while self.clock() < deadline and await self.lock.is_held(key):
|
||||
while self.clock() < deadline and await self.lock.is_held(target):
|
||||
await self.sleep(self.poll_interval_seconds)
|
||||
return await reread()
|
||||
|
||||
async def _refresh_with_lease_renewal(
|
||||
self,
|
||||
key: str,
|
||||
target: RefreshLockTarget,
|
||||
token: str,
|
||||
refresh: Callable[[], Awaitable[OAuthToken | None]],
|
||||
) -> OAuthToken | None:
|
||||
refresh_task: Final = asyncio.ensure_future(refresh())
|
||||
renewal_task: Final = asyncio.create_task(self._renew_lease_until_done(key, token, refresh_task))
|
||||
renewal_task: Final = asyncio.create_task(self._renew_lease_until_done(target, token, refresh_task))
|
||||
try:
|
||||
return await refresh_task
|
||||
finally:
|
||||
renewal_task.cancel()
|
||||
with suppress(asyncio.CancelledError):
|
||||
await renewal_task
|
||||
await self.lock.release(key, token)
|
||||
await self.lock.release(target, token)
|
||||
|
||||
async def _renew_lease_until_done(
|
||||
self,
|
||||
key: str,
|
||||
target: RefreshLockTarget,
|
||||
token: str,
|
||||
refresh_task: asyncio.Future[OAuthToken | None],
|
||||
) -> None:
|
||||
budget_deadline: Final = self.clock() + self.refresh_budget_seconds
|
||||
while not refresh_task.done() and self.clock() < budget_deadline:
|
||||
await self.sleep(self.lock_ttl_seconds / 2)
|
||||
if not refresh_task.done() and not await self.lock.extend(key, token, self.lock_ttl_seconds):
|
||||
if not refresh_task.done() and not await self.lock.extend(target, token, self.lock_ttl_seconds):
|
||||
return
|
||||
|
|
|
|||
|
|
@ -12,35 +12,41 @@ _EXTEND_IF_OWNER: Final = (
|
|||
)
|
||||
|
||||
|
||||
def refresh_lock_key(user_id: str, server_id: str) -> str:
|
||||
return f"mcp:refresh_lock:{user_id}:{server_id}"
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class RefreshLockTarget:
|
||||
user_id: str
|
||||
server_id: str
|
||||
|
||||
|
||||
def _lock_key(target: RefreshLockTarget) -> str:
|
||||
return f"mcp:refresh_lock:{target.user_id}:{target.server_id}"
|
||||
|
||||
|
||||
class McpRefreshLockStore(Protocol):
|
||||
async def acquire(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
async def acquire(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
|
||||
async def extend(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
async def extend(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure: ...
|
||||
|
||||
async def release(self, key: str, owner: str) -> bool | RedisFailure: ...
|
||||
async def release(self, target: RefreshLockTarget, owner: str) -> bool | RedisFailure: ...
|
||||
|
||||
async def is_held(self, key: str) -> bool | RedisFailure: ...
|
||||
async def is_held(self, target: RefreshLockTarget) -> bool | RedisFailure: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class TransportMcpRefreshLockStore:
|
||||
transport: RedisTransport
|
||||
|
||||
async def acquire(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure:
|
||||
return await self.transport.set(key, owner, ttl=ttl, nx=True)
|
||||
async def acquire(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure:
|
||||
return await self.transport.set(_lock_key(target), owner, ttl=ttl, nx=True)
|
||||
|
||||
async def extend(self, key: str, owner: str, *, ttl: timedelta) -> bool | RedisFailure:
|
||||
reply: Final = await self.transport.eval(_EXTEND_IF_OWNER, (key,), (owner, ttl_ms(ttl)))
|
||||
async def extend(self, target: RefreshLockTarget, owner: str, *, ttl: timedelta) -> bool | RedisFailure:
|
||||
reply: Final = await self.transport.eval(_EXTEND_IF_OWNER, (_lock_key(target),), (owner, ttl_ms(ttl)))
|
||||
return reply if isinstance(reply, RedisFailure) else reply == 1
|
||||
|
||||
async def release(self, key: str, owner: str) -> bool | RedisFailure:
|
||||
reply: Final = await self.transport.eval(_RELEASE_IF_OWNER, (key,), (owner,))
|
||||
async def release(self, target: RefreshLockTarget, owner: str) -> bool | RedisFailure:
|
||||
reply: Final = await self.transport.eval(_RELEASE_IF_OWNER, (_lock_key(target),), (owner,))
|
||||
return reply if isinstance(reply, RedisFailure) else reply == 1
|
||||
|
||||
async def is_held(self, key: str) -> bool | RedisFailure:
|
||||
reply: Final = await self.transport.get(key)
|
||||
async def is_held(self, target: RefreshLockTarget) -> bool | RedisFailure:
|
||||
reply: Final = await self.transport.get(_lock_key(target))
|
||||
return reply if isinstance(reply, RedisFailure) else reply is not None
|
||||
|
|
|
|||
|
|
@ -9,9 +9,12 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_distributed_lock import
|
|||
from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator import (
|
||||
LockAcquisition,
|
||||
)
|
||||
from litellm.redis.stores.mcp_credential_refresh import TransportMcpRefreshLockStore
|
||||
from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget, TransportMcpRefreshLockStore
|
||||
from tests._support.fake_redis import fake_redis_cache
|
||||
|
||||
TARGET = RefreshLockTarget(user_id="u", server_id="s")
|
||||
KEY = "mcp:refresh_lock:u:s"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def server() -> fakeredis.FakeServer:
|
||||
|
|
@ -31,25 +34,25 @@ def _lock(server: fakeredis.FakeServer, namespace: str | None = None) -> RedisDi
|
|||
async def test_acquire_stores_the_token_under_the_namespace_with_the_ttl(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
|
||||
) -> None:
|
||||
assert await _lock(server, namespace="ns").acquire("k", "tok-1", 10.0) is LockAcquisition.ACQUIRED
|
||||
assert await _lock(server, namespace="ns").acquire(TARGET, "tok-1", 10.0) is LockAcquisition.ACQUIRED
|
||||
|
||||
assert await raw.keys("*") == [b"ns:k"]
|
||||
assert await raw.get("ns:k") == b"tok-1"
|
||||
assert 9_000 < await raw.pttl("ns:k") <= 10_000
|
||||
assert await raw.keys("*") == [f"ns:{KEY}".encode()]
|
||||
assert await raw.get(f"ns:{KEY}") == b"tok-1"
|
||||
assert 9_000 < await raw.pttl(f"ns:{KEY}") <= 10_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_reports_held_when_key_already_held(server: fakeredis.FakeServer) -> None:
|
||||
lock = _lock(server)
|
||||
assert await lock.acquire("k", "tok-a", 10.0) is LockAcquisition.ACQUIRED
|
||||
assert await lock.acquire("k", "tok-b", 10.0) is LockAcquisition.HELD
|
||||
assert await lock.acquire(TARGET, "tok-a", 10.0) is LockAcquisition.ACQUIRED
|
||||
assert await lock.acquire(TARGET, "tok-b", 10.0) is LockAcquisition.HELD
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_acquire_reports_error_on_redis_error_distinct_from_held(server: fakeredis.FakeServer) -> None:
|
||||
lock = _lock(server)
|
||||
server.connected = False
|
||||
assert await lock.acquire("k", "tok", 10.0) is LockAcquisition.ERROR
|
||||
assert await lock.acquire(TARGET, "tok", 10.0) is LockAcquisition.ERROR
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -57,26 +60,26 @@ async def test_release_deletes_only_when_the_token_matches(
|
|||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
|
||||
) -> None:
|
||||
lock = _lock(server)
|
||||
await lock.acquire("k", "owner-B", 10.0)
|
||||
await lock.release("k", "owner-A")
|
||||
assert await raw.get("k") == b"owner-B"
|
||||
await lock.release("k", "owner-B")
|
||||
assert await raw.exists("k") == 0
|
||||
await lock.acquire(TARGET, "owner-B", 10.0)
|
||||
await lock.release(TARGET, "owner-A")
|
||||
assert await raw.get(KEY) == b"owner-B"
|
||||
await lock.release(TARGET, "owner-B")
|
||||
assert await raw.exists(KEY) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_lock_the_old_writer_took_is_held_renewed_and_released_by_its_owner(
|
||||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
|
||||
) -> None:
|
||||
await raw.set("ns:k", "owner-old", nx=True, px=10_000)
|
||||
await raw.set(f"ns:{KEY}", "owner-old", nx=True, px=10_000)
|
||||
lock = _lock(server, namespace="ns")
|
||||
|
||||
assert await lock.is_held("k") is True
|
||||
assert await lock.acquire("k", "owner-new", 10.0) is LockAcquisition.HELD
|
||||
assert await lock.extend("k", "owner-old", 30.0) is True
|
||||
assert 29_000 < await raw.pttl("ns:k") <= 30_000
|
||||
await lock.release("k", "owner-old")
|
||||
assert await lock.is_held("k") is False
|
||||
assert await lock.is_held(TARGET) is True
|
||||
assert await lock.acquire(TARGET, "owner-new", 10.0) is LockAcquisition.HELD
|
||||
assert await lock.extend(TARGET, "owner-old", 30.0) is True
|
||||
assert 29_000 < await raw.pttl(f"ns:{KEY}") <= 30_000
|
||||
await lock.release(TARGET, "owner-old")
|
||||
assert await lock.is_held(TARGET) is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -84,22 +87,22 @@ async def test_extend_refreshes_ttl_only_when_the_token_matches(
|
|||
server: fakeredis.FakeServer, raw: fakeredis.FakeAsyncRedis
|
||||
) -> None:
|
||||
lock = _lock(server)
|
||||
await lock.acquire("k", "owner-B", 1.0)
|
||||
assert await lock.extend("k", "owner-A", 10.0) is False
|
||||
assert await raw.pttl("k") <= 1_000
|
||||
assert await lock.extend("k", "owner-B", 10.0) is True
|
||||
assert 9_000 < await raw.pttl("k") <= 10_000
|
||||
await lock.acquire(TARGET, "owner-B", 1.0)
|
||||
assert await lock.extend(TARGET, "owner-A", 10.0) is False
|
||||
assert await raw.pttl(KEY) <= 1_000
|
||||
assert await lock.extend(TARGET, "owner-B", 10.0) is True
|
||||
assert 9_000 < await raw.pttl(KEY) <= 10_000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_extend_release_and_is_held_degrade_on_redis_error(server: fakeredis.FakeServer) -> None:
|
||||
lock = _lock(server)
|
||||
await lock.acquire("k", "tok", 10.0)
|
||||
await lock.acquire(TARGET, "tok", 10.0)
|
||||
server.connected = False
|
||||
|
||||
assert await lock.extend("k", "tok", 10.0) is False
|
||||
await lock.release("k", "tok")
|
||||
assert await lock.is_held("k") is False
|
||||
assert await lock.extend(TARGET, "tok", 10.0) is False
|
||||
await lock.release(TARGET, "tok")
|
||||
assert await lock.is_held(TARGET) is False
|
||||
|
||||
server.connected = True
|
||||
assert await lock.is_held("k") is True
|
||||
assert await lock.is_held(TARGET) is True
|
||||
|
|
|
|||
|
|
@ -11,8 +11,9 @@ from litellm.proxy.mcp_server.outbound_credentials.redis_refresh_coordinator imp
|
|||
LockAcquisition,
|
||||
RedisRefreshCoordinator,
|
||||
)
|
||||
from litellm.redis.stores.mcp_credential_refresh import RefreshLockTarget
|
||||
|
||||
_KEY = "mcp:refresh_lock:u:s"
|
||||
_TARGET = RefreshLockTarget(user_id="u", server_id="s")
|
||||
|
||||
|
||||
class _FakeLock:
|
||||
|
|
@ -71,9 +72,9 @@ async def test_winner_acquires_and_releases_with_the_same_token_and_never_reread
|
|||
coord = RedisRefreshCoordinator(lock, new_token=lambda: "tok")
|
||||
result = await coord.run("u", "s", refresh, reread)
|
||||
assert result is refreshed
|
||||
assert lock.acquire_calls == [(_KEY, "tok", 10.0)]
|
||||
assert lock.acquire_calls == [(_TARGET, "tok", 10.0)]
|
||||
# released with the SAME token it acquired with, so it can only delete its own lock
|
||||
assert lock.released == [(_KEY, "tok")]
|
||||
assert lock.released == [(_TARGET, "tok")]
|
||||
assert reread_calls == []
|
||||
|
||||
|
||||
|
|
@ -90,7 +91,7 @@ async def test_winner_releases_even_when_refresh_raises():
|
|||
coord = RedisRefreshCoordinator(lock, new_token=lambda: "tok")
|
||||
with pytest.raises(RuntimeError):
|
||||
await coord.run("u", "s", refresh, reread)
|
||||
assert lock.released == [(_KEY, "tok")] # the lock is freed even on failure
|
||||
assert lock.released == [(_TARGET, "tok")] # the lock is freed even on failure
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -127,8 +128,8 @@ async def test_winner_renews_the_lock_while_refresh_runs():
|
|||
|
||||
result = await task
|
||||
assert result is not None and result.access_token == "new"
|
||||
assert lock.extend_calls == [(_KEY, "tok", 10.0)]
|
||||
assert lock.released == [(_KEY, "tok")]
|
||||
assert lock.extend_calls == [(_TARGET, "tok", 10.0)]
|
||||
assert lock.released == [(_TARGET, "tok")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue