multica/server/cmd/migrate/migrate_concurrent_test.go

package main

import (
	"context"
	"fmt"
	"math/rand/v2"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"golang.org/x/sync/errgroup"
)

// MUL-2956 — concurrent migration race test.
//
// PR multica-ai/multica#3658 (MUL-2923) added a Postgres advisory lock
// around the migration loop to serialize concurrent runners. This file
// is the live-Postgres test that proves the lock is actually doing its
// job. We run N goroutines that all call runMigrations against the same
// database with the same options, and assert:
//
//  1. Pending: when migrations have NOT been applied, every goroutine
//     returns nil and exactly one application of each migration lands
//     in the bookkeeping table — no duplicate-key blow-ups, no missing
//     rows, and (since our test fixtures are deliberately non-idempotent
//     bare CREATE TABLE / ALTER TABLE) no "relation already exists"
//     failures from the SQL itself, which would prove the lock isn't
//     serializing.
//  2. Already applied: rerunning the same N-way race against the just-
//     populated bookkeeping table sends every goroutine down the EXISTS
//     no-op path; nobody re-applies anything and the underlying schema
//     is unchanged.
//  3. Lock serialization: while one connection holds the same advisory
//     lock externally, every concurrent runMigrations is observed to
//     wait, and only after the external holder releases does the lock
//     get acquired. This catches the regression where the lock would
//     get attached to a random pooled connection (the bug fixed in
//     MUL-2923 / #3658) and effectively become a no-op.
//
// The test connects to whatever DATABASE_URL points at (default
// postgres://multica:multica@localhost:5432/multica?sslmode=disable),
// matching the harness pattern already used in
// server/internal/handler/handler_test.go. If Postgres is unreachable the
// suite skips cleanly, the same way every other live-Postgres test in the
// repo skips, so CI without a database sees SKIP rather than failure.
//
// Each test isolates itself by creating a unique throwaway schema
// (migrate_test_<timestamp>_<rand>) and using a unique advisory-lock
// key per run. That means the test never touches the real
// schema_migrations table and never blocks behind a real production
// migration runner sharing the same database. The schema is dropped
// during cleanup.
//
// The same isolation is what lets the other database tests in this package
// run in parallel. What they still share is the database itself: a CREATE
// INDEX CONCURRENTLY waits out every older snapshot in it, so one test's
// build can wait on another test's open transaction — briefly, since none
// of them holds one for long, and never in a cycle, since a build only
// waits for snapshots older than its own. Extensions are the one
// database-wide object they create; createTestExtension serializes that.
//
// The four race tests below stay serial. Each fills a pool of one
// connection per CPU with its sixteen runners, and four of those beside
// the package's parallel tests would take a many-core machine to the
// server's connection limit.

const (
	defaultTestDatabaseURL = "postgres://multica:multica@localhost:5432/multica?sslmode=disable"
	// concurrentRunners is the goroutine count for the race tests. Set
	// large enough that a missing lock would reliably trip on a multi-
	// core box with -race, but small enough to keep the suite fast on a
	// single shared Postgres.
	concurrentRunners = 16
	// raceTestTimeout bounds every individual concurrent step; if the
	// lock implementation regresses into a deadlock we fail loudly
	// instead of hanging the suite.
	raceTestTimeout = 60 * time.Second
)

func testDatabaseURL() string {
	if dbURL := os.Getenv("DATABASE_URL"); dbURL != "" {
		return dbURL
	}
	return defaultTestDatabaseURL
}

func openTestPool(t *testing.T) *pgxpool.Pool {
	t.Helper()
	dbURL := testDatabaseURL()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	pool, err := pgxpool.New(ctx, dbURL)
	if err != nil {
		t.Skipf("could not connect to %s: %v", dbURL, err)
	}
	if err := pool.Ping(ctx); err != nil {
		pool.Close()
		t.Skipf("database not reachable at %s: %v", dbURL, err)
	}
	t.Cleanup(pool.Close)
	return pool
}

// fixture is the per-test sandbox: a private schema, a unique advisory
// lock key, and a temp directory full of deliberately non-idempotent
// migration SQL files.
type fixture struct {
	pool       *pgxpool.Pool
	schema     string
	tableFQN   string // e.g. "migrate_test_xyz"."schema_migrations"
	lockKey    int64
	files      []string // sorted .up.sql paths
	versions   []string // matching versions
	tableNames []string // distinct test tables each migration creates
}

func newFixture(t *testing.T) *fixture {
	t.Helper()
	pool := openTestPool(t)

	// Unique schema and lock key per test invocation. We salt with both
	// nanos and a process-local random so re-running with -count=N still
	// gets a distinct sandbox per iteration even if the wall clock has
	// not visibly advanced.
	suffix := fmt.Sprintf("%d_%d", time.Now().UnixNano(), rand.Uint32())
	schema := "migrate_test_" + suffix
	tableFQN := schema + ".schema_migrations"
	// Random non-zero positive int64. The high bit is masked off to
	// keep this in the same numeric range pg_advisory_lock expects, and
	// the OR with 1 guarantees we never end up at zero. Collision with
	// the production migrationAdvisoryLockKey constant is not strictly
	// impossible — both are int64 — but the probability is on the order
	// of 1 in 2^62, which is negligible for a unit-test sandbox.
	lockKey := int64(rand.Uint64()&0x7fffffffffffffff) | 1

	ctx := context.Background()
	if _, err := pool.Exec(ctx, fmt.Sprintf(`CREATE SCHEMA %s`, pgx.Identifier{schema}.Sanitize())); err != nil {
		t.Fatalf("create schema: %v", err)
	}
	t.Cleanup(func() {
		// Use a fresh context so cleanup still runs if the test ctx was
		// cancelled. Drop CASCADE to take everything down even when a
		// half-applied migration left orphan tables behind.
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		if _, err := pool.Exec(ctx, fmt.Sprintf(`DROP SCHEMA IF EXISTS %s CASCADE`, pgx.Identifier{schema}.Sanitize())); err != nil {
			t.Logf("drop schema %s: %v", schema, err)
		}
	})

	// Build a small set of deliberately non-idempotent migrations. Each
	// one creates a distinct table inside our schema. Bare CREATE TABLE
	// (no IF NOT EXISTS) and ALTER TABLE ADD COLUMN (no IF NOT EXISTS)
	// guarantee that if two goroutines actually ran the same migration
	// in parallel, the second one would error with "relation already
	// exists" / "column already exists" — which is exactly the failure
	// signature we want the test to catch when the lock regresses.
	dir := t.TempDir()
	const numFiles = 5
	files := make([]string, 0, numFiles)
	versions := make([]string, 0, numFiles)
	tableNames := make([]string, 0, numFiles)
	for i := 0; i < numFiles; i++ {
		version := fmt.Sprintf("%03d_test_%s", i+1, suffix)
		tableName := fmt.Sprintf("t_%s_%d", suffix, i+1)
		// Reference both the schema and the table, then add a column in
		// a follow-up statement. Either statement run twice (i.e.
		// concurrent re-application by another goroutine that won the
		// race past the EXISTS check) would error.
		body := fmt.Sprintf(
			"CREATE TABLE %s.%s (id BIGSERIAL PRIMARY KEY);\n"+
				"ALTER TABLE %s.%s ADD COLUMN payload TEXT NOT NULL DEFAULT '';\n",
			pgx.Identifier{schema}.Sanitize(), pgx.Identifier{tableName}.Sanitize(),
			pgx.Identifier{schema}.Sanitize(), pgx.Identifier{tableName}.Sanitize(),
		)
		path := filepath.Join(dir, version+".up.sql")
		if err := os.WriteFile(path, []byte(body), 0o600); err != nil {
			t.Fatalf("write migration: %v", err)
		}
		files = append(files, path)
		versions = append(versions, version)
		tableNames = append(tableNames, tableName)
	}
	sort.Strings(files)

	return &fixture{
		pool:       pool,
		schema:     schema,
		tableFQN:   tableFQN,
		lockKey:    lockKey,
		files:      files,
		versions:   versions,
		tableNames: tableNames,
	}
}

func (f *fixture) opts() runOptions {
	return runOptions{
		Direction:             "up",
		Files:                 f.files,
		SchemaMigrationsTable: f.tableFQN,
		AdvisoryLockKey:       f.lockKey,
	}
}

// appliedVersions returns the versions recorded in the bookkeeping
// table, sorted ascending. Empty slice means the table is empty (or
// does not yet exist, which the helper reports as a fatal error).
func (f *fixture) appliedVersions(t *testing.T) []string {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	rows, err := f.pool.Query(ctx,
		fmt.Sprintf(`SELECT version FROM %s ORDER BY version`,
			pgx.Identifier{f.schema, "schema_migrations"}.Sanitize()))
	if err != nil {
		t.Fatalf("read schema_migrations: %v", err)
	}
	defer rows.Close()
	var got []string
	for rows.Next() {
		var v string
		if err := rows.Scan(&v); err != nil {
			t.Fatalf("scan version: %v", err)
		}
		got = append(got, v)
	}
	if err := rows.Err(); err != nil {
		t.Fatalf("rows.Err: %v", err)
	}
	return got
}

// tableExists checks that a given table is present inside the fixture
// schema. Used to confirm migrations actually executed, not just that
// the bookkeeping rows landed.
func (f *fixture) tableExists(t *testing.T, name string) bool {
	t.Helper()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	var exists bool
	if err := f.pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.tables
			WHERE table_schema = $1 AND table_name = $2
		)`, f.schema, name).Scan(&exists); err != nil {
		t.Fatalf("check table %s.%s: %v", f.schema, name, err)
	}
	return exists
}

// TestRunMigrationsConcurrentPending fires N goroutines at runMigrations
// against a fresh schema where none of the migrations have been applied
// yet. The advisory lock must serialize them so that exactly one of
// them executes each CREATE TABLE / ALTER TABLE and exactly one row
// per migration lands in schema_migrations. If the lock is broken,
// either the SQL fails ("relation already exists") or the bookkeeping
// table picks up duplicate-key violations on the version primary key.
func TestRunMigrationsConcurrentPending(t *testing.T) {
	f := newFixture(t)

	ctx, cancel := context.WithTimeout(context.Background(), raceTestTimeout)
	defer cancel()

	g, gctx := errgroup.WithContext(ctx)
	// errgroup limits concurrency only when SetLimit is called; we want
	// every goroutine running at once so they all queue on the lock.
	for i := 0; i < concurrentRunners; i++ {
		g.Go(func() error {
			return runMigrations(gctx, f.pool, f.opts())
		})
	}
	if err := g.Wait(); err != nil {
		t.Fatalf("concurrent runMigrations(up) on pending schema returned error: %v", err)
	}

	got := f.appliedVersions(t)
	if want := f.versions; !equalStrings(got, want) {
		t.Fatalf("schema_migrations contents = %v, want %v", got, want)
	}
	for _, tbl := range f.tableNames {
		if !f.tableExists(t, tbl) {
			t.Fatalf("expected table %s.%s to exist after concurrent up, missing", f.schema, tbl)
		}
	}
}

// TestRunMigrationsConcurrentAlreadyApplied first applies the
// migrations once (single-threaded, to establish a clean baseline) and
// then fires N goroutines at runMigrations again. Every goroutine must
// hit the EXISTS no-op path and return nil, the bookkeeping table must
// stay exactly the way the baseline left it, and the underlying tables
// must not have been touched (no duplicate CREATE / ALTER blow-ups).
//
// This is the path that matters in production: most pod restarts find
// the database fully migrated and just need to confirm-and-skip.
func TestRunMigrationsConcurrentAlreadyApplied(t *testing.T) {
	f := newFixture(t)

	ctx, cancel := context.WithTimeout(context.Background(), raceTestTimeout)
	defer cancel()

	// Baseline single-threaded run.
	if err := runMigrations(ctx, f.pool, f.opts()); err != nil {
		t.Fatalf("baseline runMigrations: %v", err)
	}
	baseline := f.appliedVersions(t)
	if !equalStrings(baseline, f.versions) {
		t.Fatalf("baseline schema_migrations = %v, want %v", baseline, f.versions)
	}

	// Concurrent re-run: every goroutine should hit the EXISTS no-op
	// branch and return nil.
	g, gctx := errgroup.WithContext(ctx)
	for i := 0; i < concurrentRunners; i++ {
		g.Go(func() error {
			return runMigrations(gctx, f.pool, f.opts())
		})
	}
	if err := g.Wait(); err != nil {
		t.Fatalf("concurrent runMigrations(up) on already-applied schema returned error: %v", err)
	}

	got := f.appliedVersions(t)
	if !equalStrings(got, baseline) {
		t.Fatalf("schema_migrations changed after concurrent re-run: got %v, want %v", got, baseline)
	}
}

// TestRunMigrationsAdvisoryLockSerializes proves the lock genuinely
// blocks contenders. We acquire the same advisory key on a side
// connection BEFORE spawning any runMigrations goroutine, then start N
// goroutines and watch how many of them have made it past the lock
// acquire. The expectation:
//
//   - While the side connection holds the lock, zero goroutines have
//     completed (we observe a blocked contender in pg_locks).
//   - The moment the side connection releases the lock, the goroutines
//     start unblocking and finish in well under the test timeout.
//
// If the advisory lock had regressed back to attaching to a random
// pooled connection (the original MUL-2923 bug), the side-held lock
// would not actually block a fresh pool.Acquire from grabbing its own
// connection without the lock, and the goroutines would all complete
// while the lock was still "held" — which is exactly what this test
// detects.
func TestRunMigrationsAdvisoryLockSerializes(t *testing.T) {
	f := newFixture(t)

	ctx, cancel := context.WithTimeout(context.Background(), raceTestTimeout)
	defer cancel()

	// Acquire the lock on a pinned side connection. We use a *pgx.Conn
	// (not pool.Acquire) so the lock holder is not reachable through
	// the pool the runMigrations goroutines draw from — the lock is
	// session-scoped and we want the behaviour to be "the next pool
	// connection that calls pg_advisory_lock blocks", not "the same
	// connection re-enters". (pg_advisory_lock is reentrant on the same
	// session, so re-acquiring on the same conn would not actually
	// prove serialization.)
	dbURL := testDatabaseURL()
	holder, err := pgx.Connect(ctx, dbURL)
	if err != nil {
		t.Fatalf("side connect: %v", err)
	}
	defer holder.Close(context.Background())
	if _, err := holder.Exec(ctx, "SELECT pg_advisory_lock($1)", f.lockKey); err != nil {
		t.Fatalf("side acquire lock: %v", err)
	}

	var done int64
	var startedAt = time.Now()
	g, gctx := errgroup.WithContext(ctx)
	for i := 0; i < concurrentRunners; i++ {
		g.Go(func() error {
			err := runMigrations(gctx, f.pool, f.opts())
			atomic.AddInt64(&done, 1)
			return err
		})
	}

	// Wait until PostgreSQL confirms that a contender is blocked on this exact
	// advisory lock. This proves serialization directly without making the
	// passing path pay a fixed observation delay.
	const observeWindow = 2 * time.Second
	const observeStep = 10 * time.Millisecond
	deadline := time.Now().Add(observeWindow)
	blocked := false
	for time.Now().Before(deadline) {
		if n := atomic.LoadInt64(&done); n != 0 {
			t.Fatalf("advisory lock did not block: %d/%d goroutines finished while side connection held the lock for %s",
				n, concurrentRunners, time.Since(startedAt))
		}
		if err := holder.QueryRow(ctx, `
			SELECT EXISTS (
				SELECT 1
				FROM pg_locks AS held
				JOIN pg_locks AS waiter
				  ON waiter.locktype = held.locktype
				 AND waiter.database IS NOT DISTINCT FROM held.database
				 AND waiter.classid IS NOT DISTINCT FROM held.classid
				 AND waiter.objid IS NOT DISTINCT FROM held.objid
				 AND waiter.objsubid IS NOT DISTINCT FROM held.objsubid
				 AND waiter.mode = held.mode
				WHERE held.pid = pg_backend_pid()
				  AND held.locktype = 'advisory'
				  AND held.granted
				  AND NOT waiter.granted
			)
		`).Scan(&blocked); err != nil {
			t.Fatalf("observe advisory-lock waiter: %v", err)
		}
		if blocked {
			break
		}
		time.Sleep(observeStep)
	}
	if !blocked {
		t.Fatal("no runMigrations contender reached the held advisory lock")
	}

	// Release the side lock and wait for all goroutines to finish.
	if _, err := holder.Exec(ctx, "SELECT pg_advisory_unlock($1)", f.lockKey); err != nil {
		t.Fatalf("side release lock: %v", err)
	}
	if err := g.Wait(); err != nil {
		t.Fatalf("concurrent runMigrations after lock release returned error: %v", err)
	}

	// Sanity: state ended up as the pending case did — exactly one
	// application of every migration.
	if got, want := f.appliedVersions(t), f.versions; !equalStrings(got, want) {
		t.Fatalf("schema_migrations after lock-release race = %v, want %v", got, want)
	}
}

// TestRunMigrationsConcurrentMixedPoolStress runs the pending case
// against a deliberately under-sized pool to put pressure on the
// "every runner needs its own pinned connection for the lock" code
// path. If runMigrations ever regresses into using pool.Exec (which
// could give the lock and the migration steps different connections),
// this test will deadlock or produce SQL races. Pool size strictly
// less than runners is the interesting configuration.
func TestRunMigrationsConcurrentMixedPoolStress(t *testing.T) {
	dbURL := testDatabaseURL()
	cfg, err := pgxpool.ParseConfig(dbURL)
	if err != nil {
		t.Skipf("parse DATABASE_URL: %v", err)
	}
	// Small pool: half the runner count, minimum 2. This forces
	// runners to wait on pgxpool.Acquire AND on pg_advisory_lock,
	// exercising the same connection lifecycle a real multi-replica
	// startup would.
	cfg.MaxConns = int32(concurrentRunners / 2)
	if cfg.MaxConns < 2 {
		cfg.MaxConns = 2
	}
	ctx, cancel := context.WithTimeout(context.Background(), raceTestTimeout)
	defer cancel()
	pool, err := pgxpool.NewWithConfig(ctx, cfg)
	if err != nil {
		t.Skipf("could not open small pool: %v", err)
	}
	defer pool.Close()
	if err := pool.Ping(ctx); err != nil {
		t.Skipf("small pool not reachable: %v", err)
	}

	// Reuse the standard fixture's schema/files/lock-key wiring but
	// substitute the small-pool *pgxpool.Pool so the test exercises a
	// different connection budget.
	big := newFixture(t)
	f := *big
	f.pool = pool

	g, gctx := errgroup.WithContext(ctx)
	var startedOnce sync.Once
	startedAt := time.Time{}
	for i := 0; i < concurrentRunners; i++ {
		g.Go(func() error {
			startedOnce.Do(func() { startedAt = time.Now() })
			return runMigrations(gctx, f.pool, f.opts())
		})
	}
	if err := g.Wait(); err != nil {
		t.Fatalf("small-pool concurrent runMigrations error after %s: %v", time.Since(startedAt), err)
	}
	if got, want := big.appliedVersions(t), big.versions; !equalStrings(got, want) {
		t.Fatalf("small-pool schema_migrations = %v, want %v", got, want)
	}
}

func equalStrings(a, b []string) bool {
	if len(a) != len(b) {
		return false
	}
	for i := range a {
		if a[i] != b[i] {
			return false
		}
	}
	return true
}

// TestRunMigrationsRejectsInvalidDirection pins the direction
// whitelist contract: anything other than "up" or "down" must error
// before runMigrations touches the pool. This prevents the subtle bug
// where an empty or typo'd direction silently fell through to the
// "down" branch (`opts.Direction == "up"` is false → else branch
// handles it as a rollback).
//
// The check runs ahead of any pool/conn use, so passing nil is safe
// and lets this case execute without a live Postgres.
func TestRunMigrationsRejectsInvalidDirection(t *testing.T) {
	t.Parallel()
	bad := []string{"", "UP", "DOWN", "rollback", "x", " up "}
	for _, dir := range bad {
		err := runMigrations(context.Background(), nil, runOptions{Direction: dir})
		if err == nil {
			t.Errorf("direction %q: want error, got nil", dir)
			continue
		}
		if !strings.Contains(err.Error(), "invalid direction") {
			t.Errorf("direction %q: error %q does not mention 'invalid direction'", dir, err)
		}
	}
}