From 2d79d5dba70cf0979e1615208dfe2b2572e45526 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 30 Oct 2026 15:57:53 -0700 Subject: [PATCH] refactor(proxy): give database shutdown one lifespan owner (#48506) * refactor(proxy): give database shutdown one lifespan owner Co-Authored-By: Claude Opus 5.5 * 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 --------- Co-authored-by: Claude Opus 5.5 --- litellm/proxy/_types.py | 2 + litellm/proxy/db/routing_prisma_wrapper.py | 10 +- litellm/proxy/proxy_server.py | 122 +++++++------- litellm/proxy/utils.py | 73 +++++++- litellm/repositories/database.py | 36 ++++ .../database/test_proxy_database_lifespan.py | 158 ++++++++++++++++++ .../proxy/db/test_routing_prisma_wrapper.py | 49 ++++++ .../unit/proxy/proxy_server/test_lifecycle.py | 37 ++-- tests/unit/proxy/test_proxy_server.py | 4 +- ...test_proxy_server_endpoints_and_startup.py | 75 ++++++++- .../proxy/utils/prisma_and_spend/conftest.py | 1 + .../test_prisma_client_lifecycle.py | 156 ++++++++++++++++- tests/unit/repositories/test_database.py | 36 ++++ 13 files changed, 664 insertions(+), 95 deletions(-) create mode 100644 litellm/repositories/database.py create mode 100644 tests/integration/database/test_proxy_database_lifespan.py create mode 100644 tests/unit/repositories/test_database.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 73b12399af4..07588053959 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -124,6 +124,7 @@ from .types_utils.utils import get_instance_fn, validate_custom_validate_return_ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.repositories.database import Database from litellm.tracing import TraceReceiver Span = _Span | Any @@ -133,6 +134,7 @@ else: class ProxyLifespanState(TypedDict): tracing_receiver: ReadOnly["TraceReceiver | None"] + database: ReadOnly["Database | None"] class ReconcileOutcome(NamedTuple): diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 52a0be3d49e..50e6aa54da4 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -4,6 +4,7 @@ to a writer client. Used when DATABASE_URL_READ_REPLICA is configured; otherwise PrismaClient uses the writer-only PrismaWrapper directly. """ +import asyncio import os from collections.abc import Callable from datetime import timedelta @@ -233,8 +234,13 @@ class RoutingPrismaWrapper: await self._reader.start_token_refresh_task() async def stop_token_refresh_task(self) -> None: - await self._writer.stop_token_refresh_task() - await self._reader.stop_token_refresh_task() + results: Final = await asyncio.gather( + self._writer.stop_token_refresh_task(), + self._reader.stop_token_refresh_task(), + return_exceptions=True, + ) + if (error := next((result for result in results if isinstance(result, BaseException)), None)) is not None: + raise error async def recreate_prisma_client( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ea6b51b228b..fd242b148ee 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -174,6 +174,7 @@ from litellm.proxy.common_utils.token_count_settings import ( ) from litellm.proxy.management_helpers.auto_router_availability import AutoRouterCatalogEntry, build_auto_router_catalog from litellm.proxy.management_helpers.model_group_custom_fields import model_group_custom_fields +from litellm.repositories import database as database_runtime from litellm.router_strategy.auto_router.auto_router_model_naming import ( GATED_AUTO_ROUTER_CAPABILITIES, STRATEGY_ROUTER_PARAM_FIELDS, @@ -1350,13 +1351,7 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: - # Drain the SGR fold first: it lives in memory, so an un-drained interval - # is lost, and a write attempted after disconnect raises - # ClientNotConnectedError rather than persisting anything. Ordering this - # inside the same guard is what keeps the two from drifting apart. await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) - verbose_proxy_logger.debug("Disconnecting from Prisma") - await prisma_client.disconnect() else: await _drain_audit_tasks_on_shutdown() @@ -1479,23 +1474,38 @@ async def initialize_shared_aiohttp_session() -> "ClientSession | None": _initialize_shared_aiohttp_session: Final = initialize_shared_aiohttp_session -async def _connect_to_count_stored_values() -> SupportsRawQueries: +async def _connect_to_count_stored_values(resources: AsyncExitStack) -> SupportsRawQueries: client: Final = prisma_client or PrismaClient( database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj ) + if prisma_client is None: + _own_database(resources, client) await client.connect() return client.writer_db +def _own_database(resources: AsyncExitStack, client: PrismaClient | None) -> database_runtime.Database | None: + if client is None: + return None + database: Final = database_runtime.Database(prisma=client) + resources.push_async_callback(database.aclose) + return database + + +def _forget_prisma_client() -> None: + global prisma_client # noqa: PLW0603 # routes still read this global until they move to DatabaseDependency + prisma_client = None # rebind-ok: a closed client must not survive into the next lifespan entry + + @asynccontextmanager async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: with GracefulShutdownManager.capture_sigterm(): - async with _proxy_startup_event(app) as state: + async with AsyncExitStack() as resources, _proxy_lifespan(app, resources) as state: yield state @asynccontextmanager -async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]: +async def _proxy_lifespan(app: FastAPI, resources: AsyncExitStack) -> AsyncGenerator[ProxyLifespanState, None]: global \ prisma_client, \ master_key, \ @@ -1514,6 +1524,9 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat shared_aiohttp_session import json + resources.callback(_forget_prisma_client) + existing_database: Final = _own_database(resources, prisma_client) + init_verbose_loggers() running_loop: Final = asyncio.get_running_loop() @@ -1592,21 +1605,24 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat if decoded_worker_config is not None: await initialize_from_worker_config(decoded_worker_config) - enforce_master_key_boot_verdict( - await with_stored_secrets_counted( - master_key_boot_verdict( - master_key=master_key, - environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), - general_settings=general_settings, - config_file_path=user_config_file_path, - override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, - salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, - database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + async with AsyncExitStack() as boot_database: + enforce_master_key_boot_verdict( + await with_stored_secrets_counted( + master_key_boot_verdict( + master_key=master_key, + environment_master_key=os.getenv(MASTER_KEY_ENV_VAR), + general_settings=general_settings, + config_file_path=user_config_file_path, + override_env_is_on=get_secret_bool(WEAK_OR_UNSET_MASTER_KEY_OVERRIDE_ENV_VAR) is True, + salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, + database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None, + ), + count_values_encrypted_with=partial( + count_values_encrypted_with_or_none, partial(_connect_to_count_stored_values, boot_database) + ), ), - count_values_encrypted_with=partial(count_values_encrypted_with_or_none, _connect_to_count_stored_values), - ), - announce=announce_on_stderr_at_exit, - ) + announce=announce_on_stderr_at_exit, + ) # check if DATABASE_URL in environment - load from there if prisma_client is None: @@ -1619,6 +1635,8 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat ) ) + database: Final = existing_database or _own_database(resources, prisma_client) + await migrate_if_requested( environ=os.environ, master_key=master_key, @@ -1628,15 +1646,16 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat ) if prisma_client is not None: + startup_prisma: Final = prisma_client - async def _run_pw_migration(): + async def _run_pw_migration() -> None: try: - result: Final = await migrate_passwords_to_scrypt_async(prisma_client) + result: Final = await migrate_passwords_to_scrypt_async(startup_prisma) verbose_proxy_logger.info("Password migration: %s", result) except Exception as e: verbose_proxy_logger.warning("Password migration skipped: %s", e) - asyncio.create_task(_run_pw_migration()) + await resources.enter_async_context(database_runtime.background_task(_run_pw_migration())) async def _run_agent_grant_id_migration() -> None: from litellm.proxy.agent_endpoints.agent_registry import ( @@ -1647,7 +1666,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat for attempt in range(3): try: result = await global_agent_registry.migrate_legacy_grant_ids( - table=object_permission_table(prisma_client) + table=object_permission_table(startup_prisma) ) if result.rewritten: verbose_proxy_logger.info( @@ -1667,7 +1686,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat if attempt < 2: await asyncio.sleep(5) - asyncio.create_task(_run_agent_grant_id_migration()) + await resources.enter_async_context(database_runtime.background_task(_run_agent_grant_id_migration())) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1853,7 +1872,7 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat enabled=tracing_enabled, settings=tracing_settings, ) as receiver: - state: Final[ProxyLifespanState] = {"tracing_receiver": receiver} + state: Final[ProxyLifespanState] = {"tracing_receiver": receiver, "database": database} from litellm.proxy.admin_mcp import admin_mcp_lifespan telemetry_settings: Final = load_telemetry_settings() @@ -1907,30 +1926,6 @@ async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanStat except Exception as e: verbose_proxy_logger.error("Error closing shared aiohttp session: %s", e) - # Shutdown event - stop RDS IAM token refresh background task - if ( - prisma_client is not None - and hasattr(prisma_client, "db") - and hasattr(prisma_client.db, "stop_token_refresh_task") - ): - try: - await prisma_client.db.stop_token_refresh_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping token refresh task: %s", e) - - # Shutdown event - stop Prisma DB health watchdog task - if prisma_client is not None and hasattr(prisma_client, "stop_db_health_watchdog_task"): - try: - await prisma_client.stop_db_health_watchdog_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping DB health watchdog task: %s", e) - - if prisma_client is not None and hasattr(prisma_client, "stop_view_setup_task"): - try: - await prisma_client.stop_view_setup_task() - except Exception as e: - verbose_proxy_logger.error("Error stopping the spend view setup task: %s", e) - await _drain_spend_event_producer_on_shutdown() if telemetry_start is not None: telemetry_start.cancel() @@ -12124,17 +12119,20 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, user_api_key_cache: UserApiKeyCache, ) -> PrismaClient | None: - """ - - Sets up prisma client - - Adds necessary views to proxy - """ + if database_url is None: + return None + client: Final = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) + async with AsyncExitStack() as setup_resources: + setup_resources.push_async_callback(database_runtime.Database(prisma=client).aclose) + connected: Final = await cls._initialize_prisma_client(client) + if connected is not None: + setup_resources.pop_all() + return connected + + @staticmethod + async def _initialize_prisma_client(prisma_client: PrismaClient) -> PrismaClient | None: connected_client: PrismaClient | None = None try: - if database_url is None: - return None - - prisma_client = PrismaClient(database_url=database_url, proxy_logging_obj=proxy_logging_obj) - try: await prisma_client.connect() except Exception as e: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0b0698faa9a..efbfbc94f47 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -5197,6 +5197,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_WAIT_SECONDS: Final = 5.0 +_SHUTDOWN_ENGINE_EXIT_WAIT_SECONDS: Final = 5.0 class _ForcedRecreateDeclined(Exception): @@ -5281,6 +5283,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] = [] @@ -5382,6 +5385,9 @@ class PrismaClient: else: 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 self._db_last_reconnect_attempt_ts: float = 0.0 @@ -5423,6 +5429,7 @@ class PrismaClient: self._watching_engine: bool = False self._engine_confirmed_dead: bool = False self._engine_wait_thread: threading.Thread | None = None + self._engine_wait_threads: tuple[threading.Thread, ...] = () verbose_proxy_logger.debug("Success - Created Prisma Client") @property @@ -6631,6 +6638,14 @@ class PrismaClient: start_time: Final = time.time() try: await self.db.disconnect() + if self._closing: + 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 @@ -6649,6 +6664,43 @@ class PrismaClient: ) raise e + def _spawn_engine_task(self, work: Coroutine[object, object, object]) -> None: + if self._closing: + work.close() + return + task: Final = asyncio.create_task(work) + self._engine_tasks = self._engine_tasks | {task} + task.add_done_callback(self._engine_task_finished) + + def _engine_task_finished(self, task: asyncio.Task[object]) -> None: + self._engine_tasks = self._engine_tasks - {task} + if not task.cancelled() and (error := task.exception()) is not None: + verbose_proxy_logger.error("Prisma engine background task failed: %s", error) + + async def _stop_engine_tasks(self) -> None: + tasks: Final = tuple(self._engine_tasks) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + await self._wait_for_inflight_reconnect() + + async def _wait_for_inflight_reconnect(self) -> None: + 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=_SHUTDOWN_RECONNECT_WAIT_SECONDS) + if waiter not in done: + waiter.cancel() + verbose_proxy_logger.warning("Shutdown stopped waiting for an in-flight DB reconnect") + + async def stop_background_tasks(self) -> None: + self._closing = True + self._stop_engine_watcher() + async with contextlib.AsyncExitStack() as cleanup: + cleanup.push_async_callback(self._stop_engine_tasks) + cleanup.push_async_callback(self.db.stop_token_refresh_task) + cleanup.push_async_callback(self.stop_view_setup_task) + cleanup.push_async_callback(self.stop_db_health_watchdog_task) + def _get_engine_pid(self) -> int: """Get the PID of the writer's engine subprocess, or 0 if unavailable. @@ -6742,7 +6794,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6762,6 +6814,9 @@ class PrismaClient: name=f"prisma-engine-waitpid-{pid}", ) thread.start() + self._engine_wait_threads = tuple(thread for thread in self._engine_wait_threads if thread.is_alive()) + ( + thread, + ) self._engine_wait_thread = thread return True @@ -6823,7 +6878,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6885,7 +6940,7 @@ class PrismaClient: self._engine_confirmed_dead = True self._reap_all_zombies() self._cleanup_engine_watcher() - asyncio.create_task( + self._spawn_engine_task( self.attempt_db_reconnect( reason="engine_process_death", force=True, @@ -6958,7 +7013,7 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + if self._closing or self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: return pid: Final = self._get_engine_pid() if pid == 0: @@ -6985,7 +7040,7 @@ class PrismaClient: pid, ) self._watching_engine = True - asyncio.create_task(self._poll_engine_proc()) + self._spawn_engine_task(self._poll_engine_proc()) def _stop_engine_watcher(self) -> None: """Stop watching the engine process and clean up all resources.""" @@ -7007,7 +7062,7 @@ class PrismaClient: """ self._engine_confirmed_dead = False self._cleanup_engine_watcher() - asyncio.create_task(self._start_engine_watcher()) + self._spawn_engine_task(self._start_engine_watcher()) async def _run_reconnect_cycle( self, @@ -7214,6 +7269,8 @@ class PrismaClient: force_recreate: bool = False, stale_read_engine: "_StaleReadEngine | None" = None, ) -> bool: + if self._closing: + return False now: Final = time.time() if ( force is False @@ -7404,7 +7461,7 @@ class PrismaClient: - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. """ - if self._db_health_watchdog_enabled is not True: + if self._closing or self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug("Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED") return if self._db_health_watchdog_task is not None: @@ -7438,7 +7495,7 @@ class PrismaClient: verbose_proxy_logger.info("Stopped Prisma DB health watchdog") def start_view_setup_task(self) -> None: - if self._view_setup_task is not None: + if self._closing or self._view_setup_task is not None: return self._view_setup_task = asyncio.create_task(self._run_view_setup()) diff --git a/litellm/repositories/database.py b/litellm/repositories/database.py new file mode 100644 index 00000000000..0ffc5762e80 --- /dev/null +++ b/litellm/repositories/database.py @@ -0,0 +1,36 @@ +import asyncio +from collections.abc import AsyncGenerator, Coroutine +from contextlib import asynccontextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Final + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +@dataclass(frozen=True, slots=True) +class Database: + prisma: "PrismaClient" + + async def aclose(self) -> None: + verbose_proxy_logger.debug("Disconnecting from Prisma") + try: + await self.prisma.stop_background_tasks() + 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 +async def background_task(work: Coroutine[object, object, None]) -> AsyncGenerator[None, None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/integration/database/test_proxy_database_lifespan.py b/tests/integration/database/test_proxy_database_lifespan.py new file mode 100644 index 00000000000..a5bf62b283a --- /dev/null +++ b/tests/integration/database/test_proxy_database_lifespan.py @@ -0,0 +1,158 @@ +import asyncio +import os +import subprocess +import sys +import threading +import uuid +from pathlib import Path +from typing import Final, Literal, TypeAlias + +import psutil +import pytest + +from tests.integration._support.client import eventually +from tests.integration._support.database import read_rows, scratch_database + +pytestmark: Final = pytest.mark.timeout(180) +ROOT: Final = Path(__file__).resolve().parents[3] +PUBLICLY_KNOWN_KEY: Final = "sk-" + "1234" +Mode: TypeAlias = Literal["late", "unsafe"] +CONFIG: Final = ( + "model_list: []\n" + "general_settings:\n" + " disable_prisma_schema_update: true\n" + " disable_spend_logs: true\n" + " disable_model_info_refresh: true\n" +) + + +def _assert_database_idle() -> None: + assert ( + eventually( + lambda: read_rows( + "SELECT pid FROM pg_stat_activity WHERE datname = current_database() AND pid <> pg_backend_pid()", + (), + ), + lambda rows: not rows, + seconds=5, + ) + == [] + ) + assert psutil.Process().children(recursive=True) == [] + assert tuple(thread.name for thread in threading.enumerate() if thread.name.startswith("prisma-engine-")) == () + + +async def _assert_lifespan_closed() -> None: + from litellm.proxy import proxy_server + + await asyncio.sleep(0) + assert proxy_server.prisma_client is None + assert ( + tuple( + task.get_coro().__qualname__ + for task in asyncio.all_tasks() + if task.get_coro().__qualname__.startswith( + ("PrismaClient.", "PrismaWrapper.", "_proxy_lifespan.._run_") + ) + ) + == () + ) + await asyncio.to_thread(_assert_database_idle) + + +async def _exercise_lifespans(mode: Mode) -> None: + from litellm.proxy import proxy_server + from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError + + expected: Final = UnsafeMasterKeyError if mode == "unsafe" else ValueError + match: Final = "master key" if mode == "unsafe" else "use_redis_transaction_buffer" + with pytest.raises(expected, match=match): + async with proxy_server.proxy_startup_event(proxy_server.app): + pytest.fail("Invalid configuration started the proxy") + await _assert_lifespan_closed() + + Path(os.environ["CONFIG_FILE_PATH"]).write_text(CONFIG) + with pytest.MonkeyPatch.context() as environment: + environment.setenv("LITELLM_MASTER_KEY", f"sk-{uuid.uuid4().hex}") + for cancel in (False, True, False): + try: + async with proxy_server.proxy_startup_event(proxy_server.app) as state: + database: Final = state["database"] + assert database is not None + assert await database.prisma.writer_db.query_raw("SELECT 42 AS answer") == [{"answer": 42}] + assert await database.prisma.read_db.query_raw("SELECT 43 AS answer") == [{"answer": 43}] + assert psutil.Process().children(recursive=True) + if cancel: + raise asyncio.CancelledError + except asyncio.CancelledError: + if not cancel: + raise + await _assert_lifespan_closed() + + +@pytest.mark.parametrize("mode", ("late", "unsafe")) +@pytest.mark.parametrize("replica", (False, True), ids=("writer", "writer-and-reader")) +def test_failed_startup_then_retry_releases_database_resources(tmp_path: Path, mode: Mode, replica: bool) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text(CONFIG + (" use_redis_transaction_buffer: true\n" if mode == "late" else "")) + inherited: Final = { + key: value + for key, value in os.environ.items() + if not key.startswith("REDIS_") + and key + not in ( + "WORKER_CONFIG", + "DATABASE_URL_READ_REPLICA", + "LITELLM_SALT_KEY", + "LITELLM_MIGRATE_FROM_MASTER_KEY", + "LITELLM_DANGEROUSLY_PERMIT_WEAK_OR_UNSET_MASTER_KEY", + "INTEGRATION_UPSTREAM_URL", + ) + } + with scratch_database("lifespan") as database_url: + environment: Final = { + **inherited, + "DATABASE_URL": database_url, + "DIRECT_URL": database_url, + "CONFIG_FILE_PATH": str(config), + "LITELLM_MASTER_KEY": PUBLICLY_KNOWN_KEY if mode == "unsafe" else f"sk-{uuid.uuid4().hex}", + "LITELLM_LOCAL_MODEL_COST_MAP": "True", + "LITELLM_TELEMETRY_ENABLED": "false", + "PRISMA_HEALTH_WATCHDOG_ENABLED": "true", + "PYTHONPATH": os.pathsep.join((str(ROOT), str(ROOT / "enterprise"))), + **({"DATABASE_URL_READ_REPLICA": database_url} if replica else {}), + } + schema: Final = subprocess.run( + ( + sys.executable, + "-P", + "-m", + "prisma", + "db", + "push", + "--skip-generate", + "--schema", + str(ROOT / "litellm/proxy/schema.prisma"), + ), + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=60, + ) + assert schema.returncode == 0, schema.stderr + probe: Final = subprocess.run( + ( + sys.executable, + "-P", + "-c", + "import asyncio; from tests.integration.database.test_proxy_database_lifespan import _exercise_lifespans; " + f"asyncio.run(_exercise_lifespans({mode!r}))", + ), + cwd=ROOT, + env=environment, + capture_output=True, + text=True, + timeout=90, + ) + assert probe.returncode == 0, probe.stdout + probe.stderr diff --git a/tests/unit/proxy/db/test_routing_prisma_wrapper.py b/tests/unit/proxy/db/test_routing_prisma_wrapper.py index 911a6d42ed3..73225acd2c8 100644 --- a/tests/unit/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/unit/proxy/db/test_routing_prisma_wrapper.py @@ -254,6 +254,55 @@ def test_token_refresh_delegates_to_both_writer_and_reader(): reader.stop_token_refresh_task.assert_awaited_once() +def test_reader_token_refresh_stops_even_when_the_writer_task_failed() -> None: + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + writer: Final = MagicMock() + reader: Final = MagicMock() + writer.stop_token_refresh_task = AsyncMock(side_effect=RuntimeError("refresh failed")) + reader.stop_token_refresh_task = AsyncMock() + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + + with pytest.raises(RuntimeError, match="refresh failed"): + asyncio.run(routing.stop_token_refresh_task()) + + reader.stop_token_refresh_task.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error", (RuntimeError("refresh failed"), asyncio.CancelledError())) +async def test_reader_cleanup_finishes_before_writer_stop_failure_propagates(error: BaseException) -> None: + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + + reader_stopping: Final = asyncio.Event() + release_reader: Final = asyncio.Event() + reader_stopped: Final = asyncio.Event() + + async def stop_reader() -> None: + reader_stopping.set() + await release_reader.wait() + reader_stopped.set() + + writer: Final = MagicMock() + reader: Final = MagicMock() + writer.stop_token_refresh_task = AsyncMock(side_effect=error) + reader.stop_token_refresh_task = AsyncMock(side_effect=stop_reader) + routing: Final = RoutingPrismaWrapper(writer=writer, reader=reader) + stopping: Final = asyncio.create_task(routing.stop_token_refresh_task()) + try: + await reader_stopping.wait() + await asyncio.sleep(0) + assert not stopping.done() + release_reader.set() + with pytest.raises(type(error)): + await stopping + assert reader_stopped.is_set() + finally: + release_reader.set() + stopping.cancel() + await asyncio.gather(stopping, return_exceptions=True) + + def test_routed_actions_falls_back_to_writer_for_unknown_methods(): from litellm.proxy.db.routing_prisma_wrapper import _RoutedActions diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index a8706c5bccd..404c31f59e1 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -24,6 +24,7 @@ import logging import os import subprocess from collections.abc import Awaitable, Callable +from contextlib import aclosing from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -35,7 +36,6 @@ from typing_extensions import TypedDict import litellm.proxy.proxy_server as ps from litellm.proxy.management_helpers.audit_logs import track_audit_task from litellm.proxy.proxy_server import ( - _proxy_startup_event, ProxyStartupEvent, _resolve_pydantic_type, _resolve_typed_dict_type, @@ -51,6 +51,7 @@ from litellm.proxy.proxy_server import ( run_ollama_serve, save_worker_config, ) +from litellm.repositories import database as database_runtime from .conftest import normalize @@ -101,9 +102,10 @@ def test_cleanup_router_config_variables_fails_on_unknown_attr_raises(): @pytest.mark.asyncio -async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): +async def test_database_lifespan_disconnects_prisma_after_proxy_shutdown(monkeypatch): fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock() + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) monkeypatch.setattr(ps, "master_key", "sk-x", raising=False) @@ -117,7 +119,9 @@ async def test_proxy_shutdown_event_disconnects_prisma_and_resets(monkeypatch): monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() + fake_prisma.disconnect.assert_not_awaited() observed = { "disconnect_called": fake_prisma.disconnect.await_count == 1, @@ -204,6 +208,7 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) async def _record_flush(client, accumulator): @@ -222,7 +227,8 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - await proxy_shutdown_event() + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await proxy_shutdown_event() assert calls == ["flush", "disconnect"] @@ -239,6 +245,7 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor fake_prisma = MagicMock() fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + fake_prisma.stop_background_tasks = AsyncMock() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) async def _record_flush(client, accumulator): @@ -262,7 +269,8 @@ async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_befor monkeypatch.setattr(litellm, "success_callback", [], raising=False) track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) - await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + async with aclosing(database_runtime.Database(prisma=fake_prisma)): + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) assert calls == ["flush", "audit", "disconnect"] @@ -318,9 +326,10 @@ async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_t @pytest.mark.asyncio -async def test_proxy_shutdown_event_prisma_disconnect_raises_error(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() monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) fake_jwt = MagicMock() @@ -332,8 +341,12 @@ async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): monkeypatch.setattr(litellm, "cache", None, raising=False) monkeypatch.setattr(litellm, "success_callback", [], raising=False) - with pytest.raises(RuntimeError, match="db gone"): - await proxy_shutdown_event() + 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 # --------------------------------------------------------------------------- @@ -1008,7 +1021,9 @@ async def test_proxy_shutdown_commits_spend_before_waiting_on_a_slow_logging_tas task_may_finish: Final = asyncio.Event() task_finished: Final = asyncio.Event() - flushes: Final[list[tuple[str, bool]]] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion + flushes: Final[ + list[tuple[str, bool]] + ] = [] # mutable-ok: records what each shutdown flush saw, which is the assertion async def slow_logging_task() -> None: await task_may_finish.wait() @@ -1060,7 +1075,7 @@ def test_otel_global_provider_published_after_callback_init(): span was orphaned. The publish (``publish_global_otel_v2_provider``) must therefore appear after ``_initialize_startup_logging`` in the lifespan source. """ - wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(") publish_pos = source.find("publish_global_otel_v2_provider(") @@ -1146,7 +1161,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): The call must follow Prisma setup so DB-backed deployments do not false-positive. Direct ``_warn_budget_without_db`` tests cover the warning behavior itself. """ - wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event) + wrapped = getattr(ps._proxy_lifespan, "__wrapped__", ps._proxy_lifespan) source = inspect.getsource(wrapped) budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:") warn_pos = source.find("warn_budget_without_db(") diff --git a/tests/unit/proxy/test_proxy_server.py b/tests/unit/proxy/test_proxy_server.py index 139cfbbe58c..b3725f01512 100644 --- a/tests/unit/proxy/test_proxy_server.py +++ b/tests/unit/proxy/test_proxy_server.py @@ -116,7 +116,7 @@ async def test_admin_mcp_failure_still_closes_proxy_resources( message: Final = "LITELLM_LICENSE" if failure_phase == "license" else f"{failure_phase} failed" async def run_lifespan() -> None: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} if failure_phase == "serving": raise RuntimeError("serving failed") if failure_phase == "cancelled": @@ -168,7 +168,7 @@ async def test_proxy_shutdown_drains_active_admin_tool_before_closing_connector( transport=httpx2.ASGITransport(app=admin_mcp_proxy), base_url="http://localhost:4000" ) as client: async with admin_mcp_proxy.router.lifespan_context(admin_mcp_proxy) as state: - assert state == {"tracing_receiver": None} + assert state == {"tracing_receiver": None, "database": None} request: Final = asyncio.create_task(client.post( "/admin/mcp", headers={"Authorization": "Bearer admin", "Accept": "application/json, text/event-stream"}, diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index bc3a228321a..46da234c8d3 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -1654,6 +1654,9 @@ class MockPrisma: async def disconnect(self): pass + async def stop_background_tasks(self): + pass + mock_prisma = MockPrisma() @@ -1764,7 +1767,7 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError from litellm.proxy.proxy_server import proxy_startup_event - async def unreachable(): + async def unreachable(resources): raise ConnectionError("database is down") _, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings) @@ -1809,7 +1812,7 @@ async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values database = _DatabaseWithOneStoredCredential(encrypt_value_helper("sk-provider", new_encryption_key=encrypted_with)) - async def connected(): + async def connected(resources): return database _, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": PUBLICLY_KNOWN_KEY}) @@ -1844,6 +1847,15 @@ async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_delet class _PrismaClientWhoseDatabaseRejectsQueries: + def __init__(self) -> None: + self.disconnect_calls = 0 + + async def disconnect(self) -> None: + self.disconnect_calls += 1 + + async def stop_background_tasks(self) -> None: + pass + class _Database: async def query_raw(self, query, *args): raise RuntimeError("permission denied for table LiteLLM_CredentialsTable") @@ -1858,7 +1870,8 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc from litellm.proxy.proxy_server import proxy_startup_event _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) - monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + prisma: Final = _PrismaClientWhoseDatabaseRejectsQueries() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) with ( @@ -1871,6 +1884,25 @@ async def test_proxy_startup_stops_when_the_requested_migration_fails(monkeypatc notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message] assert len(notices) == 1 assert "Could not migrate stored values" in notices[0] + assert prisma.disconnect_calls == 1 + + +@pytest.mark.asyncio +async def test_proxy_startup_after_a_failed_migration_does_not_reuse_the_closed_client(monkeypatch, tmp_path): + from fastapi import FastAPI + + from litellm.proxy.proxy_server import proxy_startup_event + + _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _PrismaClientWhoseDatabaseRejectsQueries()) + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", PUBLICLY_KNOWN_KEY) + with pytest.raises(RuntimeError, match="permission denied"): + async with proxy_startup_event(FastAPI()): + pass + + monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "") + async with proxy_startup_event(FastAPI()) as state: + assert state["database"] is None @pytest.mark.asyncio @@ -1917,11 +1949,11 @@ async def test_proxy_startup_boots_an_unsafe_master_key_under_the_override(monke class _ShutdownAwarePrisma(MockPrisma): def __init__(self): super().__init__() - self.stop_view_setup_task = AsyncMock() + self.stop_background_tasks = AsyncMock() @pytest.mark.asyncio -async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): +async def test_proxy_shutdown_stops_database_background_tasks(monkeypatch, tmp_path): import yaml from fastapi import FastAPI @@ -1936,11 +1968,11 @@ async def test_proxy_shutdown_stops_the_view_setup_task(monkeypatch, tmp_path): monkeypatch.setattr(proxy_server_module, "store_model_in_db", False) async with proxy_startup_event(FastAPI()): - stopped_while_serving = fake_prisma.stop_view_setup_task.await_count + stopped_while_serving = fake_prisma.stop_background_tasks.await_count actual = { "stopped_while_serving": stopped_while_serving, - "stopped_after_shutdown": fake_prisma.stop_view_setup_task.await_count, + "stopped_after_shutdown": fake_prisma.stop_background_tasks.await_count, } assert actual == {"stopped_while_serving": 0, "stopped_after_shutdown": 1} @@ -14231,6 +14263,8 @@ async def test_team_window_spend_row_carries_the_request_start_time(): def _mock_startup_prisma_client(health_check_error=None, connect_error=None): client = MagicMock() client.connect = AsyncMock(side_effect=connect_error) + client.disconnect = AsyncMock() + client.stop_background_tasks = AsyncMock() client.db.start_token_refresh_task = AsyncMock() client.check_view_exists = AsyncMock() client._set_spend_logs_row_count_in_proxy_state = AsyncMock() @@ -14343,6 +14377,8 @@ async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(mon mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) with pytest.raises(httpx.ReadTimeout): await _run_setup_prisma_client(mock_client) + mock_client.stop_background_tasks.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() @pytest.mark.asyncio @@ -14363,6 +14399,31 @@ async def test_setup_prisma_client_returns_none_when_connect_itself_fails(monkey assert result is None assert mock_client.start_db_health_watchdog_task.await_count == 0 assert mock_client.health_check.await_count == 0 + mock_client.stop_background_tasks.assert_awaited_once() + mock_client.disconnect.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_cancelling_prisma_setup_closes_the_client_before_publishing_it(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DISABLE_PRISMA_HEALTH_CHECK_ON_STARTUP", "False") + checking: Final = asyncio.Event() + client: Final = _mock_startup_prisma_client() + + async def health_check() -> None: + checking.set() + await asyncio.Event().wait() + + client.health_check = health_check + setup: Final = asyncio.create_task(_run_setup_prisma_client(client)) + try: + await asyncio.wait_for(checking.wait(), timeout=2) + finally: + setup.cancel() + with pytest.raises(asyncio.CancelledError): + await setup + + client.stop_background_tasks.assert_awaited_once() + client.disconnect.assert_awaited_once() async def _run_scheduled_background_jobs(): diff --git a/tests/unit/proxy/utils/prisma_and_spend/conftest.py b/tests/unit/proxy/utils/prisma_and_spend/conftest.py index aefb95516af..afc282a14d4 100644 --- a/tests/unit/proxy/utils/prisma_and_spend/conftest.py +++ b/tests/unit/proxy/utils/prisma_and_spend/conftest.py @@ -139,6 +139,7 @@ def mock_prisma_client() -> MagicMock: client.db.is_connected = MagicMock(return_value=False) client.db.connect = AsyncMock() client.db.disconnect = AsyncMock() + client.db.stop_token_refresh_task = AsyncMock() client.db.query_raw = AsyncMock(return_value=[{"?column?": 1}]) client.db.execute_raw = AsyncMock() client.db.tx = MagicMock() 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 9aa57c7a19b..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,162 @@ from __future__ import annotations import asyncio import logging -from typing import Any -from unittest.mock import AsyncMock, MagicMock +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 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("engine_dead", (False, True)) +async def test_shutdown_waits_for_hung_request_reconnect_to_time_out( + prisma_client: PrismaClient, monkeypatch: pytest.MonkeyPatch, engine_dead: bool +) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql://test:test@localhost/test") + reconnecting: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + + async def hung_recreate(*args: object, **kwargs: object) -> None: + reconnecting.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + prisma_client.db.recreate_prisma_client = AsyncMock(side_effect=hung_recreate) + prisma_client._engine_confirmed_dead = engine_dead + with patch.object(asyncio.get_running_loop(), "time", return_value=1000.0) as clock: + reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="request", force=True, force_recreate=True) + ) + await reconnecting.wait() + closing: Final = asyncio.create_task(Database(prisma=prisma_client).aclose()) + try: + await asyncio.sleep(0) + assert not closing.done() + prisma_client.db.disconnect.assert_not_awaited() + + clock.return_value = 1000.0 + prisma_client._db_watchdog_reconnect_timeout_seconds + 1 + assert await reconnect is False + await closing + finally: + reconnect.cancel() + closing.cancel() + await asyncio.gather(reconnect, closing, return_exceptions=True) + + assert cancelled.is_set() + assert not prisma_client._db_reconnect_lock.locked() + prisma_client.db.disconnect.assert_awaited_once() + assert await prisma_client.attempt_db_reconnect(reason="request", force=True) is False + prisma_client.db.recreate_prisma_client.assert_awaited_once() + + +@pytest.mark.asyncio +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()) + try: + await asyncio.sleep(0) + assert not closing.done() + prisma_client.db.disconnect.assert_not_awaited() + + clock.return_value = 1000.0 + _SHUTDOWN_RECONNECT_WAIT_SECONDS + 1 + await closing + finally: + closing.cancel() + await asyncio.gather(closing, return_exceptions=True) + + prisma_client.db.disconnect.assert_awaited_once() + prisma_client._db_reconnect_lock.release() + await asyncio.sleep(0) + assert not prisma_client._db_reconnect_lock.locked() + + +@pytest.mark.asyncio +async def test_shutdown_rejects_a_reconnect_already_waiting_for_the_lock(prisma_client: PrismaClient) -> None: + await prisma_client._db_reconnect_lock.acquire() + reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) + ) + await asyncio.sleep(0) + closing: Final = asyncio.create_task(prisma_client.stop_background_tasks()) + await asyncio.sleep(0) + prisma_client._db_reconnect_lock.release() + + await asyncio.wait_for(closing, timeout=2) + assert await reconnect is False + assert await prisma_client.attempt_db_reconnect(reason="engine_process_death", force=True) is False + prisma_client.db.connect.assert_not_awaited() + prisma_client.db.disconnect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_shutdown_waits_for_spawned_engine_work_and_prevents_rearming(prisma_client: PrismaClient) -> None: + running: Final = asyncio.Event() + cancelled: Final = asyncio.Event() + released: Final = asyncio.Event() + finished: Final = asyncio.Event() + + async def engine_work() -> None: + running.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + await released.wait() + finished.set() + + prisma_client._spawn_engine_task(engine_work()) + await asyncio.wait_for(running.wait(), timeout=2) + closing: Final = asyncio.create_task(prisma_client.stop_background_tasks()) + try: + await asyncio.wait_for(cancelled.wait(), timeout=2) + assert not finished.is_set() + assert not closing.done() + finally: + released.set() + await asyncio.wait_for(closing, timeout=2) + + assert finished.is_set() + prisma_client._handle_writer_engine_replaced() + await asyncio.sleep(0) + 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 new file mode 100644 index 00000000000..957c4cefbfb --- /dev/null +++ b/tests/unit/repositories/test_database.py @@ -0,0 +1,36 @@ +import logging +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.repositories.database import Database + + +@pytest.mark.asyncio +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 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