From 093ceb576d7d9743bbe6ca061791884084143916 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:54:19 -0700 Subject: [PATCH] fix(proxy): do not requeue a daily spend batch whose commit already left for postgres (#42786) * fix(proxy): do not requeue a daily spend batch whose commit already left for postgres Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): settle an interrupted daily spend commit from the shutdown flush instead of blocking the cancelled tick Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): burst two workers and SIGTERM during daily spend COMMIT, expect exactly once Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/db/db_spend_update_writer.py | 114 ++++++++- .../daily_spend_update_queue.py | 14 ++ .../integration/spend/test_shutdown_flush.py | 132 +++++++++- .../proxy/db/test_db_spend_update_writer.py | 238 ++++++++++++++++++ 4 files changed, 486 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e38214c98a6..d3ac37a3e0f 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -12,7 +12,8 @@ import os import random import time import traceback -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence +from contextvars import ContextVar from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload @@ -269,6 +270,80 @@ def _spend_update_tx(prisma_client: PrismaClient) -> _SpendTransactionManager: return tx +_daily_spend_commit_started: Final[ContextVar[asyncio.Event | None]] = ContextVar( + "_daily_spend_commit_started", default=None +) + + +def _mark_daily_spend_commit_started() -> None: + started: Final = _daily_spend_commit_started.get() + if started is not None: + started.set() + + +def _mark_daily_spend_commit_finished() -> None: + started: Final = _daily_spend_commit_started.get() + if started is not None: + started.clear() + + +def _start_daily_spend_commit( + commit_started: asyncio.Event, commit: Callable[[], Coroutine[object, object, None]] +) -> "asyncio.Task[None]": + token: Final = _daily_spend_commit_started.set(commit_started) + try: + return asyncio.ensure_future(commit()) + finally: + _daily_spend_commit_started.reset(token) + + +def _track_interrupted_commit(commits: set[asyncio.Task[None]], settle: Coroutine[object, object, None]) -> None: + task: Final = asyncio.ensure_future(settle) + commits.add(task) + task.add_done_callback(commits.discard) + + +async def _settle_interrupted_commits(commits: set[asyncio.Task[None]]) -> None: + while commits: + await asyncio.wait(tuple(commits)) + + +async def _restore_tag_spend_the_commit_left_behind( + commit_task: "asyncio.Task[None]", + redis_update_buffer: RedisUpdateBuffer, + transactions: dict[str, DailyTagSpendTransaction], +) -> None: + await asyncio.wait({commit_task}) + if commit_task.cancelled() or commit_task.exception() is None: + return + await redis_update_buffer.restore_transactions_to_redis( + daily_tag_spend_update_transactions=transactions, + ) + + +async def _requeue_daily_spend_the_commit_left_behind( + commit_task: "asyncio.Task[None]", + queue: DailySpendUpdateQueue, + entity_type: str, + transactions: dict[str, BaseDailySpendTransaction], +) -> None: + await asyncio.wait({commit_task}) + if commit_task.cancelled() or not transactions: + return + failure: Final = commit_task.exception() + if failure is None: + return + spend_log_error( + "Spend tracking - daily %s spend commit interrupted by shutdown failed. Re-queued %d rows for the " + "shutdown flush. Error: %s", + entity_type, + len(transactions), + str(failure), + exc=failure, + ) + await queue.add_update(transactions) + + # The per-team advisory lock the team endpoints hold while changing a roster (TEAM_ADVISORY_LOCK_SQL), # so the roster check below cannot interleave with their writes. A row lock would deadlock with the # access-group endpoints, which lock a team row after an access-group lock. @@ -391,6 +466,9 @@ class DBSpendUpdateWriter: self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() self.window_spend_update_queue = WindowSpendUpdateQueue() + self.interrupted_tag_commits: set[asyncio.Task[None]] = ( + set() + ) # mutable-ok: same registry as DailySpendUpdateQueue.interrupted_commits async def update_database( # LiteLLM management object fields @@ -1606,17 +1684,24 @@ class DBSpendUpdateWriter: proxy_logging_obj: ProxyLogging, ) -> None: transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions() - commit_task: Final = asyncio.ensure_future( - commit( + commit_started: Final = asyncio.Event() + commit_task: Final = _start_daily_spend_commit( + commit_started, + lambda: commit( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions), - ) + ), ) try: await asyncio.shield(commit_task) except asyncio.CancelledError: + if commit_started.is_set(): + queue.track_interrupted_commit( + _requeue_daily_spend_the_commit_left_behind(commit_task, queue, entity_type, transactions) + ) + raise commit_task.cancel() if transactions: await queue.add_update(transactions) @@ -1841,23 +1926,36 @@ class DBSpendUpdateWriter: The drain is destructive, so a failed commit must push the transactions back for the next tick or their spend is lost permanently. """ + await _settle_interrupted_commits(self.interrupted_tag_commits) daily_tag_spend_update_transactions: Final = ( await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() ) if not daily_tag_spend_update_transactions: return - commit_task: Final = asyncio.ensure_future( - DBSpendUpdateWriter.update_daily_tag_spend( + commit_started: Final = asyncio.Event() + commit_task: Final = _start_daily_spend_commit( + commit_started, + lambda: DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, - ) + ), ) try: await asyncio.shield(commit_task) except BaseException: # noqa: BLE001 # a cancel must restore the drained rows before its rollback returns + if commit_started.is_set(): + _track_interrupted_commit( + self.interrupted_tag_commits, + _restore_tag_spend_the_commit_left_behind( + commit_task, + self.redis_update_buffer, + daily_tag_spend_update_transactions, + ), + ) + raise commit_task.cancel() await self.redis_update_buffer.restore_transactions_to_redis( daily_tag_spend_update_transactions=daily_tag_spend_update_transactions, @@ -2382,6 +2480,8 @@ class DBSpendUpdateWriter: sql, params = build_bulk_upsert(table=table, batch=merged_batch) async with _spend_update_tx(prisma_client) as transaction: await transaction.execute_raw(sql, *params) + _mark_daily_spend_commit_started() + _mark_daily_spend_commit_finished() except Exception as batch_error: if _spend_commit_failure_is_requeue_safe(batch_error): spend_log_error( diff --git a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py index c6381cd070b..f911d5a6767 100644 --- a/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/daily_spend_update_queue.py @@ -1,4 +1,5 @@ import asyncio +from collections.abc import Coroutine from copy import deepcopy from typing import Final @@ -57,6 +58,18 @@ class DailySpendUpdateQueue(BaseUpdateQueue): self.update_queue: asyncio.Queue[dict[str, BaseDailySpendTransaction]] = asyncio.Queue( maxsize=LITELLM_ASYNCIO_QUEUE_MAXSIZE ) + self.interrupted_commits: set[asyncio.Task[None]] = ( + set() + ) # mutable-ok: registry of in-flight commit outcomes, entries leave via their done callback + + def track_interrupted_commit(self, settle: Coroutine[object, object, None]) -> None: + task: Final = asyncio.ensure_future(settle) + self.interrupted_commits.add(task) + task.add_done_callback(self.interrupted_commits.discard) + + async def settle_interrupted_commits(self) -> None: + while self.interrupted_commits: + await asyncio.wait(tuple(self.interrupted_commits)) async def add_update(self, update: dict[str, BaseDailySpendTransaction]): """Enqueue an update.""" @@ -81,6 +94,7 @@ class DailySpendUpdateQueue(BaseUpdateQueue): self, ) -> dict[str, BaseDailySpendTransaction]: """Get all updates from the queue and return all updates aggregated by daily_transaction_key. Works for both user and team spend updates.""" + await self.settle_interrupted_commits() updates: Final = await self.flush_all_updates_from_in_memory_queue() if len(updates) > 0: verbose_proxy_logger.info( diff --git a/tests/integration/spend/test_shutdown_flush.py b/tests/integration/spend/test_shutdown_flush.py index d27275f5213..b744c1f4b5f 100644 --- a/tests/integration/spend/test_shutdown_flush.py +++ b/tests/integration/spend/test_shutdown_flush.py @@ -4,6 +4,7 @@ import signal import threading import uuid from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -13,16 +14,18 @@ import httpx import psycopg import pytest import yaml - from integration._support.client import Gateway, delete_key_if_present, eventually, string_value from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from psycopg import sql REQUESTS_WHILE_BLOCKED: Final = 6 CANCEL_LOG_LINE: Final = "in-flight scheduled job(s) for shutdown" BATCH_DRAINED_LOG_LINE: Final = f"flushed {REQUESTS_WHILE_BLOCKED} daily spend update items from in-memory queue" MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new" +COMMIT_DELAY_SECONDS: Final = 15 +BURST_REQUESTS: Final = 30 def _api_requests(table: str, column: str, identity: str) -> int: @@ -44,6 +47,50 @@ def _waiting_on(table: str) -> int: return waiting +def _committing_daily_user_spend() -> int: + rows: Final = read_rows( + "SELECT count(*)::int AS committing FROM pg_stat_activity " + "WHERE query='COMMIT' AND state='active' AND wait_event='PgSleep' AND pid IN " + "(SELECT pid FROM pg_locks WHERE relation = %s::regclass AND mode='RowExclusiveLock')", + ('"LiteLLM_DailyUserSpend"',), + ) + committing: Final = rows[0]["committing"] + assert isinstance(committing, int) + return committing + + +def _install_slow_commit(user_id: str, fails_once: bool) -> str: + suffix: Final = f"slow_commit_{uuid.uuid4().hex}" + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute(sql.SQL("CREATE SEQUENCE {}").format(sql.Identifier(suffix))) + connection.execute( + sql.SQL( + "CREATE FUNCTION {}() RETURNS trigger LANGUAGE plpgsql AS $slow$ " + "BEGIN PERFORM pg_sleep({}); " + "IF {} AND nextval({}) = 1 THEN RAISE EXCEPTION 'integration: first COMMIT fails'; END IF; " + "RETURN NULL; END $slow$" + ).format( + sql.Identifier(suffix), sql.Literal(COMMIT_DELAY_SECONDS), sql.Literal(fails_once), sql.Literal(suffix) + ) + ) + connection.execute( + sql.SQL( + 'CREATE CONSTRAINT TRIGGER {} AFTER INSERT OR UPDATE ON "LiteLLM_DailyUserSpend" ' + "DEFERRABLE INITIALLY DEFERRED FOR EACH ROW WHEN (NEW.user_id = {}) EXECUTE FUNCTION {}()" + ).format(sql.Identifier(suffix), sql.Literal(user_id), sql.Identifier(suffix)) + ) + return suffix + + +def _drop_slow_commit(suffix: str) -> None: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as connection: + connection.execute( + sql.SQL('DROP TRIGGER IF EXISTS {} ON "LiteLLM_DailyUserSpend"').format(sql.Identifier(suffix)) + ) + connection.execute(sql.SQL("DROP FUNCTION IF EXISTS {}()").format(sql.Identifier(suffix))) + connection.execute(sql.SQL("DROP SEQUENCE IF EXISTS {}").format(sql.Identifier(suffix))) + + def _provider(request: Request) -> Reply: if request.method != "POST": return Reply(status=404, body=b'{"error":"not scripted"}') @@ -77,6 +124,17 @@ class _Shutdown: def daily_user_requests(self) -> int: return _api_requests("LiteLLM_DailyUserSpend", "user_id", self.owner) + def spend_logs(self) -> int: + rows: Final = read_rows('SELECT count(*)::int AS total FROM "LiteLLM_SpendLogs" WHERE "user"=%s', (self.owner,)) + total: Final = rows[0]["total"] + assert isinstance(total, int) + return total + + def burst(self, requests: int) -> None: + with ThreadPoolExecutor(max_workers=8) as pool: + for outcome in pool.map(lambda _: self.chat(), range(requests)): + assert outcome is None + def logged(self, line: str, times: int = 1) -> bool: return self.owned.log.read_text(errors="replace").count(line) >= times @@ -121,7 +179,15 @@ def _config_with_pool_limit(tmp_path: Path, pool_limit: int) -> Path: @contextmanager -def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int) -> Iterator[_Shutdown]: +def _proxy_with_one_seeded_row( + gateway: Gateway, + tmp_path: Path, + pool_limit: int, + cancel_timeout_seconds: int = 5, + settle_seconds: int = 0, + requests: int = REQUESTS_WHILE_BLOCKED, + workers: int = 1, +) -> Iterator[_Shutdown]: owner: Final = f"integration-owner-{uuid.uuid4().hex}" with gateway.scenario() as scenario, wire_server(_provider) as wire: model: Final = scenario.model(api_base=wire.url + "/v1", num_retries=0) @@ -133,9 +199,10 @@ def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int "LITELLM_LOG": "DEBUG", "GRACEFUL_SHUTDOWN_TIMEOUT": "1", "SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS": "1", - "SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": "5", + "SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS": str(cancel_timeout_seconds), }, config=_config_with_pool_limit(tmp_path, pool_limit), + workers=workers, ) as owned: key: Final = string_value( owned.gateway.post("/key/generate", {"user_id": owner, "team_id": team, "models": [model]})["key"] @@ -145,8 +212,18 @@ def _proxy_with_one_seeded_row(gateway: Gateway, tmp_path: Path, pool_limit: int shutdown.chat() eventually(shutdown.daily_user_requests, lambda total: total == 1, seconds=60) yield shutdown - assert _api_requests("LiteLLM_DailyUserSpend", "user_id", owner) == 1 + REQUESTS_WHILE_BLOCKED - assert _api_requests("LiteLLM_DailyTeamSpend", "team_id", team) == 1 + REQUESTS_WHILE_BLOCKED + written: Final = 1 + requests + if settle_seconds: + eventually( + lambda: ( + _api_requests("LiteLLM_DailyUserSpend", "user_id", owner), + _api_requests("LiteLLM_DailyTeamSpend", "team_id", team), + ), + lambda totals: totals == (written, written), + seconds=settle_seconds, + ) + assert _api_requests("LiteLLM_DailyUserSpend", "user_id", owner) == written + assert _api_requests("LiteLLM_DailyTeamSpend", "team_id", team) == written @pytest.mark.covers("quota_management.spend_tracking.shutdown_cancel_keeps_in_flight_daily_batch") @@ -187,3 +264,48 @@ def test_daily_spend_batch_cancelled_while_waiting_for_a_row_lock_is_written_exa lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _waiting_on("LiteLLM_DailyUserSpend") == 1, holder.rollback, ) + + +@pytest.mark.parametrize( + ("cancel_timeout_seconds", "commit_fails_once"), + [ + pytest.param(60, False, id="cancel_budget_outlives_commit"), + pytest.param(5, False, id="commit_outlives_cancel_budget"), + pytest.param(60, True, id="commit_fails_within_cancel_budget"), + pytest.param(5, True, id="commit_fails_after_cancel_budget"), + ], +) +def test_daily_spend_batch_cancelled_while_postgres_is_committing_it_is_written_exactly_once( + gateway: Gateway, tmp_path: Path, cancel_timeout_seconds: int, commit_fails_once: bool +) -> None: + with ( + _proxy_with_one_seeded_row( + gateway, tmp_path, pool_limit=10, cancel_timeout_seconds=cancel_timeout_seconds, settle_seconds=90 + ) as shutdown, + psycopg.connect(os.environ["DATABASE_URL"]) as memberships, + ): + suffix: Final = _install_slow_commit(shutdown.owner, fails_once=commit_fails_once) + try: + shutdown.chat_while_spend_update_is_blocked(memberships, "LiteLLM_TeamMembership") + memberships.rollback() + shutdown.terminate_once( + lambda: shutdown.logged(BATCH_DRAINED_LOG_LINE) and _committing_daily_user_spend() == 1, + lambda: None, + ) + finally: + _drop_slow_commit(suffix) + + +def test_daily_spend_burst_across_two_workers_survives_shutdown_during_commit_exactly_once( + gateway: Gateway, tmp_path: Path +) -> None: + with _proxy_with_one_seeded_row( + gateway, tmp_path, pool_limit=10, settle_seconds=120, requests=BURST_REQUESTS, workers=2 + ) as shutdown: + suffix: Final = _install_slow_commit(shutdown.owner, fails_once=False) + try: + shutdown.burst(BURST_REQUESTS) + eventually(shutdown.spend_logs, lambda total: total == 1 + BURST_REQUESTS, seconds=60) + shutdown.terminate_once(lambda: _committing_daily_user_spend() >= 1, lambda: None) + finally: + _drop_slow_commit(suffix) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index daa07c8224a..bf3b9aed234 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -4501,3 +4501,241 @@ async def test_tag_batch_drained_from_redis_and_cancelled_mid_flight_is_restored await asyncio.wait_for(db.rolled_back.wait(), timeout=5) assert db.transaction_outcomes == ["rollback"] assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == [] + + +class _CommittingDailySpendFakeDB(_DailySpendFakeDB): + """Runs the daily upsert at once but holds the COMMIT until released. A COMMIT that has left + the client lands on the server whether or not the client keeps waiting for the reply.""" + + def __init__(self) -> None: + super().__init__(failing_table=None) + self.committing = asyncio.Event() + self.commit_release = asyncio.Event() + self.transaction_outcomes: list[str] = [] + + @asynccontextmanager + async def _tx(self) -> AsyncIterator["_CommittingDailySpendFakeDB"]: + try: + yield self + except BaseException: + self.transaction_outcomes.append("rollback") + raise + self.committing.set() + try: + await self.commit_release.wait() + finally: + self.transaction_outcomes.append("commit") + + +@pytest.mark.parametrize(("queue_name", "entity_type", "entity_id_field", "table"), _DAILY_SPEND_ENTITIES) +@pytest.mark.asyncio +async def test_cancel_that_lands_while_the_daily_batch_is_committing_waits_for_the_commit_and_does_not_requeue_it( + queue_name: str, entity_type: str, entity_id_field: str, table: str +): + """Shutdown cancels the tick after the COMMIT has left for Postgres. The server finishes that + commit whatever the client does, so putting the batch back on the queue makes the final flush + write the same spend a second time. The tick has to wait for the commit's outcome instead.""" + db_writer = DBSpendUpdateWriter() + queue = _DAILY_SPEND_QUEUES[queue_name](db_writer) + await queue.add_update({"key-a": _daily_entity_txn(entity_id_field)}) + await queue.add_update({"key-a": _daily_entity_txn(entity_id_field)}) + db = _CommittingDailySpendFakeDB() + + def flush(prisma_db: _DailySpendFakeDB): + return db_writer._flush_daily_spend_queue( + queue=queue, + entity_type=entity_type, + commit=_DAILY_SPEND_COMMITS[entity_type], + n_retry_times=0, + prisma_client=_WindowSpendFakePrisma(prisma_db), + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(flush(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + finished, _ = await asyncio.wait({tick}, timeout=0.2) + assert finished == {tick}, "the cancelled tick must hand the in-flight commit's outcome to the next flush" + with pytest.raises(asyncio.CancelledError): + tick.result() + assert len(queue.interrupted_commits) == 1 + + db.commit_release.set() + await queue.settle_interrupted_commits() + + assert db.transaction_outcomes == ["commit"] + (upsert,) = _daily_upserts(db, table) + assert _row_values(upsert, "api_requests") == [2] + assert queue.update_queue.empty(), "a batch whose COMMIT already left for the server must not be requeued" + + final_db = _DailySpendFakeDB(failing_table=None) + await flush(final_db) + assert _daily_upserts(final_db, table) == [], "the final flush must not write the committed batch again" + + +@pytest.mark.asyncio +async def test_tag_batch_drained_from_redis_and_cancelled_while_committing_is_not_restored(): + """Same in-flight COMMIT as the in-memory path, but the drained rows live in Redis. Restoring + them after the server committed writes the tag spend twice on the next tick.""" + db_writer = DBSpendUpdateWriter() + drained = {"key-a": cast(DailyTagSpendTransaction, _daily_entity_txn("tag"))} + redis_buffer = _DrainedTagRedisBuffer(drained) + db_writer.redis_update_buffer = cast(RedisUpdateBuffer, redis_buffer) + db = _CommittingDailySpendFakeDB() + + tick = asyncio.ensure_future( + db_writer._drain_and_commit_daily_tag_spend_from_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + ) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + finished, _ = await asyncio.wait({tick}, timeout=0.2) + assert finished == {tick}, "the cancelled drain must hand the in-flight commit's outcome to the next drain" + with pytest.raises(asyncio.CancelledError): + tick.result() + assert len(db_writer.interrupted_tag_commits) == 1 + + db.commit_release.set() + (settle,) = tuple(db_writer.interrupted_tag_commits) + await settle + + assert db.transaction_outcomes == ["commit"] + assert redis_buffer.restored == [], ( + "a tag batch whose COMMIT already left for the server must not be restored to Redis" + ) + + +class _CommitFailingDailySpendFakeDB(_CommittingDailySpendFakeDB): + """COMMIT leaves for the server but the reply comes back as a failure.""" + + @asynccontextmanager + async def _tx(self) -> AsyncIterator["_CommitFailingDailySpendFakeDB"]: + yield self + self.committing.set() + await self.commit_release.wait() + self.transaction_outcomes.append("commit_failed") + raise Exception("connection reset") + + +@pytest.mark.asyncio +async def test_cancel_while_committing_requeues_the_batch_when_the_commit_itself_fails(): + """Waiting for the in-flight commit's outcome must not swallow a real commit failure: + the batch still goes back on the queue and the next flush writes it once.""" + db_writer = DBSpendUpdateWriter() + queue = db_writer.daily_spend_update_queue + await queue.add_update({"key-a": _daily_txn()}) + await queue.add_update({"key-a": _daily_txn()}) + db = _CommitFailingDailySpendFakeDB() + + def flush(prisma_db: _DailySpendFakeDB): + return db_writer._flush_daily_spend_queue( + queue=queue, + entity_type="user", + commit=DBSpendUpdateWriter.update_daily_user_spend, + n_retry_times=0, + prisma_client=_WindowSpendFakePrisma(prisma_db), + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(flush(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + finished, _ = await asyncio.wait({tick}, timeout=0.2) + assert finished == {tick}, "the cancelled tick must not eat the shutdown budget waiting on the commit" + with pytest.raises(asyncio.CancelledError): + tick.result() + + db.commit_release.set() + await queue.settle_interrupted_commits() + + assert db.transaction_outcomes == ["commit_failed"] + assert not queue.update_queue.empty(), "a batch whose COMMIT came back failed must be requeued" + + final_db = _DailySpendFakeDB(failing_table=None) + await flush(final_db) + (upsert,) = _daily_upserts(final_db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "api_requests") == [2] + assert queue.update_queue.empty() + + +@pytest.mark.asyncio +async def test_shutdown_flush_that_lands_before_the_interrupted_commit_resolves_still_writes_a_failed_batch_once(): + """The cancelled tick returns right away, so a COMMIT can still be in flight when the + shutdown flush runs. If that commit later fails, the flush must first settle it, pick the + requeued rows back up, and write them exactly once instead of losing them.""" + db_writer = DBSpendUpdateWriter() + queue = db_writer.daily_spend_update_queue + await queue.add_update({"key-a": _daily_txn()}) + await queue.add_update({"key-a": _daily_txn()}) + db = _CommitFailingDailySpendFakeDB() + + def flush(prisma_db: _DailySpendFakeDB): + return db_writer._flush_daily_spend_queue( + queue=queue, + entity_type="user", + commit=DBSpendUpdateWriter.update_daily_user_spend, + n_retry_times=0, + prisma_client=_WindowSpendFakePrisma(prisma_db), + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(flush(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(tick, timeout=5) + assert db.transaction_outcomes == [], "the COMMIT is still on the wire when the shutdown flush starts" + + final_db = _DailySpendFakeDB(failing_table=None) + shutdown_flush = asyncio.ensure_future(flush(final_db)) + finished, _ = await asyncio.wait({shutdown_flush}, timeout=0.2) + assert finished == set(), "the shutdown flush must wait for the interrupted commit's outcome" + assert _daily_upserts(final_db, "LiteLLM_DailyUserSpend") == [] + + db.commit_release.set() + await asyncio.wait_for(shutdown_flush, timeout=5) + + (upsert,) = _daily_upserts(final_db, "LiteLLM_DailyUserSpend") + assert _row_values(upsert, "api_requests") == [2] + assert queue.update_queue.empty() + + +@pytest.mark.asyncio +async def test_shutdown_drain_that_lands_before_the_interrupted_tag_commit_resolves_restores_a_failed_batch(): + """Same ordering for the Redis tag path: the shutdown drain must settle the interrupted + commit before the destructive drain, or a commit that fails late is never restored.""" + db_writer = DBSpendUpdateWriter() + drained = {"key-a": cast(DailyTagSpendTransaction, _daily_entity_txn("tag"))} + redis_buffer = _DrainedTagRedisBuffer(drained) + db_writer.redis_update_buffer = cast(RedisUpdateBuffer, redis_buffer) + db = _CommitFailingDailySpendFakeDB() + + def drain(prisma_db: _DailySpendFakeDB): + return db_writer._drain_and_commit_daily_tag_spend_from_redis( + prisma_client=_WindowSpendFakePrisma(prisma_db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + tick = asyncio.ensure_future(drain(db)) + await asyncio.wait_for(db.committing.wait(), timeout=5) + tick.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(tick, timeout=5) + assert db.transaction_outcomes == [] + + final_db = _DailySpendFakeDB(failing_table=None) + shutdown_drain = asyncio.ensure_future(drain(final_db)) + finished, _ = await asyncio.wait({shutdown_drain}, timeout=0.2) + assert finished == set(), "the shutdown drain must wait for the interrupted commit's outcome" + assert _daily_upserts(final_db, "LiteLLM_DailyTagSpend") == [] + + db.commit_release.set() + await asyncio.wait_for(shutdown_drain, timeout=5) + + assert redis_buffer.restored == [drained], "a tag batch whose COMMIT came back failed must be restored to Redis" + (upsert,) = _daily_upserts(final_db, "LiteLLM_DailyTagSpend") + assert _row_values(upsert, "api_requests") == [1]