From 71a68d1679f4eaf55baa5d304d637dcef3f07079 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 30 Oct 2026 15:32:56 -0700 Subject: [PATCH] fix(proxy): log database cleanup errors and bound shutdown waits Database.aclose now logs a failure in either cleanup step instead of raising out of the lifespan exit, which matches how main handled stop errors and keeps setup_prisma_client from raising out of cleanup after a failed connect Shutdown waits for an in-flight reconnect for a fixed five seconds instead of the reconnect timeout plus grace, and joins engine wait threads with a fixed limit instead of indefinitely. The threads are daemon threads, so a thread still alive after the limit cannot keep the process from exiting Co-Authored-By: Claude Fable 5.1 --- litellm/proxy/utils.py | 24 ++++++------ litellm/repositories/database.py | 6 ++- .../unit/proxy/proxy_server/test_lifecycle.py | 7 +++- .../test_prisma_client_lifecycle.py | 37 +++++++++++++++++-- tests/unit/repositories/test_database.py | 22 ++++++++++- 5 files changed, 76 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 39396c263c4..653087523bb 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5187,7 +5187,8 @@ async def prefetch_config_params(prisma_client: "PrismaClient | None", param_nam _WRITER_WRITABILITY_PROBE_SQL: Final = "SELECT current_setting('transaction_read_only') AS transaction_read_only" _WRITER_WRITABILITY_PROBE_ROWS: Final = TypeAdapter(list[dict[str, object]]) _READ_ONLY_RECREATE_BACKOFF_CAP_SECONDS: Final = 600 -_SHUTDOWN_RECONNECT_GRACE_SECONDS: Final = 5.0 +_SHUTDOWN_RECONNECT_WAIT_SECONDS: Final = 5.0 +_SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS: Final = 5.0 class _ForcedRecreateDeclined(Exception): @@ -5272,6 +5273,7 @@ class PrismaClient: database_url: str, proxy_logging_obj: ProxyLogging, http_client: "HttpConfig | None" = None, + engine_exit_wait_seconds: float = _SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS, ): ## init logging object self.baseline_accounting_transactions: list[BaselineAccountingRecord] = [] @@ -5374,6 +5376,7 @@ class PrismaClient: self.db = writer_wrapper # Client to connect to Prisma db self._db_reconnect_lock = asyncio.Lock() self._closing: bool = False + self._engine_exit_wait_seconds: Final = engine_exit_wait_seconds self._engine_tasks: frozenset[asyncio.Task[object]] = frozenset() self._db_health_watchdog_task: asyncio.Task | None = None self._view_setup_task: asyncio.Task[_ViewSetupOutcome] | None = None @@ -6626,7 +6629,12 @@ class PrismaClient: try: await self.db.disconnect() if self._closing: - await asyncio.gather(*(asyncio.to_thread(thread.join) for thread in self._engine_wait_threads)) + await asyncio.gather( + *( + asyncio.to_thread(thread.join, self._engine_exit_wait_seconds) + for thread in self._engine_wait_threads + ) + ) self._engine_wait_threads = () except Exception as e: import traceback @@ -6667,20 +6675,12 @@ class PrismaClient: await self._wait_for_inflight_reconnect() async def _wait_for_inflight_reconnect(self) -> None: - # Request-driven reconnects are not engine tasks. Wait for their bounded - # reconnect cycle to release the lock before disconnecting, otherwise a - # late reconnect can reopen the client after shutdown. A reconnect stuck - # past its own timeout must not block shutdown, so the wait is bounded too. waiter: Final = asyncio.ensure_future(self._db_reconnect_lock.acquire()) waiter.add_done_callback(lambda acquired: acquired.cancelled() or self._db_reconnect_lock.release()) - done, _pending = await asyncio.wait( - {waiter}, timeout=self._db_watchdog_reconnect_timeout_seconds + _SHUTDOWN_RECONNECT_GRACE_SECONDS - ) + done, _pending = await asyncio.wait({waiter}, timeout=_SHUTDOWN_RECONNECT_WAIT_SECONDS) if waiter not in done: waiter.cancel() - verbose_proxy_logger.warning( - "Shutdown stopped waiting for an in-flight DB reconnect that outlived its timeout" - ) + verbose_proxy_logger.warning("Shutdown stopped waiting for an in-flight DB reconnect") async def stop_background_tasks(self) -> None: self._closing = True diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py index c7291ad8adc..0ffc5762e80 100644 --- a/litellm/repositories/database.py +++ b/litellm/repositories/database.py @@ -18,8 +18,12 @@ class Database: verbose_proxy_logger.debug("Disconnecting from Prisma") try: await self.prisma.stop_background_tasks() - finally: + except Exception: + verbose_proxy_logger.exception("Stopping database background tasks failed") + try: await self.prisma.disconnect() + except Exception: + verbose_proxy_logger.exception("Database disconnect failed") @asynccontextmanager diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index f25dfed196d..404c31f59e1 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -326,7 +326,7 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t @pytest.mark.asyncio -async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monkeypatch): +async def test_database_lifespan_logs_disconnect_error_after_shutdown_instead_of_raising(monkeypatch, caplog): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=RuntimeError("db gone")) fake_prisma.stop_background_tasks = AsyncMock() @@ -341,10 +341,13 @@ async def test_database_lifespan_propagates_disconnect_error_after_shutdown(monk monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - with pytest.raises(RuntimeError, match="db gone"): + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): async with aclosing(database_runtime.Database(prisma=fake_prisma)): await proxy_shutdown_event() + fake_prisma.disconnect.assert_awaited_once() + assert "db gone" in caplog.text + # --------------------------------------------------------------------------- # _flush_spend_logs_queue_on_shutdown diff --git a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py index 764cb6469f9..adecc346d26 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py +++ b/tests/unit/proxy/utils/prisma_and_spend/test_prisma_client_lifecycle.py @@ -14,12 +14,13 @@ from __future__ import annotations import asyncio import logging +import threading from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.proxy.utils import PrismaClient +from litellm.proxy.utils import _SHUTDOWN_RECONNECT_WAIT_SECONDS, PrismaClient from litellm.repositories.database import Database @@ -68,7 +69,10 @@ async def test_shutdown_waits_for_hung_request_reconnect_to_time_out( @pytest.mark.asyncio -async def test_shutdown_stops_waiting_for_a_reconnect_stuck_past_its_timeout(prisma_client: PrismaClient) -> None: +async def test_shutdown_stops_waiting_for_a_held_reconnect_lock_after_the_fixed_limit( + prisma_client: PrismaClient, +) -> None: + assert _SHUTDOWN_RECONNECT_WAIT_SECONDS < prisma_client._db_watchdog_reconnect_timeout_seconds await prisma_client._db_reconnect_lock.acquire() with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock: closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose()) @@ -77,7 +81,7 @@ async def test_shutdown_stops_waiting_for_a_reconnect_stuck_past_its_timeout(pri assert not closing.done() prisma_client.db.disconnect.assert_not_awaited() - clock.return_value = 1000.0 + prisma_client._db_watchdog_reconnect_timeout_seconds + 60 + clock.return_value = 1000.0 + _SHUTDOWN_RECONNECT_WAIT_SECONDS + 1 await closing finally: closing.cancel() @@ -140,6 +144,33 @@ async def test_shutdown_waits_for_spawned_engine_work_and_prevents_rearming(pris assert not prisma_client._engine_tasks +@pytest.mark.asyncio +async def test_shutdown_disconnect_returns_when_an_engine_wait_thread_never_exits( + patched_prisma_import: MagicMock, mock_prisma_client: MagicMock, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False) + monkeypatch.delenv("IAM_TOKEN_DB_AUTH", raising=False) + client: Final = PrismaClient( + database_url="postgresql://test:test@localhost:5432/test", + proxy_logging_obj=MagicMock(failure_handler=AsyncMock()), + engine_exit_wait_seconds=0.05, + ) + client.db = mock_prisma_client.db + release: Final = threading.Event() + stuck: Final = threading.Thread(target=release.wait, daemon=True) + stuck.start() + client._engine_wait_threads = (stuck,) + client._closing = True + try: + await asyncio.wait_for(client.disconnect(), timeout=2) + finally: + release.set() + stuck.join(timeout=2) + + client.db.disconnect.assert_awaited_once() + assert client._engine_wait_threads == () + + _PROBE_SQL = "SELECT to_regclass($1) IS NOT NULL AS present" diff --git a/tests/unit/repositories/test_database.py b/tests/unit/repositories/test_database.py index 008eb1124e8..957c4cefbfb 100644 --- a/tests/unit/repositories/test_database.py +++ b/tests/unit/repositories/test_database.py @@ -1,3 +1,4 @@ +import logging from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -7,12 +8,29 @@ from litellm.repositories.database import Database @pytest.mark.asyncio -async def test_database_disconnects_even_when_stopping_background_tasks_fails() -> None: +async def test_database_disconnects_and_logs_when_stopping_background_tasks_fails( + caplog: pytest.LogCaptureFixture, +) -> None: prisma: Final = MagicMock() prisma.stop_background_tasks = AsyncMock(side_effect=RuntimeError("watchdog failed")) prisma.disconnect = AsyncMock() - with pytest.raises(RuntimeError, match="watchdog failed"): + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): await Database(prisma=prisma).aclose() prisma.disconnect.assert_awaited_once() + assert "watchdog failed" in caplog.text + + +@pytest.mark.asyncio +async def test_database_aclose_returns_and_logs_when_disconnect_fails(caplog: pytest.LogCaptureFixture) -> None: + prisma: Final = MagicMock() + prisma.stop_background_tasks = AsyncMock() + prisma.disconnect = AsyncMock(side_effect=RuntimeError("engine hung")) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + await Database(prisma=prisma).aclose() + + prisma.stop_background_tasks.assert_awaited_once() + prisma.disconnect.assert_awaited_once() + assert "engine hung" in caplog.text