diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py index 671ce0f5c3e..89857859134 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py @@ -13,6 +13,7 @@ from litellm_proxy_extras.prisma_toolchain import ( ) MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big") +PRISMA_ADVISORY_LOCK_KEY: Final = 72707369 if TYPE_CHECKING: import psycopg @@ -32,7 +33,7 @@ def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]: if not database_url: return environment migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL")) - return {**environment, "DATABASE_URL": migration_url} + return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"} @dataclass(frozen=True, slots=True) @@ -55,10 +56,15 @@ class MigrationCoordinator: def check_connection(self) -> None: self.connection.execute("SELECT 1") + def release_statement_snapshot(self) -> None: + # a parameterized statement leaves its portal snapshot registered on the backend until the + # next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away + self.connection.execute("SELECT 1") + def acquire_prisma_lock(self) -> None: deadline: Final = time.monotonic() + migration_lock_timeout() while time.monotonic() < deadline: - if _try_lock(self.connection, 72707369): + if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY): return time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic()))) raise RuntimeError( @@ -95,6 +101,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N ) +@contextmanager +def prisma_migration_lock(database_url: str) -> Generator[None, None, None]: + """Hold the coordinator lock and Prisma's own lock key as transaction-level locks on one + pinned connection while a Prisma command runs with its session advisory lock disabled, so a + transaction pooler can never strand that lock on an idle backend after the command exits.""" + with migration_lock(database_url) as coordinator: + coordinator.acquire_prisma_lock() + coordinator.release_statement_snapshot() + yield + + @contextmanager def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]: """A session-level, non-blocking hold of the migration coordinator lock on an autocommit diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 2a109130ff8..4a311331fc0 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -9,6 +9,7 @@ import tempfile import threading import time from collections.abc import Callable +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, replace from pathlib import Path from typing import TYPE_CHECKING, Final, Optional, Protocol, Union @@ -1021,7 +1022,11 @@ class ProxyExtrasDBManager: def _setup_database_v2(use_migrate: bool) -> bool: if not use_migrate: return ProxyExtrasDBManager._run_database_v2(False) - from litellm_proxy_extras.migration_lock import migration_environment, migration_lock + from litellm_proxy_extras.migration_lock import ( + migration_environment, + migration_lock, + prisma_migration_lock, + ) from litellm_proxy_extras.migration_recovery import ( baseline_current_schema, recover_completed_migration, @@ -1055,7 +1060,9 @@ class ProxyExtrasDBManager: migration_environment(_get_prisma_env()), ) - while not ProxyExtrasDBManager._run_database_v2(True, recover_completed, baseline_existing): + while not ProxyExtrasDBManager._run_database_v2( + True, recover_completed, baseline_existing, lambda: prisma_migration_lock(lock_url) + ): continue return True @@ -1064,6 +1071,7 @@ class ProxyExtrasDBManager: use_migrate: bool, recover_completed: Callable[[str], bool] = lambda name: False, baseline_existing: "Callable[[str], None] | None" = None, + deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext, ) -> bool: """ v2 migration resolver (opt-in via --use_v2_migration_resolver). @@ -1126,11 +1134,12 @@ class ProxyExtrasDBManager: while not budget.exhausted: attempt_started = time.monotonic() try: - result = prisma_toolchain.run_prisma( - [_get_prisma_command(), "migrate", "deploy"], - timeout=deploy_timeout, - env=migration_env, - ) + with deploy_lock(): + result = prisma_toolchain.run_prisma( + [_get_prisma_command(), "migrate", "deploy"], + timeout=deploy_timeout, + env=migration_env, + ) logger.info(f"prisma migrate deploy stdout: {result.stdout}") return True @@ -1155,6 +1164,7 @@ class ProxyExtrasDBManager: e, budget, time.monotonic() - attempt_started, + deploy_lock=deploy_lock, ) if next_budget.attempts_left < budget.attempts_left: @@ -1187,15 +1197,19 @@ class ProxyExtrasDBManager: return match.group(1) if match else None @staticmethod - def _v2_roll_back_migration_best_effort(migration_name: str) -> None: + def _v2_roll_back_migration_best_effort( + migration_name: str, + deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext, + ) -> None: from litellm_proxy_extras.migration_lock import migration_environment try: - prisma_toolchain.run_prisma( - [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], - timeout=prisma_command_timeout(), - env=migration_environment(_get_prisma_env()), - ) + with deploy_lock(): + prisma_toolchain.run_prisma( + [_get_prisma_command(), "migrate", "resolve", "--rolled-back", migration_name], + timeout=prisma_command_timeout(), + env=migration_environment(_get_prisma_env()), + ) except (subprocess.CalledProcessError, subprocess.TimeoutExpired): pass @@ -1204,6 +1218,7 @@ class ProxyExtrasDBManager: error: subprocess.CalledProcessError, budget: "_MigrateAttemptBudget", attempt_seconds: float = 0.0, + deploy_lock: Callable[[], AbstractContextManager[None]] = nullcontext, ) -> "_MigrateAttemptBudget": """Recover from one failed `prisma migrate deploy`, and price the pass. @@ -1225,14 +1240,14 @@ class ProxyExtrasDBManager: "deadlock race, rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() if ledger_logs and _MIGRATION_LOCK_TIMEOUT_MARKER in ledger_logs: logger.info( "Migration %s timed out waiting for a table lock, rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() if ProxyExtrasDBManager._failed_migration_recovered(migration_name, started_at): logger.info( @@ -1276,7 +1291,7 @@ class ProxyExtrasDBManager: "rolling its ledger row back and retrying", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() if migration_name and _MIGRATION_LOCK_TIMEOUT_MARKER in stderr: @@ -1286,7 +1301,7 @@ class ProxyExtrasDBManager: "its ledger row back and retrying. Check for long-running transactions on the migrated tables.", migration_name, ) - ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name) + ProxyExtrasDBManager._v2_roll_back_migration_best_effort(migration_name, deploy_lock) return budget.spend() raise RuntimeError( diff --git a/tests/integration/database/test_prisma_migration_lock.py b/tests/integration/database/test_prisma_migration_lock.py new file mode 100644 index 00000000000..0c6365b2b5b --- /dev/null +++ b/tests/integration/database/test_prisma_migration_lock.py @@ -0,0 +1,58 @@ +import os +import secrets +from typing import Final + +import psycopg +import pytest +from litellm_proxy_extras.migration_lock import MIGRATION_LOCK_KEY, PRISMA_ADVISORY_LOCK_KEY, prisma_migration_lock + +pytestmark = pytest.mark.skipif("DATABASE_URL" not in os.environ, reason="needs a live Postgres") + + +def _can_take_prismas_session_lock(connection: psycopg.Connection[tuple[object, ...]]) -> bool: + row: Final = connection.execute("SELECT pg_try_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)).fetchone() + taken: Final = row is not None and row[0] is True + if taken: + connection.execute("SELECT pg_advisory_unlock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)) + return taken + + +def _granted_advisory_keys(connection: psycopg.Connection[tuple[object, ...]]) -> frozenset[int]: + rows: Final = connection.execute( + "SELECT (classid::bigint << 32) | objid::bigint FROM pg_locks WHERE locktype = 'advisory' AND granted" + ).fetchall() + return frozenset(int(row[0]) for row in rows) + + +def _index_builds_concurrently_now(connection: psycopg.Connection[tuple[object, ...]]) -> bool: + table: Final = f'"lit_cic_{secrets.token_hex(4)}"' + connection.execute(f"CREATE TABLE {table} (a int)") + try: + connection.execute("SET lock_timeout = '5s'") + connection.execute(f"CREATE INDEX CONCURRENTLY ON {table} (a)") + except psycopg.errors.LockNotAvailable: + return False + finally: + connection.execute("RESET lock_timeout") + connection.execute(f"DROP TABLE {table}") + return True + + +def test_prisma_migration_lock_does_not_stall_a_concurrent_index_build(): + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer: + with prisma_migration_lock(os.environ["DATABASE_URL"]): + assert _index_builds_concurrently_now(peer) + + +def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leaves_nothing_behind(): + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer: + assert _can_take_prismas_session_lock(peer) + + with prisma_migration_lock(os.environ["DATABASE_URL"]): + held: Final = _granted_advisory_keys(peer) + assert PRISMA_ADVISORY_LOCK_KEY in held + assert MIGRATION_LOCK_KEY in held + assert not _can_take_prismas_session_lock(peer) + + assert _granted_advisory_keys(peer).isdisjoint({PRISMA_ADVISORY_LOCK_KEY, MIGRATION_LOCK_KEY}) + assert _can_take_prismas_session_lock(peer) diff --git a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py index e60f049a8ed..75ea6e2a698 100644 --- a/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py +++ b/tests/unit/litellm_proxy_extras/test_litellm_proxy_extras_utils.py @@ -5,6 +5,7 @@ import re import subprocess import sys import threading +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Final, NoReturn, Optional @@ -2083,3 +2084,101 @@ def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch assert ProxyExtrasDBManager.repair_invalid_indexes() is False + +@pytest.mark.parametrize( + "environment", + ( + {"DATABASE_URL": "postgresql://pool/db"}, + {"DATABASE_URL": "postgresql://pool/db?pgbouncer=true", "DIRECT_URL": "postgresql://writer/db"}, + ), + ids=("pooled-only", "pooled-with-direct"), +) +def test_prisma_runs_its_migrations_without_a_session_advisory_lock(environment): + from litellm_proxy_extras.migration_lock import migration_environment + + migrated = migration_environment(environment) + + assert migrated["PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK"] == "1" + assert "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK" not in migration_environment({"PRISMA_OFFLINE_MODE": "true"}) + + +class TestMigrateDeployRunsUnderTheCoordinatorLock: + """`prisma migrate deploy` only runs while the injected lock is held, and the lock is + released again before the resolver handles the outcome, success or failure.""" + + @staticmethod + def _recording_lock(events: list[str]): + from contextlib import contextmanager + + @contextmanager + def lock(): + events.append("lock") + try: + yield + finally: + events.append("unlock") + + return lock + + def test_deploy_runs_between_lock_and_unlock(self, monkeypatch, tmp_path): + import litellm_proxy_extras.utils as utils_module + + events: list[str] = [] + harness = _MigrateDeployHarness(monkeypatch, tmp_path, ["ok"]) + original_run = harness._fake_run + + def run_and_record(cmd, **kwargs): + events.append("deploy") + return original_run(cmd, **kwargs) + + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record) + + assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True + assert events == ["lock", "deploy", "unlock"] + + def test_a_failed_deploy_releases_the_lock_before_the_retry(self, monkeypatch, tmp_path): + import litellm_proxy_extras.utils as utils_module + + events: list[str] = [] + harness = _MigrateDeployHarness( + monkeypatch, tmp_path, ["Error: P1002\nTimed out waiting for the advisory lock", "ok"] + ) + original_run = harness._fake_run + + def run_and_record(cmd, **kwargs): + events.append("deploy") + return original_run(cmd, **kwargs) + + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record) + + assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True + assert events == ["lock", "deploy", "unlock", "lock", "deploy", "unlock"] + + def test_a_deadlock_rollback_runs_under_the_coordinator_lock( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + import litellm_proxy_extras.utils as utils_module + + migration_name: Final = "20260921190000_agent_identity" + stderr: Final = ( + "Error: P3018\n\n" + "A migration failed to apply. New migrations cannot be applied before the error is recovered from.\n\n" + f"Migration name: {migration_name}\n\n" + "Database error code: 40P01\n\n" + f"Database error:\n{_P3009_DEADLOCK_LOGS}" + ) + events: Final[list[str]] = [] + harness: Final = _MigrateDeployHarness(monkeypatch, tmp_path, [stderr, "ok"]) + original_run: Final = harness._fake_run + + def run_and_record(cmd: Sequence[str], *, timeout: float, env: Mapping[str, str]) -> _FakeCompleted: + if cmd[1:] == ["migrate", "deploy"]: + events.append("deploy") + elif cmd[1:4] == ["migrate", "resolve", "--rolled-back"]: + events.append("rollback") + return original_run(cmd, timeout=timeout, env=env) + + monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", run_and_record) + + assert ProxyExtrasDBManager._run_database_v2(True, deploy_lock=self._recording_lock(events)) is True + assert events == ["lock", "deploy", "unlock", "lock", "rollback", "unlock", "lock", "deploy", "unlock"]