From d466d27a4f3d60c83c7c02e403454a9b5a926063 Mon Sep 17 00:00:00 2001 From: Scott Robinson Date: Thu, 6 Aug 2026 10:03:48 +0000 Subject: [PATCH] fix(migrate): Serialize concurrent migrations Two replicas running `extenddb migrate` can both evaluate pending work before either records it, then apply the same migration concurrently. PostgreSQL can reject the duplicate DDL in its system catalogs, making an otherwise idempotent container startup unsafe. Serialize the complete migration decision and apply sequence with a namespaced PostgreSQL transaction-level advisory lock. The lock is held by an explicit transaction on a dedicated catalog connection, so it covers the version read and data-migration pending check, catalog and data migrations and final ledger observation. A peer acquires the lock only after the holder finishes, then observes no pending work. - Try `pg_try_advisory_xact_lock` first so a contending process can explain why it is waiting, then use `pg_advisory_xact_lock` for the blocking acquisition. - Keep the explicit transaction open for the full migration interval. Transaction-pooling proxies retain one backend until rollback. This avoids session locks split across pooled backend sessions. Direct PostgreSQL, RDS, Aurora, PgBouncer transaction pooling, and RDS Proxy follow the same transaction-scoped lock contract. - Bound the wait with a five-minute transaction-local `lock_timeout`. PostgreSQL SQLSTATE 55P03 becomes an actionable error naming the migration advisory lock and explaining that the holder may be wedged. - Flush waiting and acquired messages so container logs show progress while the process is still blocked rather than only after it exits. - Verify the granted row in `pg_locks`, including the two-integer key subtype, ExclusiveLock mode, and granted state. A key or lock-mode mismatch fails closed before migrations run. - Release by rolling back the dedicated transaction, then close the connection as a backstop. PostgreSQL also rolls the transaction back and releases it when the process or backend dies without cleanup. - Reject a second in-process acquire instead of deadlocking on a new transaction waiting for the lock already held by the first. - Scope the lock to the catalog database. Replicas sharing that catalog serialize even when migration statements touch the data database; deployments using different catalogs remain independent. - Lock init's schema phase too, so init and migrate cannot race. Two fresh init processes still do not queue. The second fails earlier while creating the already-existing catalog database. - Add default no-op lock methods to `Bootstrapper`. Existing external backends preserve current behavior until they opt into a backend lock. A MinimalBootstrapper test pins defaults and trait object safety. The advisory lock deliberately does not claim to make one migration atomic with its ledger write. The custom runner still applies SQL and records the filename in separate commits. Both apply sites carry a TODO for #221, which removes the migration files' internal BEGIN/COMMIT and lets sqlx commit each migration with its ledger row. Current recovery properties are narrower: catalog 001 is normally shielded by its version write, data 001 has an adoption guard, data 002 is repeatable, and replaying data 003 drops the token table. This gap must close before another migration relies on atomic recording. Add deterministic PostgreSQL CLI coverage: - Hold an uncommitted `schema_history` row so the first migrator commits 002's DDL and blocks exactly before its ledger insert. Start a second migrator and prove it waits. Release the barrier and verify both processes succeed, exactly one applies 002, and one ledger row exists. - Assert `gsi_pending` is committed while the 002 ledger row remains invisible, pinning the known apply-versus-ledger gap. - SIGKILL the holder without release, observe the waiting peer acquire the lock, then release the ledger barrier and verify recovery. This proves lock release on connection death. - Hold the lock from an external session and prove migrate explains the wait before the holder releases it, then completes successfully. Make subprocess reads unbuffered and cleanup exception-safe so tests cannot miss Python-buffered data or leak children after failure. Explicitly flush the Rust messages that the live tests observe. Integrate main, including container-readiness and SQLite changes. Resolve the CLI runner overlap as a union: lifecycle, container-readiness, and migration-concurrency files are excluded from the shared pytest run and executed together in the dedicated PostgreSQL CLI section. Verification on the final tree: - fmt and clippy `-D warnings` clean - release build clean - 669 workspace tests pass, with 3 ignored - 28 live PostgreSQL CLI tests pass, with 1 Unix-socket skip - all three migration-concurrency tests pass on PostgreSQL 16.10 - a one-second timeout mutation emits the exact diagnostic in 2.21s; the committed value is restored to five minutes and release rebuilt --- crates/app/src/cmd_init.rs | 70 ++-- crates/app/src/cmd_migrate.rs | 19 + crates/storage-postgres/src/bootstrapper.rs | 155 +++++++- crates/storage-postgres/src/migrations.rs | 16 + crates/storage/src/bootstrapper.rs | 124 ++++++ devtools/run-tests | 4 +- tests/test_cli_migrate_concurrency.py | 393 ++++++++++++++++++++ 7 files changed, 755 insertions(+), 26 deletions(-) create mode 100644 tests/test_cli_migrate_concurrency.py 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)