test(migrations): make the lock-timeout retry test independent of Prisma lock order

The test now holds the Prisma migration lock until every replica is waiting on it, so no replica can spend its
retry attempts before the others reach the gated migration

tests/e2e/migrations/containers.py

@@ -41,7 +41,6 @@ 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
 
@@ -212,7 +211,6 @@ class Containers:
             port: Final = int(docker("port", name, "4000/tcp").rsplit(":", 1)[1])
             replica: Final = Replica(
                 name=name,
-                identity=name,
                 transport=HttpTransport(f"http://127.0.0.1:{port}", MASTER_KEY, 15),
                 output=directory,
             )

tests/e2e/migrations/database.py

@@ -3,7 +3,6 @@ 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
@@ -98,27 +97,6 @@ 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]:

tests/e2e/migrations/test_startup.py

@@ -1,5 +1,4 @@
 from contextlib import ExitStack
-from itertools import accumulate, repeat
 from typing import Final
 
 import pytest
@@ -14,19 +13,13 @@ from .checks import (
     start_replicas,
 )
 from .containers import Containers, failed, ready, until, waiting
-from .database import COORDINATOR_LOCK, GATE_KEY, PRISMA_LOCK, Database, Databases, restricted_user
+from .database import COORDINATOR_LOCK, 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(
@@ -62,28 +55,18 @@ class TestMigrationStartup:
     ) -> 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)
-                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),
-                )
+                with database.lock(PRISMA_LOCK):
+                    replicas: Final = start_replicas(stack, containers, database, (GATED,))
+                    until(
+                        "every replica waiting for the Prisma migration lock",
+                        lambda: all("Waiting for the Prisma migration lock" in replica.logs() for replica in replicas),
+                    )
+                waiting(replicas, 20)
                 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)
-                    ),
+                    lambda: any("timed out waiting for a table lock" in replica.logs() for replica in replicas),
                 )
-                waiting((owner, *followers), 20)
-            ready((owner, *followers), database)
+            ready(replicas, database)
             assert_completed_after_lock_timeouts(database, GATED)
 
     @meta(