Fix migration race condition in concurrent container startups

cmd/server/main.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"database/sql"
 	"log"
 	"net/http"
 	"os"
@@ -294,38 +295,45 @@ func runMigrations(db *database.DB) error {
 			continue
 		}
 
-		log.Printf("Applying migration: %s", entry.Name())
-
-		// Read migration file
-		migrationSQL, err := os.ReadFile("migrations/" + entry.Name())
-		if err != nil {
-			return err
-		}
-
-		// Start transaction for this migration
-		tx, err := db.Begin()
-		if err != nil {
+		// Try to claim this migration using INSERT ... ON CONFLICT
+		// This ensures only one instance will successfully claim and execute the migration
+		var claimedVersion string
+		err = db.QueryRow(
+			`INSERT INTO schema_migrations (version, applied_at) 
+			 VALUES ($1, CURRENT_TIMESTAMP) 
+			 ON CONFLICT (version) DO NOTHING 
+			 RETURNING version`,
+			entry.Name(),
+		).Scan(&claimedVersion)
+
+		if err == sql.ErrNoRows {
+			// ON CONFLICT happened - another instance already claimed this migration
+			// This is expected in concurrent scenarios, not an error
+			log.Printf("Migration %s already claimed by another instance, skipping", entry.Name())
+			continue
+		} else if err != nil {
+			// Unexpected database error
 			return err
 		}
 
-		// Execute migration
-		if _, err := tx.Exec(string(migrationSQL)); err != nil {
-			tx.Rollback()
-			return err
-		}
+		log.Printf("Applying migration: %s", entry.Name())
 
-		// Record migration as applied
-		_, err = tx.Exec(
-			"INSERT INTO schema_migrations (version) VALUES ($1)",
-			entry.Name(),
-		)
+		// Read migration file
+		migrationSQL, err := os.ReadFile("migrations/" + entry.Name())
 		if err != nil {
-			tx.Rollback()
+			// Migration was claimed but can't be read - attempt to remove the claim
+			if _, delErr := db.Exec("DELETE FROM schema_migrations WHERE version = $1", entry.Name()); delErr != nil {
+				log.Printf("Warning: failed to remove claim for %s after read error: %v", entry.Name(), delErr)
+			}
 			return err
 		}
 
-		// Commit transaction
-		if err := tx.Commit(); err != nil {
+		// Execute migration DDL
+		// Note: We don't use a transaction here because we've already inserted the tracking row
+		// If the migration fails, the tracking row remains as a record that it was attempted
+		if _, err := db.Exec(string(migrationSQL)); err != nil {
+			log.Printf("✗ Failed to apply migration %s: %v", entry.Name(), err)
+			// Leave the tracking row to prevent re-attempts; manual intervention required
 			return err
 		}
 

cmd/server/migrations_test.go

@@ -0,0 +1,122 @@
+package main
+
+import (
+	"database/sql"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+
+	_ "github.com/lib/pq"
+	"github.com/techbuzzz/agent-shaker/internal/database"
+)
+
+// TestMigrationConcurrentSafety verifies that concurrent migration attempts
+// don't cause conflicts or duplicate executions
+func TestMigrationConcurrentSafety(t *testing.T) {
+	// Skip if no test database is available
+	dbURL := os.Getenv("TEST_DATABASE_URL")
+	if dbURL == "" {
+		t.Skip("TEST_DATABASE_URL not set, skipping integration test")
+	}
+
+	// Clean up any existing test schema
+	cleanupDB(t, dbURL)
+
+	// Create temporary test migration file
+	tmpDir := t.TempDir()
+	migrationFile := filepath.Join(tmpDir, "001_test.sql")
+	migrationSQL := `CREATE TABLE IF NOT EXISTS test_table (id SERIAL PRIMARY KEY, name TEXT);`
+	if err := os.WriteFile(migrationFile, []byte(migrationSQL), 0644); err != nil {
+		t.Fatalf("Failed to create test migration: %v", err)
+	}
+
+	// Change to temp directory for migration discovery
+	oldWd, _ := os.Getwd()
+	migrationsDir := filepath.Join(tmpDir, "migrations")
+	if err := os.Mkdir(migrationsDir, 0755); err != nil {
+		t.Fatalf("Failed to create migrations dir: %v", err)
+	}
+	if err := os.WriteFile(filepath.Join(migrationsDir, "001_test.sql"), []byte(migrationSQL), 0644); err != nil {
+		t.Fatalf("Failed to create migration in migrations dir: %v", err)
+	}
+	defer os.Chdir(oldWd)
+	os.Chdir(tmpDir)
+
+	// Test concurrent migration attempts
+	const numConcurrent = 5
+	var wg sync.WaitGroup
+	errors := make(chan error, numConcurrent)
+
+	for i := 0; i < numConcurrent; i++ {
+		wg.Add(1)
+		go func(id int) {
+			defer wg.Done()
+
+			// Each goroutine creates its own DB connection
+			db, err := database.NewDB(dbURL)
+			if err != nil {
+				errors <- err
+				return
+			}
+			defer db.Close()
+
+			// Attempt to run migrations
+			if err := runMigrations(db); err != nil {
+				errors <- err
+			}
+		}(i)
+	}
+
+	wg.Wait()
+	close(errors)
+
+	// Check for any errors
+	for err := range errors {
+		t.Errorf("Migration error: %v", err)
+	}
+
+	// Verify migration was applied exactly once
+	db, err := database.NewDB(dbURL)
+	if err != nil {
+		t.Fatalf("Failed to connect to verify: %v", err)
+	}
+	defer db.Close()
+
+	// Check schema_migrations table
+	var count int
+	err = db.QueryRow("SELECT COUNT(*) FROM schema_migrations WHERE version = '001_test.sql'").Scan(&count)
+	if err != nil {
+		t.Fatalf("Failed to query migrations: %v", err)
+	}
+	if count != 1 {
+		t.Errorf("Expected migration to be recorded exactly once, got %d times", count)
+	}
+
+	// Verify the test table was created
+	var exists bool
+	err = db.QueryRow(`
+		SELECT EXISTS (
+			SELECT FROM information_schema.tables 
+			WHERE table_name = 'test_table'
+		)
+	`).Scan(&exists)
+	if err != nil {
+		t.Fatalf("Failed to check table existence: %v", err)
+	}
+	if !exists {
+		t.Error("Expected test_table to exist after migration")
+	}
+}
+
+func cleanupDB(t *testing.T, dbURL string) {
+	db, err := sql.Open("postgres", dbURL)
+	if err != nil {
+		t.Fatalf("Failed to open DB for cleanup: %v", err)
+	}
+	defer db.Close()
+
+	// Drop test tables
+	_, _ = db.Exec("DROP TABLE IF EXISTS test_table CASCADE")
+	_, _ = db.Exec("DROP TABLE IF EXISTS schema_migrations CASCADE")
+}

docs/MIGRATION_CONCURRENCY.md

@@ -0,0 +1,84 @@
+# Migration Concurrency Safety
+
+## Problem
+
+The original migration system had a race condition when multiple instances started simultaneously:
+
+1. **Instance A** reads `schema_migrations`, sees migration X is not applied
+2. **Instance B** reads `schema_migrations`, sees migration X is not applied
+3. **Instance A** executes migration X DDL successfully
+4. **Instance B** executes migration X DDL (may succeed if idempotent)
+5. **Instance A** inserts `version=X` into `schema_migrations` → SUCCESS
+6. **Instance B** tries to insert `version=X` into `schema_migrations` → PRIMARY KEY CONFLICT → CRASH
+
+Even though the migration was successfully applied by Instance A, Instance B would crash, preventing the application from starting.
+
+## Solution
+
+The implementation now uses PostgreSQL's `INSERT ... ON CONFLICT DO NOTHING RETURNING` pattern to atomically claim migrations:
+
+```go
+var claimedVersion string
+err = db.QueryRow(
+    `INSERT INTO schema_migrations (version, applied_at) 
+     VALUES ($1, CURRENT_TIMESTAMP) 
+     ON CONFLICT (version) DO NOTHING 
+     RETURNING version`,
+    entry.Name(),
+).Scan(&claimedVersion)
+
+if err != nil {
+    // sql.ErrNoRows means ON CONFLICT happened - another instance claimed this
+    log.Printf("Migration %s already claimed by another instance, skipping", entry.Name())
+    continue
+}
+
+// Only the instance that successfully claimed the migration reaches here
+// and executes the DDL
+```
+
+### How It Works
+
+1. Each instance attempts to insert a row into `schema_migrations` for the migration
+2. The `ON CONFLICT DO NOTHING` clause prevents errors if the row already exists
+3. The `RETURNING version` clause returns the version only if the INSERT succeeded
+4. If another instance already claimed the migration:
+   - `ON CONFLICT` prevents the insert
+   - No rows are returned
+   - `Scan()` returns `sql.ErrNoRows`
+   - The instance safely skips this migration
+5. Only one instance will successfully claim and execute each migration
+
+### Benefits
+
+- **No race conditions**: The database enforces atomicity
+- **No crashes**: Concurrent instances gracefully skip already-claimed migrations
+- **Simple**: No need for advisory locks or external coordination
+- **Docker-safe**: Works correctly with `docker-compose up --scale app=3`
+- **Production-safe**: Multiple pods/containers can start simultaneously
+
+### Trade-offs
+
+- Migrations are no longer wrapped in a single transaction
+- The tracking row is inserted *before* executing the DDL
+- If a migration fails:
+  - The tracking row remains (prevents re-attempts)
+  - Manual intervention is required to fix and continue
+  - This is intentional: failed migrations should not auto-retry
+
+### Testing Concurrent Safety
+
+Run the integration test with a test database:
+
+```bash
+export TEST_DATABASE_URL="postgresql://user:pass@localhost/testdb"
+go test ./cmd/server -v -run TestMigrationConcurrentSafety
+```
+
+Or test manually with Docker Compose:
+
+```bash
+docker-compose up --scale app=3 -d
+docker-compose logs app
+# All three instances should start without migration conflicts
+```