use try advisory lock

database/postgres/README.md

@@ -9,6 +9,7 @@
 | `x-statement-timeout` | `StatementTimeout` | Abort any statement that takes more than the specified number of milliseconds |
 | `x-multi-statement` | `MultiStatementEnabled` | Enable multi-statement execution (default: false) |
 | `x-multi-statement-max-size` | `MultiStatementMaxSize` | Maximum size of single statement in bytes (default: 10MB) |
+| `x-lock-retry-max-interval` | `Locking` | When acquiring a lock fails, retries are used with an exponential backoff. This parameter specifies what is the maximum interval between retries in milliseconds (default: 1000ms). Values below the [initial retry interval](./postgres.go#L39) (100ms) will be ignored. |
 | `dbname` | `DatabaseName` | The name of the database to connect to |
 | `search_path` | | This variable specifies the order in which schemas are searched when an object is referenced by a simple name with no schema specified. |
 | `user` | | The user to sign in as |

database/postgres/postgres.go

@@ -15,6 +15,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/golang-migrate/migrate/v4"
 	"github.com/golang-migrate/migrate/v4/database"
 	"github.com/golang-migrate/migrate/v4/database/multistmt"
@@ -32,6 +33,9 @@ var (
 
 	DefaultMigrationsTable       = "schema_migrations"
 	DefaultMultiStatementMaxSize = 10 * 1 << 20 // 10 MB
+
+	DefaultLockInitialRetryInterval = 100 * time.Millisecond
+	DefaultLockMaxRetryInterval     = 1000 * time.Millisecond
 )
 
 var (
@@ -51,6 +55,17 @@ type Config struct {
 	migrationsTableName   string
 	StatementTimeout      time.Duration
 	MultiStatementMaxSize int
+	Locking               LockConfig
+}
+
+type LockConfig struct {
+	// InitialRetryInterval the initial (minimum) retry interval used for exponential backoff
+	// to try acquire a lock
+	InitialRetryInterval time.Duration
+
+	// MaxRetryInterval the maximum retry interval. Once the exponential backoff reaches this limit,
+	// the retry interval remains the same
+	MaxRetryInterval time.Duration
 }
 
 type Postgres struct {
@@ -200,13 +215,29 @@ func (p *Postgres) Open(url string) (database.Driver, error) {
 		}
 	}
 
+	lockConfig := LockConfig{
+		InitialRetryInterval: DefaultLockInitialRetryInterval,
+		MaxRetryInterval:     DefaultLockMaxRetryInterval,
+	}
+	if s := purl.Query().Get("x-lock-retry-max-interval"); len(s) > 0 {
+		maxRetryIntervalMillis, err := strconv.Atoi(s)
+		if err != nil {
+			return nil, fmt.Errorf("unable to parse option x-lock-retry-max-interval: %w", err)
+		}
+		maxRetryInterval := time.Duration(maxRetryIntervalMillis) * time.Millisecond
+		if maxRetryInterval > DefaultLockInitialRetryInterval {
+			lockConfig.MaxRetryInterval = maxRetryInterval
+		}
+	}
+
 	px, err := WithInstance(db, &Config{
 		DatabaseName:          purl.Path,
 		MigrationsTable:       migrationsTable,
 		MigrationsTableQuoted: migrationsTableQuoted,
 		StatementTimeout:      time.Duration(statementTimeout) * time.Millisecond,
 		MultiStatementEnabled: multiStatementEnabled,
 		MultiStatementMaxSize: multiStatementMaxSize,
+		Locking:               lockConfig,
 	})
 
 	if err != nil {
@@ -229,24 +260,45 @@ func (p *Postgres) Close() error {
 	return nil
 }
 
+// Lock tries to acquire an advisory lock and retries indefinitely with an exponential backoff strategy
 // https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
 func (p *Postgres) Lock() error {
 	return database.CasRestoreOnErr(&p.isLocked, false, true, database.ErrLocked, func() error {
-		aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
-		if err != nil {
-			return err
-		}
+		backOff := p.config.Locking.nonStopBackoff()
+		err := backoff.Retry(func() error {
+			ok, err := p.tryLock()
+			if err != nil {
+				return fmt.Errorf("p.tryLock: %w", err)
+			}
 
-		// This will wait indefinitely until the lock can be acquired.
-		query := `SELECT pg_advisory_lock($1)`
-		if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
-			return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
-		}
+			if ok {
+				return nil
+			}
 
-		return nil
+			return fmt.Errorf("could not acquire lock") // causes retry
+		}, backOff)
+
+		return err
 	})
 }
 
+func (p *Postgres) tryLock() (bool, error) {
+	aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
+	if err != nil {
+		return false, err
+	}
+
+	// https://www.postgresql.org/docs/current/functions-admin.html#FUNCTIONS-ADVISORY-LOCKS
+	// should always return true or false
+	query := `SELECT pg_try_advisory_lock($1)`
+	var ok bool
+	if err := p.conn.QueryRowContext(context.Background(), query, aid).Scan(&ok); err != nil {
+		return false, &database.Error{OrigErr: err, Err: "pg_try_advisory_lock failed", Query: []byte(query)}
+	}
+
+	return ok, nil
+}
+
 func (p *Postgres) Unlock() error {
 	return database.CasRestoreOnErr(&p.isLocked, true, false, database.ErrNotLocked, func() error {
 		aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
@@ -485,3 +537,13 @@ func (p *Postgres) ensureVersionTable() (err error) {
 
 	return nil
 }
+
+func (l *LockConfig) nonStopBackoff() backoff.BackOff {
+	b := backoff.NewExponentialBackOff()
+	b.InitialInterval = l.InitialRetryInterval
+	b.MaxInterval = l.MaxRetryInterval
+	b.MaxElapsedTime = 0 // this backoff won't stop
+	b.Reset()
+
+	return b
+}

database/postgres/postgres_test.go

@@ -96,6 +96,7 @@ func Test(t *testing.T) {
 	t.Run("testFailToCreateTableWithoutPermissions", testFailToCreateTableWithoutPermissions)
 	t.Run("testCheckBeforeCreateTable", testCheckBeforeCreateTable)
 	t.Run("testParallelSchema", testParallelSchema)
+	t.Run("testPostgresConcurrentMigrations", testPostgresConcurrentMigrations)
 	t.Run("testPostgresLock", testPostgresLock)
 	t.Run("testWithInstanceConcurrent", testWithInstanceConcurrent)
 	t.Run("testWithConnection", testWithConnection)
@@ -628,6 +629,49 @@ func testParallelSchema(t *testing.T) {
 	})
 }
 
+func testPostgresConcurrentMigrations(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
+		// GIVEN - a set of concurrent processes running migrations
+		const concurrency = 3
+		var wg sync.WaitGroup
+
+		ip, port, err := c.FirstPort()
+		if err != nil {
+			t.Fatal(err)
+		}
+		addr := pgConnectionString(ip, port, "x-lock-retry-max-interval=2000")
+
+		// WHEN
+		for i := 0; i < concurrency; i++ {
+			wg.Add(1)
+
+			go func() {
+				defer wg.Done()
+
+				p := &Postgres{}
+				d, err := p.Open(addr)
+				if err != nil {
+					t.Error(err)
+				}
+				defer func() {
+					if err := d.Close(); err != nil {
+						t.Error(err)
+					}
+				}()
+
+				m, err := migrate.NewWithDatabaseInstance("file://./examples/migrations", "postgres", d)
+				if err != nil {
+					t.Error(err)
+				}
+				dt.TestMigrate(t, m)
+			}()
+		}
+
+		wg.Wait()
+		// THEN
+	})
+}
+
 func testPostgresLock(t *testing.T) {
 	dktesting.ParallelTest(t, specs, func(t *testing.T, c dktest.ContainerInfo) {
 		ip, port, err := c.FirstPort()

go.mod

@@ -8,7 +8,7 @@ require (
 	github.com/Azure/go-autorest/autorest/adal v0.9.16
 	github.com/ClickHouse/clickhouse-go v1.4.3
 	github.com/aws/aws-sdk-go v1.49.6
-	github.com/cenkalti/backoff/v4 v4.1.2
+	github.com/cenkalti/backoff/v4 v4.3.0
 	github.com/cockroachdb/cockroach-go/v2 v2.1.1
 	github.com/dhui/dktest v0.4.6
 	github.com/docker/docker v28.3.3+incompatible

go.sum

@@ -731,8 +731,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/boombuler/barcode v1.0.0/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
 github.com/boombuler/barcode v1.0.1/go.mod h1:paBWMcWSl3LHKBqUq+rly7CNSldXjb2rDl3JlRe0mD8=
-github.com/cenkalti/backoff/v4 v4.1.2 h1:6Yo7N8UP2K6LWZnW94DLVSSrbobcWdVzAYOisuDPIFo=
-github.com/cenkalti/backoff/v4 v4.1.2/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
+github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
+github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/census-instrumentation/opencensus-proto v0.3.0/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
 github.com/census-instrumentation/opencensus-proto v0.4.1/go.mod h1:4T9NM4+4Vw91VeyqjLS6ao50K5bOcLKN6Q42XnYaRYw=