From fbcdacc44659813b05dd2f9d145ec63575e2843d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 21 Apr 2026 23:29:17 -0700 Subject: [PATCH] [Fix] Proxy: reconnect Prisma DB without blocking the event loop When the DB becomes unreachable the reconnect path calls `prisma.disconnect()`, which ultimately invokes prisma-client-py's synchronous `subprocess.Popen.wait()` on the query engine subprocess. That call does not yield to asyncio, so the event loop freezes for however long the Rust engine takes to shut down (30-120+ seconds in production when the engine is stuck on TCP close). During the freeze `/health/liveliness` becomes unresponsive, and in Kubernetes the liveness probe fails and the pod is SIGKILL'd. Replace `disconnect()` in the reconnect paths with a direct, non-blocking kill of the engine subprocess (SIGTERM -> 0.5s asyncio-yielding sleep -> SIGKILL) followed by a fresh Prisma client and a new `connect()`. Both `recreate_prisma_client` and the formerly-separate "direct reconnect" path go through the same kill-then-recreate flow. Also validate `_get_engine_pid` returns an int (defensive; prevents a MagicMock leak under unit-test mocking). Tests that encoded the old blocking behavior are updated or removed; the deleted `test_lightweight_reconnect_skips_kill_on_successful_disconnect` invariant ("don't kill on successful disconnect") was part of the bug. --- litellm/proxy/db/prisma_client.py | 19 +-- litellm/proxy/utils.py | 23 ++-- .../proxy/test_prisma_engine_watchdog.py | 77 +++++++----- .../proxy/db/test_prisma_self_heal.py | 119 +++++++++--------- 4 files changed, 135 insertions(+), 103 deletions(-) diff --git a/litellm/proxy/db/prisma_client.py b/litellm/proxy/db/prisma_client.py index 73735796eb..a8942f92a1 100644 --- a/litellm/proxy/db/prisma_client.py +++ b/litellm/proxy/db/prisma_client.py @@ -52,7 +52,9 @@ class PrismaWrapper: engine = self._original_prisma._engine process = getattr(engine, "process", None) if engine is not None else None if process is not None: - return process.pid + pid = process.pid + if isinstance(pid, int): + return pid except (AttributeError, TypeError): pass return 0 @@ -217,15 +219,18 @@ class PrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, http_client: Optional[Any] = None ): - """Disconnect and reconnect the Prisma client with a new database URL.""" + """Disconnect and reconnect the Prisma client with a new database URL. + + Kills the old engine subprocess directly (SIGTERM → SIGKILL) rather than + calling `disconnect()`. prisma-client-py's `disconnect()` calls a + synchronous `subprocess.Popen.wait()` that can freeze the asyncio event + loop for 30-120+ seconds when the engine is stuck on TCP close, + breaking `/health/liveliness` and causing Kubernetes pod restarts. + """ from prisma import Prisma # type: ignore old_engine_pid = self._get_engine_pid() - - try: - await self._original_prisma.disconnect() - except Exception as e: - verbose_proxy_logger.warning(f"Failed to disconnect Prisma client: {e}") + if old_engine_pid > 0: await self._kill_engine_process(old_engine_pid) if http_client is not None: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f21a729f55..48931a0c28 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4129,17 +4129,20 @@ class PrismaClient: ) async def _do_direct_reconnect() -> None: - old_pid = self._get_engine_pid() - try: - await self.db.disconnect() - except Exception as disconnect_err: - verbose_proxy_logger.warning( - "Prisma DB disconnect before reconnect failed: %s", - disconnect_err, + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error( + "DATABASE_URL not set; cannot reconnect Prisma client." ) - await PrismaWrapper._kill_engine_process(old_pid) - - await self.db.connect() + raise RuntimeError("DATABASE_URL not set") + # Fresh Prisma client + new engine subprocess. The previous + # "lightweight" path called `disconnect()` which blocks the + # event loop on `subprocess.Popen.wait()`; since that call + # ends up killing the engine anyway, we do it non-blockingly + # via `_kill_engine_process` inside `recreate_prisma_client`. + self._cleanup_engine_watcher() + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() await self.db.query_raw("SELECT 1") await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py index 786167b948..0d241f7574 100644 --- a/tests/litellm/proxy/test_prisma_engine_watchdog.py +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -254,32 +254,48 @@ async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( @pytest.mark.asyncio -async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( +async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive( engine_client, ) -> None: - """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" - engine_client._engine_pid = 1234 + """Direct reconnect (engine alive) calls recreate_prisma_client + SELECT 1. - with patch.object(engine_client, "_is_engine_alive", return_value=True): + The old "lightweight" path called `disconnect()` + `connect()`, which + blocks the event loop on the sync `process.wait()` inside aclose(). + The fix routes both engine-alive and engine-dead paths through + `recreate_prisma_client`, which non-blockingly kills the old engine. + """ + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=True), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await engine_client._run_reconnect_cycle(timeout_seconds=5.0) - engine_client.db.connect.assert_awaited_once() + engine_client.db.recreate_prisma_client.assert_awaited_once_with( + "postgresql://test" + ) engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") - engine_client.db.recreate_prisma_client.assert_not_awaited() + engine_client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio -async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( +async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown( engine_client, ) -> None: - """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + """When the engine PID is not tracked, direct reconnect still runs.""" engine_client._engine_pid = 0 + engine_client._start_engine_watcher = AsyncMock() - await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) - engine_client.db.connect.assert_awaited_once() + engine_client.db.recreate_prisma_client.assert_awaited_once_with( + "postgresql://test" + ) engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") - engine_client.db.recreate_prisma_client.assert_not_awaited() + engine_client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio @@ -473,36 +489,38 @@ def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): @pytest.mark.asyncio -async def test_escalation_after_consecutive_lightweight_failures(engine_client): - """After N consecutive lightweight reconnect failures, _engine_confirmed_dead +async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client): + """After N consecutive direct reconnect failures, _engine_confirmed_dead is set to True so _run_reconnect_cycle takes the heavy reconnect path.""" engine_client._reconnect_escalation_threshold = 3 engine_client._consecutive_reconnect_failures = 0 engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test + engine_client._start_engine_watcher = AsyncMock(return_value=None) - # Make lightweight reconnect fail every time - engine_client.db.disconnect = AsyncMock(return_value=None) - engine_client.db.connect = AsyncMock(side_effect=Exception("connect failed")) + # Make direct reconnect fail every time + engine_client.db.recreate_prisma_client = AsyncMock( + side_effect=Exception("recreate failed") + ) # Run 3 failed reconnect attempts - for i in range(3): - result = await engine_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=5.0 - ) - assert result is False + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + for _ in range(3): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) + assert result is False assert engine_client._consecutive_reconnect_failures == 3 - # Next attempt should escalate: _engine_confirmed_dead set to True before _run_reconnect_cycle + # Next attempt should escalate to the heavy path (recreate_prisma_client still + # the call, but via the _engine_confirmed_dead branch that also re-arms the watcher). engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) - engine_client._start_engine_watcher = AsyncMock(return_value=None) with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): result = await engine_client._attempt_reconnect_inside_lock( force=True, reason="test_escalation", timeout_seconds=5.0 ) - # Heavy reconnect should have been attempted (recreate_prisma_client called) engine_client.db.recreate_prisma_client.assert_awaited_once() @@ -511,15 +529,16 @@ async def test_successful_reconnect_resets_failure_counter(engine_client): """A successful reconnect resets _consecutive_reconnect_failures to 0.""" engine_client._consecutive_reconnect_failures = 2 engine_client._db_reconnect_cooldown_seconds = 0 + engine_client._start_engine_watcher = AsyncMock() # Make reconnect succeed - engine_client.db.disconnect = AsyncMock(return_value=None) - engine_client.db.connect = AsyncMock(return_value=None) + engine_client.db.recreate_prisma_client = AsyncMock(return_value=None) engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) - result = await engine_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=5.0 - ) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await engine_client._attempt_reconnect_inside_lock( + force=True, reason="test", timeout_seconds=5.0 + ) assert result is True assert engine_client._consecutive_reconnect_failures == 0 diff --git a/tests/test_litellm/proxy/db/test_prisma_self_heal.py b/tests/test_litellm/proxy/db/test_prisma_self_heal.py index fb215e5477..57693a1a17 100644 --- a/tests/test_litellm/proxy/db/test_prisma_self_heal.py +++ b/tests/test_litellm/proxy/db/test_prisma_self_heal.py @@ -34,18 +34,18 @@ async def test_attempt_db_reconnect_should_succeed(mock_proxy_logging): client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() - result = await client.attempt_db_reconnect( - reason="unit_test_reconnect_success", - force=True, - ) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + result = await client.attempt_db_reconnect( + reason="unit_test_reconnect_success", + force=True, + ) assert result is True - client.db.disconnect.assert_awaited_once() - client.db.connect.assert_awaited_once() + client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") client.db.query_raw.assert_awaited_once_with("SELECT 1") @@ -140,15 +140,19 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt( ) client._db_last_reconnect_attempt_ts = 0.0 client._db_reconnect_cooldown_seconds = 10 - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() # Use a counter-based mock to avoid StopIteration when time.time() is called # more times than expected (varies by Python version / internal code paths). fake_clock = iter(range(100, 10000)) - with patch( - "litellm.proxy.utils.time.time", side_effect=lambda: float(next(fake_clock)) + with ( + patch( + "litellm.proxy.utils.time.time", + side_effect=lambda: float(next(fake_clock)), + ), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), ): result = await client.attempt_db_reconnect( reason="unit_test_cooldown_timestamp_after_attempt", @@ -162,23 +166,28 @@ async def test_attempt_db_reconnect_should_set_cooldown_after_attempt( @pytest.mark.asyncio -async def test_run_reconnect_cycle_watchdog_should_use_direct_db_ops( +async def test_run_reconnect_cycle_watchdog_should_use_recreate_prisma_client( mock_proxy_logging, ): + """Direct reconnect goes through recreate_prisma_client (which non-blockingly + kills the old engine) instead of calling disconnect() — see issue #26191. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.disconnect = AsyncMock(side_effect=AssertionError("wrapper disconnect used")) - client.connect = AsyncMock(side_effect=AssertionError("wrapper connect used")) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) + client.db.disconnect = AsyncMock( + side_effect=AssertionError("disconnect must not be called") + ) + client.db.recreate_prisma_client = AsyncMock(return_value=None) client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + client._start_engine_watcher = AsyncMock() - await client._run_reconnect_cycle(timeout_seconds=None) + with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}): + await client._run_reconnect_cycle(timeout_seconds=None) - client.db.disconnect.assert_awaited_once() - client.db.connect.assert_awaited_once() + client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") client.db.query_raw.assert_awaited_once_with("SELECT 1") + client.db.disconnect.assert_not_awaited() @pytest.mark.asyncio @@ -189,19 +198,22 @@ async def test_run_reconnect_cycle_watchdog_should_use_default_timeout_budget( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) client._db_watchdog_reconnect_timeout_seconds = 0.1 - client.db.disconnect = AsyncMock(return_value=None) + client._start_engine_watcher = AsyncMock() - async def _slow_connect(): + async def _slow_recreate(_db_url): await asyncio.sleep(0.08) async def _slow_query(_query: str): await asyncio.sleep(0.08) return [{"result": 1}] - client.db.connect = AsyncMock(side_effect=_slow_connect) + client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate) client.db.query_raw = AsyncMock(side_effect=_slow_query) - with pytest.raises(asyncio.TimeoutError): + with ( + pytest.raises(asyncio.TimeoutError), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await client._run_reconnect_cycle(timeout_seconds=None) @@ -212,19 +224,22 @@ async def test_run_reconnect_cycle_timeout_should_use_single_overall_budget( client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(return_value=None) + client._start_engine_watcher = AsyncMock() - async def _slow_connect(): + async def _slow_recreate(_db_url): await asyncio.sleep(0.08) async def _slow_query(_query: str): await asyncio.sleep(0.08) return [{"result": 1}] - client.db.connect = AsyncMock(side_effect=_slow_connect) + client.db.recreate_prisma_client = AsyncMock(side_effect=_slow_recreate) client.db.query_raw = AsyncMock(side_effect=_slow_query) - with pytest.raises(asyncio.TimeoutError): + with ( + pytest.raises(asyncio.TimeoutError), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + ): await client._run_reconnect_cycle(timeout_seconds=0.1) @@ -319,42 +334,32 @@ async def test_db_health_watchdog_start_stop_lifecycle(mock_proxy_logging): @pytest.mark.asyncio -async def test_lightweight_reconnect_kills_engine_on_disconnect_failure( +async def test_recreate_prisma_client_kills_old_engine_without_disconnect( mock_proxy_logging, ): - """Lightweight reconnect must kill the old engine PID when disconnect() fails.""" + """recreate_prisma_client SIGTERMs the old engine PID directly rather than + calling `disconnect()`, which blocks the asyncio event loop on the sync + `subprocess.Popen.wait()` inside prisma-client-py — see issue #26191. + """ client = PrismaClient( database_url="mock://test", proxy_logging_obj=mock_proxy_logging ) - client.db.disconnect = AsyncMock(side_effect=Exception("disconnect failed")) - client.db.connect = AsyncMock(return_value=None) - client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + disconnect_mock = AsyncMock( + side_effect=AssertionError("disconnect must not be called on reconnect path") + ) + client.db._original_prisma.disconnect = disconnect_mock with ( - patch.object(client, "_get_engine_pid", return_value=9999), - patch("os.kill") as mock_kill, - patch("asyncio.sleep", new_callable=AsyncMock), + patch.object(client.db, "_get_engine_pid", return_value=9999), + patch("litellm.proxy.db.prisma_client.os.kill") as mock_kill, + patch("litellm.proxy.db.prisma_client.asyncio.sleep", new_callable=AsyncMock), ): - await client._run_reconnect_cycle(timeout_seconds=5.0) + # Return a Prisma instance whose connect() is awaitable. + fake_new_prisma = MagicMock() + fake_new_prisma.connect = AsyncMock(return_value=None) + with patch("prisma.Prisma", return_value=fake_new_prisma): + await client.db.recreate_prisma_client("postgresql://test") mock_kill.assert_any_call(9999, signal.SIGTERM) - client.db.connect.assert_awaited_once() - client.db.query_raw.assert_awaited_once_with("SELECT 1") - - -@pytest.mark.asyncio -async def test_lightweight_reconnect_skips_kill_on_successful_disconnect( - mock_proxy_logging, -): - """Lightweight reconnect must NOT kill when disconnect() succeeds.""" - client = PrismaClient( - database_url="mock://test", proxy_logging_obj=mock_proxy_logging - ) - client.db.disconnect = AsyncMock(return_value=None) - client.db.connect = AsyncMock(return_value=None) - client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) - - with patch("os.kill") as mock_kill: - await client._run_reconnect_cycle(timeout_seconds=5.0) - - mock_kill.assert_not_called() + disconnect_mock.assert_not_awaited() + fake_new_prisma.connect.assert_awaited_once()