fix(prisma): add consecutive-failures gate before watchdog reconnect

A single failed SELECT 1 probe triggered an immediate engine kill and
respawn. A transient probe failure from brief event loop contention or
a momentary external condition now causes a full reconnect cycle,
turning a self-clearing hiccup into a hard outage.

This adds a consecutive probe failure counter to the watchdog loop.
Reconnect fires only after K consecutive probe failures, where K is
read from PRISMA_WATCHDOG_FAILURES_BEFORE_RECONNECT (default 1 to
preserve current behavior). Any successful probe resets the counter.
The counter is also reset when a reconnect is triggered so the gate
re-arms cleanly for the next failure window.
This commit is contained in:
gvisco 2026-06-02 17:00:57 +02:00
parent f48a87ef12
commit 3617858bdf
2 changed files with 108 additions and 4 deletions

View file

@ -2925,6 +2925,10 @@ class PrismaClient:
self._reconnect_escalation_threshold: int = max(
1, int(os.getenv("PRISMA_RECONNECT_ESCALATION_THRESHOLD", "3"))
)
self._watchdog_failures_before_reconnect: int = max(
1, int(os.getenv("PRISMA_WATCHDOG_FAILURES_BEFORE_RECONNECT", "1"))
)
self._consecutive_probe_failures: int = 0
self._engine_pidfd: int = -1
self._engine_pid: int = 0
self._watching_engine: bool = False
@ -4775,16 +4779,29 @@ class PrismaClient:
self.db.query_raw("SELECT 1"),
timeout=self._db_health_watchdog_probe_timeout_seconds,
)
self._consecutive_probe_failures = 0
except asyncio.CancelledError:
break
except Exception as e:
if isinstance(
e, asyncio.TimeoutError
) or PrismaDBExceptionHandler.is_database_connection_error(e):
await self.attempt_db_reconnect(
reason="db_health_watchdog_connection_error",
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
)
self._consecutive_probe_failures += 1
if (
self._consecutive_probe_failures
>= self._watchdog_failures_before_reconnect
):
self._consecutive_probe_failures = 0
await self.attempt_db_reconnect(
reason="db_health_watchdog_connection_error",
timeout_seconds=self._db_watchdog_reconnect_timeout_seconds,
)
else:
verbose_proxy_logger.debug(
"Prisma DB watchdog probe failure %d/%d; deferring reconnect.",
self._consecutive_probe_failures,
self._watchdog_failures_before_reconnect,
)
else:
verbose_proxy_logger.debug(
"Prisma DB health watchdog observed non-DB error: %s", e

View file

@ -486,3 +486,90 @@ async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
# The flag must STILL be True so the next attempt re-enters the heavy
# branch instead of silently demoting to the lightweight path.
assert client._engine_confirmed_dead is True
# ---------------------------------------------------------------------------
# Consecutive-failures gate (PR 3)
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_db_health_watchdog_defers_reconnect_below_threshold(mock_proxy_logging):
"""With threshold=3, two consecutive probe failures must NOT trigger reconnect."""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.query_raw = AsyncMock(side_effect=asyncio.TimeoutError())
client.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
client._db_health_watchdog_probe_timeout_seconds = 0.2
client._watchdog_failures_before_reconnect = 3
with patch(
"litellm.proxy.utils.asyncio.sleep",
AsyncMock(side_effect=[None, None, asyncio.CancelledError()]),
):
await client._db_health_watchdog_loop()
client.attempt_db_reconnect.assert_not_awaited()
assert client._consecutive_probe_failures == 2
@pytest.mark.asyncio
async def test_db_health_watchdog_triggers_reconnect_at_threshold(mock_proxy_logging):
"""With threshold=3, exactly three consecutive failures must trigger one reconnect."""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.query_raw = AsyncMock(side_effect=asyncio.TimeoutError())
client.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
client._db_health_watchdog_probe_timeout_seconds = 0.2
client._db_watchdog_reconnect_timeout_seconds = 7.0
client._watchdog_failures_before_reconnect = 3
with patch(
"litellm.proxy.utils.asyncio.sleep",
AsyncMock(side_effect=[None, None, None, asyncio.CancelledError()]),
):
await client._db_health_watchdog_loop()
client.attempt_db_reconnect.assert_awaited_once_with(
reason="db_health_watchdog_connection_error",
timeout_seconds=7.0,
)
@pytest.mark.asyncio
async def test_db_health_watchdog_resets_failure_counter_on_successful_probe(
mock_proxy_logging,
):
"""A successful probe resets the failure counter so two failures before and two
after do not add up to a threshold of three."""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db.query_raw = AsyncMock(
side_effect=[
asyncio.TimeoutError(),
asyncio.TimeoutError(),
[{"1": 1}],
asyncio.TimeoutError(),
asyncio.TimeoutError(),
]
)
client.attempt_db_reconnect = AsyncMock(return_value=True)
client._db_health_watchdog_interval_seconds = 1
client._db_health_watchdog_probe_timeout_seconds = 0.2
client._watchdog_failures_before_reconnect = 3
with patch(
"litellm.proxy.utils.asyncio.sleep",
AsyncMock(
side_effect=[None, None, None, None, None, asyncio.CancelledError()]
),
):
await client._db_health_watchdog_loop()
client.attempt_db_reconnect.assert_not_awaited()
assert client._consecutive_probe_failures == 2