From d113e4a0dd586fbebb0ad5d94c0a820e6465b03f Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 15 Oct 2026 14:47:33 -0700 Subject: [PATCH] fix(proxy-extras): hold only Prisma's lock key during migrate deploy so it never waits on the request log index build (#46660) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_proxy_extras/migration_lock.py | 41 ++++++++++++++----- tests/e2e/migrations/test_startup.py | 8 ++-- .../database/test_prisma_migration_lock.py | 25 ++++++++++- 3 files changed, 58 insertions(+), 16 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py index ac54e3df9ae..85d971cfc06 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py +++ b/litellm-proxy-extras/litellm_proxy_extras/migration_lock.py @@ -18,6 +18,25 @@ MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big") PRISMA_ADVISORY_LOCK_KEY: Final = 72707369 _SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate") + +@dataclass(frozen=True, slots=True) +class AdvisoryLock: + key: int + name: str + holder: str + + +COORDINATOR_LOCK: Final = AdvisoryLock( + MIGRATION_LOCK_KEY, + "v2 migration coordinator lock", + "another v2 migration resolver", +) +PRISMA_LOCK: Final = AdvisoryLock( + PRISMA_ADVISORY_LOCK_KEY, + "Prisma migration lock", + "another Prisma migration", +) + if TYPE_CHECKING: import psycopg @@ -100,7 +119,10 @@ class MigrationCoordinator: @contextmanager -def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, None]: +def migration_lock( + database_url: str, + lock: AdvisoryLock = COORDINATOR_LOCK, +) -> Generator[MigrationCoordinator, None, None]: import psycopg wait_seconds: Final = migration_lock_timeout() @@ -112,11 +134,11 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N try: with connection: coordinator: Final = MigrationCoordinator(connection) - logger.info("Waiting for the v2 migration coordinator lock (up to %ss)", wait_seconds) + logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds) while time.monotonic() < deadline: with connection.transaction(): - if _try_lock(connection): - logger.info("Acquired the v2 migration coordinator lock") + if _try_lock(connection, lock.key): + logger.info("Acquired the %s", lock.name) yield coordinator coordinator.check_connection() @@ -125,18 +147,17 @@ def migration_lock(database_url: str) -> Generator[MigrationCoordinator, None, N except psycopg.Error as exc: raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None raise RuntimeError( - f"Timed out waiting for another v2 migration resolver after {wait_seconds}s. " + f"Timed out waiting for {lock.holder} after {wait_seconds}s. " f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}." ) @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() + """Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled. + + This keeps deploy independent of the coordinator key held by request-log index builds.""" + with migration_lock(database_url, PRISMA_LOCK) as coordinator: coordinator.release_statement_snapshot() yield diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py index 5ec8d7690f2..1f324f9cae8 100644 --- a/tests/e2e/migrations/test_startup.py +++ b/tests/e2e/migrations/test_startup.py @@ -97,11 +97,9 @@ class TestMigrationStartup: with containers.start( database, (COMPLETE,), environment={"LITELLM_MIGRATION_LOCK_TIMEOUT": "12"} ) as replica: - until( - "a replica holding the coordinator lock while Prisma's lock is held", - lambda: bool(database.holders(COORDINATOR_LOCK)), - ) - failed((replica,), "Timed out waiting for Prisma's lock") + until("Prisma lock contention", lambda: "Waiting for the Prisma migration lock" in replica.logs()) + failed((replica,), "Timed out waiting for another Prisma migration") + assert database.holders(COORDINATOR_LOCK) == () assert database.history() == history assert not database.exists("migration_effect") with containers.start(database, (COMPLETE,)) as restarted: diff --git a/tests/integration/database/test_prisma_migration_lock.py b/tests/integration/database/test_prisma_migration_lock.py index 0c6365b2b5b..7482e96d3d3 100644 --- a/tests/integration/database/test_prisma_migration_lock.py +++ b/tests/integration/database/test_prisma_migration_lock.py @@ -51,8 +51,31 @@ def test_prisma_migration_lock_holds_prismas_key_on_a_pinned_connection_and_leav 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 MIGRATION_LOCK_KEY not 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) + + +def test_prisma_migration_lock_does_not_wait_for_the_coordinator_lock( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_MIGRATION_LOCK_TIMEOUT", "3") + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer: + peer.execute("SELECT pg_advisory_lock(%s)", (MIGRATION_LOCK_KEY,)) + + with prisma_migration_lock(os.environ["DATABASE_URL"]): + assert PRISMA_ADVISORY_LOCK_KEY in _granted_advisory_keys(peer) + + +def test_prisma_migration_lock_times_out_when_another_prisma_migration_holds_its_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("LITELLM_MIGRATION_LOCK_TIMEOUT", "3") + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as peer: + peer.execute("SELECT pg_advisory_lock(%s)", (PRISMA_ADVISORY_LOCK_KEY,)) + + with pytest.raises(RuntimeError, match="Timed out waiting for another Prisma migration"): + with prisma_migration_lock(os.environ["DATABASE_URL"]): + pytest.fail("Prisma migration lock was acquired while another migration held it")