From 842826501d4df6fce8fc8d70ade95777faa9773e Mon Sep 17 00:00:00 2001 From: Patrick Buckley Date: Sun, 31 May 2026 16:00:34 -0700 Subject: [PATCH] fix(storage): prevent concurrent-migration deadlock on the advisory lock Migrations run on every node at boot, and one migration rebuilds an index with CREATE INDEX CONCURRENTLY, which can't run in a transaction and waits for all concurrent transactions to drain. The advisory lock that serialises migrations was held inside an open transaction, so the lock-holder's own idle-in- transaction connection deadlocked the concurrent build when several nodes started together. Acquire the lock on an AUTOCOMMIT connection and poll pg_try_advisory_lock so no waiter pins a snapshot. Adds a Postgres concurrency regression test (skipped on SQLite). --- tests/test_migration_concurrency.py | 121 ++++++++++++++++++++++++++++ turnstone/core/storage/_migrate.py | 23 ++++-- 2 files changed, 139 insertions(+), 5 deletions(-) create mode 100644 tests/test_migration_concurrency.py diff --git a/tests/test_migration_concurrency.py b/tests/test_migration_concurrency.py new file mode 100644 index 00000000..34db893a --- /dev/null +++ b/tests/test_migration_concurrency.py @@ -0,0 +1,121 @@ +"""Concurrency regression test for the PostgreSQL migration advisory lock. + +Reproduces the multi-node boot scenario: several workers run migrations against +the *same fresh* PostgreSQL database simultaneously (as 10 containers do on +``docker compose up``). Migration 041 rebuilds an index with ``CREATE INDEX +CONCURRENTLY``, which cannot run inside a transaction and waits for every +concurrent transaction to drain. A previous bug held ``pg_advisory_lock`` inside +an open transaction, so the lock-holder's own ``idle in transaction`` connection +deadlocked the concurrent index build — and waiters blocked on the lock piled on +more open transactions. The fix (``turnstone/core/storage/_migrate.py``) takes +the lock on an AUTOCOMMIT connection and polls ``pg_try_advisory_lock`` so no +waiter pins a snapshot. + +PostgreSQL-only — skipped on the SQLite backend (no CONCURRENTLY, no advisory +lock path). Runs in CI's ``test-postgres`` job (``--storage-backend=postgresql``). +""" + +from __future__ import annotations + +import os +import threading +import uuid +from typing import Any + +import pytest +import sqlalchemy as sa + + +def _pg_base_url() -> str: + return os.environ.get( + "TURNSTONE_TEST_PG_URL", + "postgresql+psycopg://postgres:postgres@localhost:5432/turnstone_test", + ) + + +@pytest.fixture +def fresh_pg_url(request: pytest.FixtureRequest) -> Any: + """Create a throwaway PostgreSQL database, yield its URL, drop it after. + + Skips unless the suite is running against PostgreSQL — migrations must run + from scratch (so 041's CONCURRENTLY actually executes), which the shared + ``turnstone_test`` schema can't provide. + """ + if request.config.getoption("--storage-backend") != "postgresql": + pytest.skip("PostgreSQL-only (advisory-lock / CREATE INDEX CONCURRENTLY path)") + + base = sa.make_url(_pg_base_url()) + db_name = f"ts_migtest_{uuid.uuid4().hex[:12]}" + # CREATE/DROP DATABASE can't run in a transaction → AUTOCOMMIT admin engine. + admin = sa.create_engine(base.set(database="postgres"), isolation_level="AUTOCOMMIT") + try: + with admin.connect() as conn: + conn.execute(sa.text(f'CREATE DATABASE "{db_name}"')) + yield base.set(database=db_name) + finally: + with admin.connect() as conn: + # Terminate any lingering backends before dropping. + conn.execute( + sa.text( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " + "WHERE datname = :d AND pid <> pg_backend_pid()" + ), + {"d": db_name}, + ) + conn.execute(sa.text(f'DROP DATABASE IF EXISTS "{db_name}"')) + admin.dispose() + + +class _EngineStub: + """Minimal stand-in for a StorageBackend — run_migrations only reads _engine.""" + + def __init__(self, engine: Any) -> None: + self._engine = engine + + +def test_concurrent_run_migrations_no_deadlock(fresh_pg_url: Any) -> None: + from turnstone.core.storage._migrate import run_migrations + + n_workers = 4 + errors: list[BaseException] = [] + barrier = threading.Barrier(n_workers) + + def _worker() -> None: + engine = sa.create_engine(fresh_pg_url) + try: + barrier.wait(timeout=30) # release together → maximise overlap + run_migrations(_EngineStub(engine), "postgresql") + except BaseException as exc: # noqa: BLE001 — capture for the assertion + errors.append(exc) + finally: + engine.dispose() + + threads = [threading.Thread(target=_worker, name=f"migrate-{i}") for i in range(n_workers)] + for t in threads: + t.start() + # A join timeout is essential: under the old bug these threads deadlock, and + # we want a clean test failure, not a hung suite. + for t in threads: + t.join(timeout=60) + + stuck = [t.name for t in threads if t.is_alive()] + assert not stuck, f"migration thread(s) deadlocked (alive after 60s): {stuck}" + assert not errors, f"migration(s) raised: {errors!r}" + + # All workers converged on head, and migration 041's CONCURRENTLY rebuild ran + # (partial parent index present, low-cardinality kind index dropped). + engine = sa.create_engine(fresh_pg_url) + try: + with engine.connect() as conn: + rev = conn.execute(sa.text("SELECT version_num FROM alembic_version")).scalar() + indexes = set( + conn.execute( + sa.text("SELECT indexname FROM pg_indexes WHERE tablename = 'workstreams'") + ).scalars() + ) + finally: + engine.dispose() + + assert rev is not None + assert "idx_workstreams_parent" in indexes + assert "idx_workstreams_kind" not in indexes diff --git a/turnstone/core/storage/_migrate.py b/turnstone/core/storage/_migrate.py index 47ff6531..6188c89e 100644 --- a/turnstone/core/storage/_migrate.py +++ b/turnstone/core/storage/_migrate.py @@ -54,9 +54,19 @@ def run_migrations(storage: Any, backend: str) -> None: def _run_with_pg_lock(engine: Any, cfg: Any) -> None: """Run Alembic upgrade under a PostgreSQL advisory lock. - Advisory lock ID 7_475_283 (arbitrary, derived from 'turnstone'). - ``pg_advisory_lock`` blocks until the lock is available, so - concurrent containers wait in line rather than racing. + Advisory lock ID 7_475_283 (arbitrary, derived from 'turnstone') serializes + concurrent containers so only one applies migrations at a time. + + The lock is held on an **autocommit** connection and acquired by *polling* + ``pg_try_advisory_lock`` rather than the blocking ``pg_advisory_lock``. Both + details are load-bearing. A migration that rebuilds an index with ``CREATE + INDEX CONCURRENTLY`` (migration 041) waits for every concurrent transaction + to drain before it can finish. If the lock-holder held the lock inside an + open transaction — or waiters blocked on ``pg_advisory_lock`` inside one — + those connections sit ``idle in transaction`` and never drain, so the + concurrent build deadlocks against the very lock meant to protect it. + Autocommit keeps the holder transaction-free; polling (with a sleep that + holds no snapshot) keeps waiters transaction-free between attempts. Retries with jittered backoff if PostgreSQL is temporarily at max_connections (common during large-cluster startup stampedes). @@ -71,12 +81,15 @@ def _run_with_pg_lock(engine: Any, cfg: Any) -> None: for attempt in range(max_retries): try: with engine.connect() as conn: - conn.execute(sa.text("SELECT pg_advisory_lock(7475283)")) + conn = conn.execution_options(isolation_level="AUTOCOMMIT") + # Poll, don't block: a waiter blocked inside pg_advisory_lock + # would pin a snapshot that CREATE INDEX CONCURRENTLY waits on. + while not conn.execute(sa.text("SELECT pg_try_advisory_lock(7475283)")).scalar(): + time.sleep(random.uniform(0.5, 1.5)) # noqa: S311 try: command.upgrade(cfg, "head") finally: conn.execute(sa.text("SELECT pg_advisory_unlock(7475283)")) - conn.commit() return except Exception as exc: err_str = str(exc).lower()