diff --git a/crates/app/src/cmd_init.rs b/crates/app/src/cmd_init.rs index 1b302683..d9403413 100755 --- a/crates/app/src/cmd_init.rs +++ b/crates/app/src/cmd_init.rs @@ -203,32 +203,21 @@ pub async fn run(args: InitArgs) -> anyhow::Result { .await .map_err(|e| anyhow::anyhow!("{e:?}"))?; - // Check if catalog is already initialized. - let initialized = bootstrapper - .is_catalog_initialized() - .await - .map_err(|e| anyhow::anyhow!("{e:?}"))?; - - if initialized { - println!("--- Catalog already initialized. Use 'extenddb migrate' for pending migrations."); - } else { - bootstrapper - .run_catalog_migrations() - .await - .map_err(|e| anyhow::anyhow!("{e:?}"))?; - } - - // Record data database connection in catalog. + // Take the same lock `migrate` uses around the schema work below, so an + // init cannot race a migrate running on another replica and fail on a + // duplicate pg_type entry from concurrent `CREATE TABLE IF NOT EXISTS`. + // Two concurrent inits cannot get this far: the second aborts earlier, at + // `create_catalog_db`, because the database already exists. The catalog + // database does exist by this point, so the lock connection can be opened. bootstrapper - .record_data_connection() - .await - .map_err(|e| anyhow::anyhow!("{e:?}"))?; - - // Initialize data database schema. - bootstrapper - .run_data_migrations() + .acquire_migration_lock() .await .map_err(|e| anyhow::anyhow!("{e:?}"))?; + let migration_result = run_init_migrations(bootstrapper.as_ref()).await; + if let Err(e) = bootstrapper.release_migration_lock().await { + tracing::warn!("Failed to release migration lock: {e:?}"); + } + migration_result?; bootstrapper .bootstrap_encryption_key() @@ -300,6 +289,41 @@ pub async fn run(args: InitArgs) -> anyhow::Result { Ok(0) } +/// Apply the catalog and data schema while the migration lock is held. Split out +/// of `run` so that the lock is released on every path, including errors. +async fn run_init_migrations( + bootstrapper: &dyn extenddb_storage::bootstrapper::Bootstrapper, +) -> anyhow::Result<()> { + // Check if catalog is already initialized. + let initialized = bootstrapper + .is_catalog_initialized() + .await + .map_err(|e| anyhow::anyhow!("{e:?}"))?; + + if initialized { + println!("--- Catalog already initialized. Use 'extenddb migrate' for pending migrations."); + } else { + bootstrapper + .run_catalog_migrations() + .await + .map_err(|e| anyhow::anyhow!("{e:?}"))?; + } + + // Record data database connection in catalog. + bootstrapper + .record_data_connection() + .await + .map_err(|e| anyhow::anyhow!("{e:?}"))?; + + // Initialize data database schema. + bootstrapper + .run_data_migrations() + .await + .map_err(|e| anyhow::anyhow!("{e:?}"))?; + + Ok(()) +} + /// Extract a CLI argument value by flag name. fn extract_arg(args: &[String], flag: &str) -> Option { args.windows(2).find(|w| w[0] == flag).map(|w| w[1].clone()) diff --git a/crates/app/src/cmd_migrate.rs b/crates/app/src/cmd_migrate.rs index 46c89c61..f76f8ac6 100755 --- a/crates/app/src/cmd_migrate.rs +++ b/crates/app/src/cmd_migrate.rs @@ -53,6 +53,25 @@ pub async fn run(args: MigrateArgs) -> anyhow::Result<()> { .await .map_err(|e| anyhow::anyhow!("{e:?}"))?; + // Serialize concurrent migrators, such as several replicas starting at + // once, so they don't race each other applying schema changes. + bootstrap + .acquire_migration_lock() + .await + .map_err(|e| anyhow::anyhow!("{e:?}"))?; + let result = apply_migrations(bootstrap.as_ref(), &args).await; + if let Err(e) = bootstrap.release_migration_lock().await { + tracing::warn!("Failed to release migration lock: {e:?}"); + } + result +} + +/// Run the version checks and apply pending migrations while the migration lock +/// is held. Split out of `run` so that the lock is always released afterwards. +async fn apply_migrations( + bootstrap: &dyn extenddb_storage::bootstrapper::Bootstrapper, + args: &MigrateArgs, +) -> anyhow::Result<()> { // Show current catalog version. println!("--- Checking current catalog version..."); let current = bootstrap diff --git a/crates/storage-postgres/src/bootstrapper.rs b/crates/storage-postgres/src/bootstrapper.rs index 7b9335bf..d10a8255 100755 --- a/crates/storage-postgres/src/bootstrapper.rs +++ b/crates/storage-postgres/src/bootstrapper.rs @@ -7,6 +7,8 @@ //! teardown using PostgreSQL-specific DDL. Connection pools are created //! lazily as needed during the bootstrap sequence. +use std::time::Duration; + use async_trait::async_trait; use extenddb_storage::bootstrapper::{ AdminBootstrapResult, BootstrapConfig, Bootstrapper, @@ -18,11 +20,28 @@ use extenddb_storage::bootstrapper::{ use extenddb_storage::management_store::{OpError, OpResult}; use sqlx::PgPool; use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; -use tokio::sync::OnceCell; +use tokio::sync::{Mutex, OnceCell}; use crate::CATALOG_VERSION; use crate::migrations; +/// ExtendDB's advisory-lock namespace. PostgreSQL's two-argument form, +/// `pg_advisory_xact_lock(classid, objid)`, lets ExtendDB reserve one stable +/// `classid` by convention and assign a distinct `objid` to each internal lock. +/// Other applications in the same database must avoid choosing the same keys. +/// +/// The value itself is arbitrary. It only has to stay stable across releases and +/// differ from any other namespace we add later. +const ADVISORY_LOCK_NAMESPACE: i32 = 0x0045_4442; // 'E', 'D', 'B' +/// `objid` for the schema-migration lock, which serializes concurrent `migrate` +/// runs (for example several replicas starting at once). +const MIGRATION_LOCK_OBJID: i32 = 1; + +/// Maximum time to wait for another migrator. Normal contention should clear +/// well within this window; expiry indicates a peer that is likely wedged and +/// needs operator attention rather than an indefinitely stuck init Job. +const MIGRATION_LOCK_TIMEOUT: Duration = Duration::from_secs(5 * 60); + /// Utilities for bootstrapping a `PostgreSQL` backend store. /// /// Holds the bootstrap configuration and lazily-created connection pools. @@ -32,6 +51,9 @@ use crate::migrations; pub struct PostgresBootstrapper { config: BootstrapConfig, admin_pool: OnceCell, + /// Dedicated connection whose open transaction holds the migration advisory + /// lock. `None` when no migration lock is held. + lock_conn: Mutex>, } impl PostgresBootstrapper { @@ -42,6 +64,7 @@ impl PostgresBootstrapper { Self { config, admin_pool: OnceCell::new(), + lock_conn: Mutex::new(None), } } @@ -209,6 +232,136 @@ impl Bootstrapper for PostgresBootstrapper { migrations::run_data_migrations(&pool).await } + async fn acquire_migration_lock(&self) -> OpResult<()> { + use std::io::Write as _; + + use sqlx::Connection; + + // This is not re-entrant. A second call would open a second transaction + // and block on the lock the first one holds, deadlocking against itself. + let mut held = self.lock_conn.lock().await; + if held.is_some() { + return Err(OpError::Internal( + "Migration lock is already held by this process".to_owned(), + )); + } + + // Keep an explicit transaction open on a dedicated catalog connection + // for the whole migration. Transaction-level advisory locks are released + // automatically when that transaction ends or the connection dies. + // + // The explicit transaction is important for transaction-pooling proxies: + // they must retain one PostgreSQL backend until COMMIT/ROLLBACK, so the + // lock cannot silently move between backends as separate autocommit + // statements can. RDS Proxy also supports this path without session + // pinning. The migration statements can use separate connections because + // advisory locks coordinate globally within the catalog database. + // + // Advisory locks are scoped to a database, so migrators only serialize + // if they share this catalog database. ExtendDB replicas normally do, + // because they use the same catalog connection string. + let mut conn = + sqlx::PgConnection::connect_with(&self.app_connect_opts(&self.config.catalog_db)) + .await + .map_err(|e| { + OpError::Internal(format!("Cannot connect to take migration lock: {e}")) + })?; + sqlx::query("BEGIN") + .execute(&mut conn) + .await + .map_err(|e| OpError::Internal(format!("Cannot begin migration lock: {e}")))?; + + // Equivalent to SET LOCAL lock_timeout, but parameterized so the Rust + // duration remains the single source of truth. This applies only to the + // dedicated transaction and is discarded by rollback on release. + let lock_timeout = format!("{}ms", MIGRATION_LOCK_TIMEOUT.as_millis()); + let _: String = sqlx::query_scalar("SELECT set_config('lock_timeout', $1, true)") + .bind(&lock_timeout) + .fetch_one(&mut conn) + .await + .map_err(|e| { + OpError::Internal(format!("Cannot configure migration lock timeout: {e}")) + })?; + + // Try first so that a migrator which has to wait can say so, rather + // than sitting silent for as long as the other migration takes. + let acquired: bool = sqlx::query_scalar("SELECT pg_try_advisory_xact_lock($1, $2)") + .bind(ADVISORY_LOCK_NAMESPACE) + .bind(MIGRATION_LOCK_OBJID) + .fetch_one(&mut conn) + .await + .map_err(|e| OpError::Internal(format!("Cannot acquire migration lock: {e}")))?; + if !acquired { + println!("--- Another migrator holds the migration lock; waiting for it to finish..."); + std::io::stdout() + .flush() + .map_err(|e| OpError::Internal(format!("Cannot report migration wait: {e}")))?; + sqlx::query("SELECT pg_advisory_xact_lock($1, $2)") + .bind(ADVISORY_LOCK_NAMESPACE) + .bind(MIGRATION_LOCK_OBJID) + .execute(&mut conn) + .await + .map_err(|e| { + if let sqlx::Error::Database(db_err) = &e + && db_err.code().as_deref() == Some("55P03") + { + return OpError::Internal(format!( + "Timed out after {}s waiting for the migration advisory lock; \ + another migrator may be wedged. Check the holder's logs or \ + terminate it before retrying: {e}", + MIGRATION_LOCK_TIMEOUT.as_secs(), + )); + } + OpError::Internal(format!("Cannot acquire migration lock: {e}")) + })?; + println!("--- Migration lock acquired."); + std::io::stdout() + .flush() + .map_err(|e| OpError::Internal(format!("Cannot report migration lock: {e}")))?; + } + + // Guard against a key or lock-mode mismatch in the acquisition queries. + // `objsubid = 2` distinguishes the two-i32 advisory-lock keyspace from + // the one-i64 form, and `granted` excludes a merely waiting request. + let held_by_this_transaction: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_locks \ + WHERE locktype = 'advisory' AND classid = $1 AND objid = $2 \ + AND objsubid = 2 AND mode = 'ExclusiveLock' AND granted \ + AND pid = pg_backend_pid())", + ) + .bind(ADVISORY_LOCK_NAMESPACE) + .bind(MIGRATION_LOCK_OBJID) + .fetch_one(&mut conn) + .await + .map_err(|e| OpError::Internal(format!("Cannot verify migration lock: {e}")))?; + if !held_by_this_transaction { + return Err(OpError::Internal( + "PostgreSQL did not report the migration transaction's advisory lock after \ + acquisition; refusing to run migrations without verified serialization" + .to_owned(), + )); + } + + *held = Some(conn); + Ok(()) + } + + async fn release_migration_lock(&self) -> OpResult<()> { + use sqlx::Connection; + + let Some(mut conn) = self.lock_conn.lock().await.take() else { + return Ok(()); + }; + // This transaction contains only the advisory lock, so rollback is the + // safest release: it cannot accidentally commit future work added to the + // dedicated connection. Closing is a backstop if rollback fails. + let release = sqlx::query("ROLLBACK").execute(&mut conn).await; + let _ = conn.close().await; + release + .map(|_| ()) + .map_err(|e| OpError::Internal(format!("Cannot release migration lock: {e}"))) + } + async fn pending_data_migrations(&self) -> OpResult> { let pool = self.app_pool(&self.config.data_db).await?; migrations::pending_data_migrations(&pool).await diff --git a/crates/storage-postgres/src/migrations.rs b/crates/storage-postgres/src/migrations.rs index fd390c32..84b48417 100755 --- a/crates/storage-postgres/src/migrations.rs +++ b/crates/storage-postgres/src/migrations.rs @@ -25,6 +25,14 @@ pub(crate) async fn run_catalog_migrations(pool: &PgPool) -> OpResult<()> { .execute(pool) .await .map_err(|e| OpError::Internal(format!("Migration {filename} failed: {e}")))?; + // TODO(#221): applying this SQL and recording it are separate commits. + // A crash here can leave a migration applied but unrecorded. Catalog 001 + // is normally shielded from replay by its version write, data 001 has an + // adoption guard, and data 002 is repeatable, but those are narrow + // recovery properties: catalog 001 is not idempotent and replaying data + // 003 drops the token table. The sqlx adoption must remove the files' + // own BEGIN/COMMIT and commit each ledger row with its migration before + // another migration lands. record_migration(pool, filename).await?; } println!(" Migrations applied."); @@ -91,6 +99,14 @@ pub(crate) async fn run_data_migrations(pool: &PgPool) -> OpResult<()> { .execute(pool) .await .map_err(|e| OpError::Internal(format!("Data migration {filename} failed: {e}")))?; + // TODO(#221): applying this SQL and recording it are separate commits. + // A crash here can leave a migration applied but unrecorded. Catalog 001 + // is normally shielded from replay by its version write, data 001 has an + // adoption guard, and data 002 is repeatable, but those are narrow + // recovery properties: catalog 001 is not idempotent and replaying data + // 003 drops the token table. The sqlx adoption must remove the files' + // own BEGIN/COMMIT and commit each ledger row with its migration before + // another migration lands. record_migration(pool, filename).await?; } println!(" Data migrations applied."); diff --git a/crates/storage/src/bootstrapper.rs b/crates/storage/src/bootstrapper.rs index 88299ebe..e0b1f8bf 100755 --- a/crates/storage/src/bootstrapper.rs +++ b/crates/storage/src/bootstrapper.rs @@ -66,6 +66,28 @@ pub trait Bootstrapper: Send + Sync { /// Run data schema migrations (stream tables, sequences, etc.). async fn run_data_migrations(&self) -> OpResult<()>; + /// Acquire a cross-process lock so that concurrent `migrate` runs, such as + /// several replicas starting at once, don't race each other applying schema + /// changes. Blocks until the lock is available. + /// + /// The default is a no-op, which means a backend that does not override this + /// gets no serialization at all. Implement it if concurrent bootstrap is a + /// supported deployment for your backend. + /// + /// Implementations need not be re-entrant, and callers must not acquire + /// twice without releasing. Every acquire must be paired with + /// [`Bootstrapper::release_migration_lock`], and the backend must also + /// release the lock if the process dies without doing so. + async fn acquire_migration_lock(&self) -> OpResult<()> { + Ok(()) + } + + /// Release the lock taken by [`Bootstrapper::acquire_migration_lock`], and + /// do nothing if no lock is held. Defaults to a no-op. + async fn release_migration_lock(&self) -> OpResult<()> { + Ok(()) + } + /// Filenames of data-database migrations that have not yet been applied. /// /// Excludes a pre-tracking baseline migration that already exists and will @@ -465,3 +487,105 @@ pub mod helpers { } } } + +/// Compile-time and behaviour guard for out-of-tree backends. +/// +/// `MinimalBootstrapper` implements only the methods the trait requires. If a +/// new required method is added, or a defaulted one loses its default, this +/// stops compiling, which is the same breakage an out-of-tree backend crate +/// would hit. It also pins object safety by boxing as `dyn Bootstrapper`, and +/// checks that the defaulted migration-lock methods are usable no-ops. +#[cfg(test)] +mod out_of_tree_compat_tests { + use super::*; + use crate::management_store::OpResult; + + struct MinimalBootstrapper; + + #[async_trait] + impl Bootstrapper for MinimalBootstrapper { + async fn ensure_app_user(&self) -> OpResult<()> { + Ok(()) + } + async fn grant_app_role_to_admin(&self) -> OpResult<()> { + Ok(()) + } + async fn create_catalog_db(&self) -> OpResult<()> { + Ok(()) + } + async fn create_data_db(&self) -> OpResult<()> { + Ok(()) + } + async fn run_catalog_migrations(&self) -> OpResult<()> { + Ok(()) + } + async fn run_data_migrations(&self) -> OpResult<()> { + Ok(()) + } + async fn pending_data_migrations(&self) -> OpResult> { + Ok(Vec::new()) + } + async fn record_data_connection(&self) -> OpResult<()> { + Ok(()) + } + async fn bootstrap_encryption_key(&self) -> OpResult<()> { + Ok(()) + } + async fn bootstrap_default_account(&self) -> OpResult<()> { + Ok(()) + } + async fn bootstrap_admin_user( + &self, + env_user: Option<&str>, + _env_password: Option<&str>, + ) -> OpResult { + Ok(AdminBootstrapResult { + username: env_user.unwrap_or("admin").to_owned(), + generated_password: None, + already_existed: false, + from_env: false, + }) + } + async fn is_catalog_initialized(&self) -> OpResult { + Ok(true) + } + async fn list_table_names(&self) -> OpResult> { + Ok(Vec::new()) + } + async fn get_data_db_name(&self) -> OpResult> { + Ok(None) + } + async fn drop_databases(&self, _data_db: &str) -> OpResult<()> { + Ok(()) + } + async fn read_catalog_version(&self) -> OpResult> { + Ok(None) + } + fn expected_catalog_version(&self) -> String { + "0.0.0".to_owned() + } + fn catalog_database_name(&self) -> String { + "minimal".to_owned() + } + fn endpoint_info(&self) -> String { + "in-memory".to_owned() + } + fn catalog_connection_url(&self) -> String { + "minimal://".to_owned() + } + fn generate_backend_config_section(&self) -> String { + String::new() + } + } + + /// A backend that does not implement the migration lock still works, and + /// the defaults are no-ops rather than errors. + #[tokio::test] + async fn migration_lock_defaults_to_a_usable_no_op() { + let bootstrapper: Box = Box::new(MinimalBootstrapper); + assert!(bootstrapper.acquire_migration_lock().await.is_ok()); + assert!(bootstrapper.release_migration_lock().await.is_ok()); + // Releasing without acquiring must also be harmless. + assert!(bootstrapper.release_migration_lock().await.is_ok()); + } +} diff --git a/devtools/run-tests b/devtools/run-tests index 7c9b8947..a32ea8de 100755 --- a/devtools/run-tests +++ b/devtools/run-tests @@ -443,7 +443,7 @@ if $RUN_PYTEST; then fi echo " AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID:-}" # CLI lifecycle tests manage their own server; exclude from main suite - PYTEST_ARGS=(python3 -m pytest tests/ -v --ignore=tests/python --ignore=tests/test_cli_lifecycle.py --ignore=tests/test_cli_container_readiness.py --ignore=tests/test_gsi_async_queue.py) + PYTEST_ARGS=(python3 -m pytest tests/ -v --ignore=tests/python --ignore=tests/test_cli_lifecycle.py --ignore=tests/test_cli_container_readiness.py --ignore=tests/test_cli_migrate_concurrency.py --ignore=tests/test_gsi_async_queue.py) # Parallel execution (--parallel): distribute by file so module/class-scoped # fixtures stay within a single worker. if [[ -n "$PARALLEL" ]] && python3 -c "import xdist" 2>/dev/null; then @@ -532,7 +532,7 @@ fi if $RUN_PYTEST && [[ "$TARGET" != "real-dynamodb" && -n "${EXTENDDB_TEST_PG_CONNECTION_STRING:-}" ]]; then CLI_OUTFILE="discussions/test-cli-${HASH}.txt" echo "=== CLI lifecycle tests → ${CLI_OUTFILE} ===" - CLI_ARGS=(python3 -m pytest tests/test_cli_lifecycle.py tests/test_cli_container_readiness.py -v) + CLI_ARGS=(python3 -m pytest tests/test_cli_lifecycle.py tests/test_cli_container_readiness.py tests/test_cli_migrate_concurrency.py -v) if [[ -n "$FILTER" ]]; then CLI_ARGS+=(-k "$FILTER") fi diff --git a/tests/test_cli_migrate_concurrency.py b/tests/test_cli_migrate_concurrency.py new file mode 100644 index 00000000..04a7f31c --- /dev/null +++ b/tests/test_cli_migrate_concurrency.py @@ -0,0 +1,393 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Migration concurrency CLI tests. + +Cover the PostgreSQL advisory lock that serializes `extenddb migrate`, so two +replicas starting at once cannot race each other applying schema changes. + +Like the other CLI lifecycle tests these require a PostgreSQL instance +(EXTENDDB_TEST_PG_CONNECTION_STRING) and a built binary, and are excluded from +the backend-agnostic pytest suite. +""" + +from __future__ import annotations + +import os +import select +import signal +import subprocess +import time + +from lifecycle_helpers import ( + EXTENDDB_BINARY, + PG_ADMIN_CONN, + _init_args, + _pg_args, + _run_extenddb, +) + +# Must match ADVISORY_LOCK_NAMESPACE and MIGRATION_LOCK_OBJID in +# crates/storage-postgres/src/bootstrapper.rs. +LOCK_NAMESPACE = 0x0045_4442 +MIGRATION_LOCK_OBJID = 1 + + +def _init(cli_env): + result = _run_extenddb( + "init", *_init_args(cli_env), + config=cli_env["config_path"], + env_override={"EXTENDDB_ADMIN_PASSWORD": "TestPass1!"}, + ) + assert result.returncode == 0, result.stdout + result.stderr + + +def _migrate_cmd(cli_env): + return [ + EXTENDDB_BINARY, "migrate", "--yes", *_pg_args(), + "--config", cli_env["config_path"], + ] + + +def _decode(data): + """Decode captured subprocess output for assertions and diagnostics.""" + return data.decode("utf-8", errors="replace") + + +def _read_until(proc, expected, timeout=15.0): + """Read raw stdout until expected appears, without waiting for process exit.""" + assert proc.stdout is not None + expected_bytes = expected.encode() + output = bytearray() + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + remaining = deadline - time.monotonic() + ready, _, _ = select.select([proc.stdout], [], [], min(0.1, remaining)) + if not ready: + if proc.poll() is not None: + output.extend(os.read(proc.stdout.fileno(), 4096)) + break + continue + chunk = os.read(proc.stdout.fileno(), 4096) + if not chunk: + break + output.extend(chunk) + if expected_bytes in output: + return _decode(output) + raise AssertionError( + f"timed out waiting for {expected!r}; output so far: {_decode(output)!r}" + ) + + +def _wait_for_blocked_ledger_insert(data_db, timeout=15.0): + """Wait until a migrator is blocked recording 002 in schema_history.""" + import psycopg2 + + observer = psycopg2.connect(PG_ADMIN_CONN + "/" + data_db) + observer.autocommit = True + try: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + with observer.cursor() as cur: + cur.execute( + """ + SELECT EXISTS( + SELECT 1 FROM pg_stat_activity + WHERE datname = %s + AND state = 'active' + AND wait_event_type = 'Lock' + AND query LIKE 'INSERT INTO schema_history%%' + ) + """, + (data_db,), + ) + if cur.fetchone()[0]: + return + time.sleep(0.05) + finally: + observer.close() + raise AssertionError("migrator did not block on the schema_history insert") + + +def _assert_002_applied_but_unrecorded(data_db): + """Prove the barrier is after migration commit and before ledger commit.""" + import psycopg2 + + observer = psycopg2.connect(PG_ADMIN_CONN + "/" + data_db) + observer.autocommit = True + try: + with observer.cursor() as cur: + cur.execute("SELECT to_regclass('public.gsi_pending') IS NOT NULL") + assert cur.fetchone()[0] is True, "002 schema should already be committed" + cur.execute( + "SELECT count(*) FROM schema_history " + "WHERE filename = '002_gsi_pending.sql'" + ) + assert cur.fetchone()[0] == 0, "002 ledger row should still be uncommitted" + finally: + observer.close() + + +def _terminate(proc): + """Best-effort cleanup for a subprocess after a failed assertion.""" + if proc is None or proc.poll() is not None: + return + proc.terminate() + try: + proc.communicate(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + proc.communicate(timeout=5) + + +class TestMigrateConcurrency: + """Concurrent `migrate` runs are serialized by an advisory lock.""" + + def test_concurrent_migrate_no_race(self, cli_env): + import psycopg2 + + _init(cli_env) + data_db = cli_env["db_name"][: -len("_catalog")] + + # Simulate a pre-002 deployment so migrate has real work to do. + blocker = psycopg2.connect(PG_ADMIN_CONN + "/" + data_db) + blocker.autocommit = True + with blocker.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS gsi_pending") + cur.execute( + "DELETE FROM schema_history WHERE filename = '002_gsi_pending.sql'" + ) + + # Insert the ledger row without committing it. A migrator's ordinary + # pending check cannot see this row, so it applies 002, then its own + # INSERT blocks on the uncommitted unique-key conflict. This gives the + # test a deterministic point after migration SQL and before recording. + blocker.autocommit = False + with blocker.cursor() as cur: + cur.execute( + "INSERT INTO schema_history (filename) VALUES ('002_gsi_pending.sql')" + ) + + cmd = _migrate_cmd(cli_env) + p1 = None + p2 = None + barrier_released = False + try: + p1 = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, bufsize=0, + ) + _wait_for_blocked_ledger_insert(data_db) + _assert_002_applied_but_unrecorded(data_db) + + # The first migrator now holds the migration lock at a known point. + # The second must stop at that lock rather than reach its pending + # check or ledger insert. Reading this line before releasing the + # blocker also proves that wait visibility is flushed immediately. + p2 = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, bufsize=0, + ) + p2_prefix = _read_until(p2, "waiting for it to finish") + assert p2.poll() is None, "second migrator exited instead of waiting" + + blocker.rollback() + barrier_released = True + + out1_raw = p1.communicate(timeout=60) + out2_tail_raw = p2.communicate(timeout=60) + out1 = (_decode(out1_raw[0]), _decode(out1_raw[1])) + out2 = ( + p2_prefix + _decode(out2_tail_raw[0]), + _decode(out2_tail_raw[1]), + ) + finally: + try: + if not barrier_released: + blocker.rollback() + finally: + try: + blocker.close() + finally: + try: + _terminate(p1) + finally: + _terminate(p2) + + assert p1.returncode == 0, out1 + assert p2.returncode == 0, out2 + + # Exactly one process applied 002. The second took its pending snapshot + # only after the first committed the ledger row and released the lock. + outputs = [out1[0], out2[0]] + applied = [o for o in outputs if "Applying 002_gsi_pending.sql" in o] + assert len(applied) == 1, f"expected exactly one migrator to apply 002: {outputs}" + others = [o for o in outputs if "Applying 002_gsi_pending.sql" not in o] + assert len(others) == 1, outputs + assert ( + "Everything is up to date" in others[0] + or "002_gsi_pending.sql — already applied" in others[0] + ), f"the second migrator should have observed 002 as done: {others[0]}" + assert "Migration lock acquired" in out2[0], out2[0] + + # The migration and its ledger entry each landed exactly once. + conn = psycopg2.connect(PG_ADMIN_CONN + "/" + data_db) + try: + with conn.cursor() as cur: + cur.execute("SELECT to_regclass('public.gsi_pending') IS NOT NULL") + assert cur.fetchone()[0] is True + cur.execute( + "SELECT count(*) FROM schema_history " + "WHERE filename = '002_gsi_pending.sql'" + ) + assert cur.fetchone()[0] == 1 + finally: + conn.close() + + def test_sigkill_releases_lock_and_waiting_peer_recovers(self, cli_env): + """A killed lock holder releases the lock and exposes the ledger gap.""" + import psycopg2 + + _init(cli_env) + data_db = cli_env["db_name"][: -len("_catalog")] + + blocker = psycopg2.connect(PG_ADMIN_CONN + "/" + data_db) + blocker.autocommit = True + with blocker.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS gsi_pending") + cur.execute( + "DELETE FROM schema_history WHERE filename = '002_gsi_pending.sql'" + ) + blocker.autocommit = False + with blocker.cursor() as cur: + cur.execute( + "INSERT INTO schema_history (filename) VALUES ('002_gsi_pending.sql')" + ) + + cmd = _migrate_cmd(cli_env) + holder = None + peer = None + barrier_released = False + try: + holder = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, bufsize=0, + ) + _wait_for_blocked_ledger_insert(data_db) + _assert_002_applied_but_unrecorded(data_db) + + peer = subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, bufsize=0, + ) + peer_waiting = _read_until(peer, "waiting for it to finish") + assert peer.poll() is None, "peer exited instead of waiting for holder" + + # SIGKILL gives the holder no chance to call release. PostgreSQL must + # roll back the dedicated transaction when the connection dies, + # which releases pg_advisory_xact_lock for the waiting peer. + holder.kill() + holder_out_raw = holder.communicate(timeout=10) + holder_out = (_decode(holder_out_raw[0]), _decode(holder_out_raw[1])) + assert holder.returncode == -signal.SIGKILL, holder_out + + peer_acquired = _read_until(peer, "Migration lock acquired") + assert peer.poll() is None, "peer should next block on the ledger barrier" + + blocker.rollback() + barrier_released = True + + peer_tail_raw = peer.communicate(timeout=60) + peer_out = ( + peer_waiting + peer_acquired + _decode(peer_tail_raw[0]), + _decode(peer_tail_raw[1]), + ) + finally: + try: + if not barrier_released: + blocker.rollback() + finally: + try: + blocker.close() + finally: + try: + _terminate(holder) + finally: + _terminate(peer) + + assert peer.returncode == 0, peer_out + assert "waiting for it to finish" in peer_out[0], peer_out[0] + assert "Migration lock acquired" in peer_out[0], peer_out[0] + + conn = psycopg2.connect(PG_ADMIN_CONN + "/" + data_db) + try: + with conn.cursor() as cur: + cur.execute("SELECT to_regclass('public.gsi_pending') IS NOT NULL") + assert cur.fetchone()[0] is True + cur.execute( + "SELECT count(*) FROM schema_history " + "WHERE filename = '002_gsi_pending.sql'" + ) + assert cur.fetchone()[0] == 1 + finally: + conn.close() + + def test_migrate_waits_for_a_held_lock(self, cli_env): + """migrate blocks on a lock held elsewhere, and says so immediately. + + The lock is held from an external session, so contention is guaranteed + instead of depending on process timing. + """ + import psycopg2 + + _init(cli_env) + + # Advisory locks are scoped to a database, and migrate takes this one on + # the catalog database. Session- and transaction-level requests for the + # same key conflict with each other. + conn = psycopg2.connect(PG_ADMIN_CONN + "/" + cli_env["db_name"]) + conn.autocommit = True + proc = None + lock_held = False + try: + with conn.cursor() as cur: + cur.execute( + "SELECT pg_advisory_lock(%s, %s)", + (LOCK_NAMESPACE, MIGRATION_LOCK_OBJID), + ) + lock_held = True + proc = subprocess.Popen( + _migrate_cmd(cli_env), stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdin=subprocess.DEVNULL, bufsize=0, + ) + # Observe the explanation while the external session still holds + # the lock. This also guards the explicit stdout flush in Rust. + prefix = _read_until(proc, "waiting for it to finish") + assert proc.poll() is None, "migrate should block while the lock is held" + + with conn.cursor() as cur: + cur.execute( + "SELECT pg_advisory_unlock(%s, %s)", + (LOCK_NAMESPACE, MIGRATION_LOCK_OBJID), + ) + lock_held = False + + out_tail, err = proc.communicate(timeout=60) + out = prefix + _decode(out_tail) + err = _decode(err) + assert proc.returncode == 0, (out, err) + assert "waiting for it to finish" in out, out + assert "Migration lock acquired" in out, out + finally: + try: + if lock_held: + with conn.cursor() as cur: + cur.execute( + "SELECT pg_advisory_unlock(%s, %s)", + (LOCK_NAMESPACE, MIGRATION_LOCK_OBJID), + ) + finally: + try: + conn.close() + finally: + _terminate(proc)