fix(proxy): prevent no-Redis spend-counter reseed race

litellm/proxy/db/spend_counter_reseed.py

@@ -8,12 +8,15 @@
 deployments. This module reseeds from the authoritative DB instead.
 
 A per-counter singleflight lock collapses concurrent reseeds on the same pod
-to one DB query per cold-cache window. The lock dict is bounded LRU to cap
-memory in long-lived deployments.
+to one DB query per cold-cache window. The lock registry retains active locks
+and bounds idle entries with LRU eviction in long-lived deployments.
 """
 
 import asyncio
 from collections import OrderedDict
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
 from datetime import datetime
 from typing import TYPE_CHECKING, ClassVar, Optional
 
@@ -36,6 +39,12 @@
     from litellm.proxy.utils import PrismaClient
 
 
+@dataclass
+class _CounterLock:
+    lock: asyncio.Lock
+    users: int = 0
+
+
 class SpendCounterReseed:
     """
     Reseeds spend counters from the authoritative DB and warms the cache,
@@ -53,23 +62,60 @@ class SpendCounterReseed:
     and get_tag_objects_batch(); callers pass those values as fallback_spend.
     """
 
-    _locks: ClassVar["OrderedDict[str, asyncio.Lock]"] = OrderedDict()
+    _locks: ClassVar["OrderedDict[str, _CounterLock]"] = OrderedDict()
     _registry_lock: ClassVar[Optional[asyncio.Lock]] = None
 
     @staticmethod
-    async def _get_lock(counter_key: str) -> asyncio.Lock:
+    def _get_registry_lock() -> asyncio.Lock:
         if SpendCounterReseed._registry_lock is None:
             SpendCounterReseed._registry_lock = asyncio.Lock()
-        async with SpendCounterReseed._registry_lock:
-            lock = SpendCounterReseed._locks.get(counter_key)
-            if lock is not None:
-                SpendCounterReseed._locks.move_to_end(counter_key)
-                return lock
-            lock = asyncio.Lock()
-            SpendCounterReseed._locks[counter_key] = lock
-            if len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE:
-                SpendCounterReseed._locks.popitem(last=False)
-            return lock
+        return SpendCounterReseed._registry_lock
+
+    @staticmethod
+    def _get_or_create_lock(counter_key: str) -> _CounterLock:
+        counter_lock = SpendCounterReseed._locks.get(counter_key)
+        if counter_lock is None:
+            counter_lock = _CounterLock(lock=asyncio.Lock())
+            SpendCounterReseed._locks[counter_key] = counter_lock
+        SpendCounterReseed._locks.move_to_end(counter_key)
+        return counter_lock
+
+    @staticmethod
+    def _prune_idle_locks() -> None:
+        while len(SpendCounterReseed._locks) > SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE:
+            idle_key = next(
+                (
+                    key
+                    for key, counter_lock in SpendCounterReseed._locks.items()
+                    if counter_lock.users == 0 and not counter_lock.lock.locked()
+                ),
+                None,
+            )
+            if idle_key is None:
+                return
+            SpendCounterReseed._locks.pop(idle_key)
+
+    @staticmethod
+    async def _get_lock(counter_key: str) -> asyncio.Lock:
+        async with SpendCounterReseed._get_registry_lock():
+            counter_lock = SpendCounterReseed._get_or_create_lock(counter_key)
+            SpendCounterReseed._prune_idle_locks()
+            return counter_lock.lock
+
+    @staticmethod
+    @asynccontextmanager
+    async def _counter_lock(counter_key: str) -> AsyncIterator[None]:
+        async with SpendCounterReseed._get_registry_lock():
+            counter_lock = SpendCounterReseed._get_or_create_lock(counter_key)
+            counter_lock.users += 1
+        try:
+            lock = await SpendCounterReseed._get_lock(counter_key)
+            async with lock:
+                yield
+        finally:
+            async with SpendCounterReseed._get_registry_lock():
+                counter_lock.users -= 1
+                SpendCounterReseed._prune_idle_locks()
 
     @staticmethod
     async def from_db(prisma_client: Optional["PrismaClient"], counter_key: str) -> Optional[float]:
@@ -152,8 +198,7 @@ async def coalesced(
         Returns the spend value (including 0.0 from a fresh budget reset)
         when the DB read succeeds, or None when the DB is unavailable.
         """
-        lock = await SpendCounterReseed._get_lock(counter_key)
-        async with lock:
+        async with SpendCounterReseed._counter_lock(counter_key):
             # Re-check after acquiring the lock. Skip in-memory on a clean
             # Redis miss - in-memory is per-pod-stale.
             redis_clean_miss = False
@@ -197,7 +242,15 @@ async def coalesced(
                         value=current_value,
                     )
                 else:
-                    await spend_counter_cache.async_increment_cache(key=counter_key, value=db_spend, refresh_ttl=True)
+                    cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
+                    if cached is not None:
+                        current_value = float(cached)
+                    else:
+                        current_value = await spend_counter_cache.async_increment_cache(
+                            key=counter_key,
+                            value=db_spend,
+                            refresh_ttl=True,
+                        )
             except Exception:
                 verbose_proxy_logger.exception(
                     "SpendCounterReseed.coalesced: failed to warm counter %s",
@@ -262,8 +315,7 @@ async def coalesced_window(
         entity_id: str,
         window_start: datetime,
     ) -> Optional[float]:
-        lock = await SpendCounterReseed._get_lock(counter_key)
-        async with lock:
+        async with SpendCounterReseed._counter_lock(counter_key):
             redis_clean_miss = False
             if spend_counter_cache.redis_cache is not None:
                 try:

litellm/proxy/proxy_server.py

@@ -2248,24 +2248,25 @@ async def _repair_stale_spend_counter(counter_key: str, db_spend: float) -> None
     in-memory copy is guarded by a read-compare-write with no await in between,
     so it is atomic within the worker.
     """
-    cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
-    needs_update = True
-    if cached is not None:
-        try:
-            needs_update = float(cached) < db_spend
-        except (TypeError, ValueError):
-            needs_update = True
-    if needs_update:
-        spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
-    if spend_counter_cache.redis_cache is not None:
-        try:
-            await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend)
-        except Exception:
-            verbose_proxy_logger.debug(
-                "Unable to repair stale spend counter %s in Redis",
-                counter_key,
-                exc_info=True,
-            )
+    async with SpendCounterReseed._counter_lock(counter_key):
+        cached = spend_counter_cache.in_memory_cache.get_cache(key=counter_key)
+        needs_update = True
+        if cached is not None:
+            try:
+                needs_update = float(cached) < db_spend
+            except (TypeError, ValueError):
+                needs_update = True
+        if needs_update:
+            spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=db_spend)
+        if spend_counter_cache.redis_cache is not None:
+            try:
+                await spend_counter_cache.redis_cache.async_set_max(key=counter_key, value=db_spend)
+            except Exception:
+                verbose_proxy_logger.debug(
+                    "Unable to repair stale spend counter %s in Redis",
+                    counter_key,
+                    exc_info=True,
+                )
 
 
 async def reseed_spend_counter_from_db(counter_key: str) -> None:
@@ -2638,11 +2639,6 @@ async def _init_and_increment_spend_counter(
     2. If not found, reseed from the DB via `SpendCounterReseed.coalesced`.
        Falls back to the cached object's `.spend` via user_api_key_cache
        only if prisma is unavailable, since that value can lag the flusher.
-    3. Seed counter via async_increment_cache (not async_set_cache) to avoid a
-       check-then-set race: if two pods cold-start simultaneously, both may see
-       the counter as absent and seed it. Using increment means the worst case
-       is over-counting (conservative, blocks slightly early) rather than
-       under-counting (would allow overspend).
     4. Increment atomically (both in-memory + Redis)
     """
     await _ensure_spend_counter_initialized(
@@ -2778,11 +2774,12 @@ async def _increment_spend_counter_cache(counter_key: str, increment: float):
         )
         return current_value
 
-    return await spend_counter_cache.async_increment_cache(
-        key=counter_key,
-        value=increment,
-        refresh_ttl=True,
-    )
+    async with SpendCounterReseed._counter_lock(counter_key):
+        return await spend_counter_cache.async_increment_cache(
+            key=counter_key,
+            value=increment,
+            refresh_ttl=True,
+        )
 
 
 async def _invalidate_spend_counter(counter_key: str):

tests/test_litellm/proxy/proxy_server/test_spend_counters.py

@@ -1179,6 +1179,28 @@ async def test_increment_spend_counter_cache_redis_error_raises_and_invalidates(
     assert fake_cache.redis_cache.async_delete_cache.called is True
 
 
+@pytest.mark.asyncio
+async def test_increment_spend_counter_cache_no_redis_uses_in_memory_increment(
+    monkeypatch,
+):
+    fake_cache = _make_spend_counter_cache(
+        redis_increment_value=44.0,
+        with_redis=False,
+    )
+    monkeypatch.setattr(ps, "spend_counter_cache", fake_cache)
+
+    result = await ps._increment_spend_counter_cache(
+        counter_key="spend:key:k", increment=4.0
+    )
+
+    assert result == 44.0
+    fake_cache.async_increment_cache.assert_awaited_once_with(
+        key="spend:key:k",
+        value=4.0,
+        refresh_ttl=True,
+    )
+
+
 # ---------------------------------------------------------------------------
 # _invalidate_spend_counter
 # ---------------------------------------------------------------------------