From 27cb9e0d3ea54095be467c2af52307bbe8da0bdf Mon Sep 17 00:00:00 2001 From: Hrishikesh Yadav Date: Sun, 18 Oct 2026 11:49:16 +0530 Subject: [PATCH] fix(proxy): seed cold spend counter from source cache with SET NX to prevent spend doubling --- litellm/proxy/proxy_server.py | 38 +++++++++++---- .../proxy/proxy_server/test_spend_counters.py | 47 ++++++++++++++++++- 2 files changed, 75 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e3ea9dabe50..5c176abd3ba 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3837,11 +3837,8 @@ async def _prepare_spend_counter_increment( 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). + 3. Seed counter via SET NX under per-counter lock (matching SpendCounterReseed.coalesced) + to avoid double-counting base_spend when concurrent requests cold-start the counter. 4. Increment is returned for the caller to apply via pipeline """ with service_target(SPEND_COUNTERS_TARGET): @@ -3944,9 +3941,34 @@ async def ensure_spend_counter_initialized( ) if db_spend is None: # DB unavailable - fall back to in-process cache (may be stale). - base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) - if base_spend > 0: - await increment_spend_counter_cache(counter_key=counter_key, increment=base_spend) + async with await SpendCounterReseed._get_lock(counter_key): + if await _is_spend_counter_cache_warm(counter_key=counter_key): + return + base_spend: Final = await _get_source_cache_base_spend(source_cache_key=source_cache_key) + if base_spend > 0: + if spend_counter_cache.redis_cache is not None: + seeded: Final = await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, + value=base_spend, + nx=True, + ) + if seeded: + current_value = float(base_spend) + else: + cached: Final = await spend_counter_cache.redis_cache.async_get_cache(key=counter_key) + current_value = float(cached) if cached is not None else float(base_spend) + spend_counter_cache.in_memory_cache.set_cache( + key=counter_key, + value=current_value, + ) + record_spend_counter_value(counter_key, current_value) + else: + cached_spend: Final = spend_counter_cache.in_memory_cache.get_cache(key=counter_key) + seeded_spend: Final = ( + max(base_spend, float(cached_spend)) if cached_spend is not None else base_spend + ) + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=seeded_spend) + record_spend_counter_value(counter_key, seeded_spend) _ensure_spend_counter_initialized: Final = ensure_spend_counter_initialized diff --git a/tests/unit/proxy/proxy_server/test_spend_counters.py b/tests/unit/proxy/proxy_server/test_spend_counters.py index 01400384842..6f27f6aa1db 100644 --- a/tests/unit/proxy/proxy_server/test_spend_counters.py +++ b/tests/unit/proxy/proxy_server/test_spend_counters.py @@ -1326,14 +1326,57 @@ async def test_ensure_spend_counter_initialized_cold_seeds_from_source_cache( observed = { "source_cache_called": fake_user_cache.async_get_cache.called, - "seed_increment_called": fake_cache.redis_cache.async_increment.called, + "seed_set_cache_called": fake_cache.redis_cache.async_set_cache.called, "warm_check_done": fake_cache.redis_cache.async_get_cache.called, } assert normalize(observed) == { "source_cache_called": True, - "seed_increment_called": True, + "seed_set_cache_called": True, "warm_check_done": True, } + fake_cache.redis_cache.async_set_cache.assert_awaited_once_with( + key="spend:user:u", + value=7.0, + nx=True, + ) + assert not fake_cache.redis_cache.async_increment.called + + +@pytest.mark.asyncio +async def test_ensure_spend_counter_initialized_concurrent_cold_seeds_do_not_double_spend( + monkeypatch, +): + redis_store: dict[str, float] = {} + + async def _fake_async_get_cache(key, **kwargs): + return redis_store.get(key) + + async def _fake_async_set_cache(key, value, nx=False, **kwargs): + if nx and key in redis_store: + return False + redis_store[key] = float(value) + return True + + fake_cache = _make_spend_counter_cache() + fake_cache.redis_cache.async_get_cache = AsyncMock(side_effect=_fake_async_get_cache) + fake_cache.redis_cache.async_set_cache = AsyncMock(side_effect=_fake_async_set_cache) + + fake_user_cache = _make_user_api_key_cache(get_value={"spend": 6.0}) + monkeypatch.setattr(ps, "spend_counter_cache", fake_cache) + monkeypatch.setattr(ps, "user_api_key_cache", fake_user_cache) + monkeypatch.setattr(ps, "prisma_client", None) + monkeypatch.setattr(ps.SpendCounterReseed, "coalesced", AsyncMock(return_value=None)) + + counter_key = "spend:end_user:customer-1" + await asyncio.gather( + ps._ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1"), + ps._ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1"), + ps._ensure_spend_counter_initialized(counter_key=counter_key, source_cache_key="customer-1"), + ) + + assert redis_store[counter_key] == 6.0 + assert fake_cache.redis_cache.async_increment.call_count == 0 + # ---------------------------------------------------------------------------