From 464a4dba9dc1dadf6d7affaa985aa7050596a745 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 11 Oct 2026 17:36:47 +0000 Subject: [PATCH] test(migrations): assert lock-timeout retry and exhaustion instead of an unbounded wait (#45955) * test(migrations): assert lock-timeout retry and exhaustion instead of an unbounded wait Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(migrations): accumulate follower lock wait observations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/scripts/run_migration_tests.py | 2 +- tests/e2e/migrations/checks.py | 19 +++++++ tests/e2e/migrations/containers.py | 10 +++- tests/e2e/migrations/database.py | 37 ++++++++++++-- tests/e2e/migrations/test_startup.py | 65 +++++++++++++++++++++--- 5 files changed, 121 insertions(+), 12 deletions(-) diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py index 5a73c54e3f5..50d01f1d50b 100644 --- a/.circleci/scripts/run_migration_tests.py +++ b/.circleci/scripts/run_migration_tests.py @@ -10,7 +10,7 @@ from typing import Final from xml.etree import ElementTree SUITES: Final = { - "startup": (("test_startup.py",), 12), + "startup": (("test_startup.py",), 13), "recovery": (("test_recovery.py",), 15), "legacy": (("test_legacy.py", "test_pooling.py"), 11), "upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5), diff --git a/tests/e2e/migrations/checks.py b/tests/e2e/migrations/checks.py index 4d67dbaa20d..a1748ec34e6 100644 --- a/tests/e2e/migrations/checks.py +++ b/tests/e2e/migrations/checks.py @@ -39,6 +39,25 @@ def assert_completed(database: Database, migration: Migration = COMPLETE) -> Non assert database.query("SELECT id FROM migration_effect") == ((1,),) +@step("Check that migration {migration.name} rolled back lock timeouts and completed exactly once") +def assert_completed_after_lock_timeouts(database: Database, migration: Migration) -> None: + rows: Final = database.query( + "SELECT finished_at IS NOT NULL, rolled_back_at IS NOT NULL, applied_steps_count FROM " + "_prisma_migrations WHERE migration_name = %s ORDER BY started_at", + (migration.name,), + ) + finished: Final = tuple(row for row in rows if row[0]) + rolled_back: Final = tuple(row for row in rows if row[1]) + assert finished == ((True, False, 1),), "Expected exactly one completed migration attempt" + assert rolled_back and all(row == (False, True, 0) for row in rolled_back), ( + "Expected every failed migration attempt to be rolled back" + ) + assert all(is_finished != is_rolled_back for is_finished, is_rolled_back, _ in rows), ( + "Expected every migration attempt to be either completed or rolled back" + ) + assert database.query("SELECT id FROM migration_effect") == ((1,),) + + @step("Apply the test migration by hand and record it in _prisma_migrations") def confirmed_history(database: Database) -> str: database.execute(COMPLETE_SQL) diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py index a05941373be..c606199553f 100644 --- a/tests/e2e/migrations/containers.py +++ b/tests/e2e/migrations/containers.py @@ -41,6 +41,7 @@ def until(description: str, condition: Callable[[], bool], seconds: float = 150) @dataclass(frozen=True, slots=True) class Replica: name: str + identity: str transport: HttpTransport output: Path @@ -171,7 +172,7 @@ class Containers: "model_list: []\ngeneral_settings:\n master_key: os.environ/LITELLM_MASTER_KEY\n" ) env: Final = { - "DATABASE_URL": prisma_url(database.container_url, database.schema), + "DATABASE_URL": prisma_url(database.container_url, database.schema, application_name=name), "LITELLM_MASTER_KEY": MASTER_KEY, "LITELLM_SALT_KEY": MASTER_KEY, "LITELLM_LOCAL_MODEL_COST_MAP": "True", @@ -209,7 +210,12 @@ class Containers: "4000", ) port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1]) - replica: Final = Replica(name, HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), directory) + replica: Final = Replica( + name=name, + identity=name, + transport=HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15), + output=directory, + ) yield replica finally: try: diff --git a/tests/e2e/migrations/database.py b/tests/e2e/migrations/database.py index 929a6aee817..9bf03a216c4 100644 --- a/tests/e2e/migrations/database.py +++ b/tests/e2e/migrations/database.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Generator from contextlib import contextmanager from dataclasses import dataclass +from itertools import chain from typing import Final, LiteralString from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit from uuid import uuid4 @@ -24,10 +25,19 @@ def connect_url(url: str, name: str) -> str: return urlunsplit(urlsplit(url)._replace(path=f"/{name}", query="")) -def prisma_url(url: str, schema: str) -> str: +def prisma_url(url: str, schema: str, application_name: str | None = None) -> str: parsed: Final = urlsplit(url) - query: Final = tuple((key, value) for key, value in parse_qsl(parsed.query) if key != "schema") - return urlunsplit(parsed._replace(query=urlencode((*query, ("schema", schema))))) + query: Final = tuple( + (key, value) + for key, value in parse_qsl(parsed.query) + if key != "schema" and (application_name is None or key != "application_name") + ) + parameters: Final = query + ( + (("schema", schema),) + if application_name is None + else (("schema", schema), ("application_name", application_name)) + ) + return urlunsplit(parsed._replace(query=urlencode(parameters))) @dataclass(frozen=True, slots=True) @@ -79,6 +89,27 @@ class Database: (key >> 32, key & 0xFFFFFFFF), ) + @step("List replica identities waiting on any of the advisory locks") + def waiting_identities(self, keys: tuple[int, ...]) -> frozenset[str]: + if not keys: + return frozenset() + key_pairs: Final = tuple((key >> 32, key & 0xFFFFFFFF) for key in keys) + key_predicates: Final = sql.SQL(" OR ").join( + sql.SQL("(waiting_lock.classid = %s AND waiting_lock.objid = %s)") for _ in key_pairs + ) + parameters: Final = tuple(chain.from_iterable(key_pairs)) + statement: Final = sql.SQL( + "SELECT DISTINCT activity.application_name " + "FROM pg_stat_activity AS activity " + "JOIN pg_locks AS waiting_lock ON waiting_lock.pid = activity.pid " + "WHERE activity.datname = current_database() " + "AND waiting_lock.locktype = 'advisory' AND NOT waiting_lock.granted " + "AND waiting_lock.database = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND ({})" + ).format(key_predicates) + identities: Final = self.query(statement, parameters) + return frozenset(row[0] for row in identities if isinstance(row[0], str)) + @step("Hold an advisory lock on the test database") @contextmanager def lock(self, key: int = GATE_KEY) -> Generator[None]: diff --git a/tests/e2e/migrations/test_startup.py b/tests/e2e/migrations/test_startup.py index ff7112c3881..856192937e2 100644 --- a/tests/e2e/migrations/test_startup.py +++ b/tests/e2e/migrations/test_startup.py @@ -1,17 +1,31 @@ from contextlib import ExitStack +from itertools import accumulate, repeat from typing import Final import pytest -from .checks import COMPLETE, FATAL, GATED, assert_completed, start_replicas +from .checks import ( + COMPLETE, + FATAL, + GATED, + assert_completed, + assert_completed_after_lock_timeouts, + start_replicas, +) from .containers import Containers, failed, ready, until, waiting -from .database import PRISMA_LOCK, Database, Databases, restricted_user +from .database import GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user from .startup_models import Migration from e2e_metadata import Domain, Subject, meta pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] +def union_waiting_identities( + observed: frozenset[str], snapshot: frozenset[str] +) -> frozenset[str]: + return observed | snapshot + + class TestMigrationStartup: @pytest.mark.parametrize("replicas,v2", ((1, True), (3, True), (1, False))) @meta( @@ -42,16 +56,34 @@ class TestMigrationStartup: domain=Domain.DB, ) ) - def test_waiters_survive_prolonged_contention(self, containers: Containers, database: Database) -> None: + def test_waiters_retry_lock_timeouts_until_contention_clears( + self, containers: Containers, database: Database + ) -> None: with ExitStack() as stack: with database.lock(): owner: Final = stack.enter_context(containers.start(database, (GATED,))) until("owner blocked in migration SQL", lambda: bool(database.blocked())) followers: Final = start_replicas(stack, containers, database, (GATED,), count=2) - until("both followers attempted Prisma locking", lambda: len(database.blocked(PRISMA_LOCK)) == 2) - waiting((owner, *followers), 120) + waiting_identities: Final = accumulate( + (database.waiting_identities((PRISMA_LOCK, GATE_KEY)) for _ in repeat(None)), + union_waiting_identities, + initial=frozenset(), + ) + follower_identities: Final = frozenset(follower.identity for follower in followers) + until( + "both followers entered migration locking", + lambda: follower_identities <= next(waiting_identities), + ) + until( + "a replica logged a migration lock-timeout retry", + lambda: any( + "timed out waiting for a table lock" in replica.logs() + for replica in (owner, *followers) + ), + ) + waiting((owner, *followers), 20) ready((owner, *followers), database) - assert_completed(database, GATED) + assert_completed_after_lock_timeouts(database, GATED) @meta( Subject( @@ -72,6 +104,27 @@ class TestMigrationStartup: ready((restarted,), database) assert_completed(database) + @meta( + Subject( + domain=Domain.DB, + ) + ) + def test_lock_timeout_exhaustion_then_restart(self, containers: Containers, database: Database) -> None: + with database.lock(): + with containers.start( + database, (GATED,), environment={"LITELLM_MIGRATION_DDL_LOCK_TIMEOUT": "2"} + ) as replica: + until("owner blocked in migration SQL", lambda: bool(database.blocked())) + failed((replica,), "attempts that made no progress") + assert not database.exists("migration_effect") + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE migration_name = %s AND finished_at IS NOT NULL", + (GATED.name,), + ) == ((0,),) + with containers.start(database, (GATED,)) as restarted: + ready((restarted,), database) + assert_completed_after_lock_timeouts(database, GATED) + @meta( Subject( domain=Domain.DB,