feat: an option for crdb to wait util the lock is released to mimic pg locks behaviour

database/cockroachdb/README.md

@@ -2,18 +2,21 @@
 
 `cockroachdb://user:password@host:port/dbname?query` (`cockroach://`, and `crdb-postgres://` work, too)
 
-| URL Query  | WithInstance Config | Description |
-|------------|---------------------|-------------|
-| `x-migrations-table` | `MigrationsTable` | Name of the migrations table |
-| `x-lock-table` | `LockTable` | Name of the table which maintains the migration lock |
-| `x-force-lock` | `ForceLock` | Force lock acquisition to fix faulty migrations which may not have released the schema lock (Boolean, default is `false`) |
-| `dbname` | `DatabaseName` | The name of the database to connect to |
-| `user` | | The user to sign in as |
-| `password` | | The user's password |
-| `host` | | The host to connect to. Values that start with / are for unix domain sockets. (default is localhost) |
-| `port` | | The port to bind to. (default is 5432) |
-| `connect_timeout` | | Maximum wait for connection, in seconds. Zero or not specified means wait indefinitely. |
-| `sslcert` | | Cert file location. The file must contain PEM encoded data. |
-| `sslkey` | | Key file location. The file must contain PEM encoded data. |
-| `sslrootcert` | | The location of the root certificate file. The file must contain PEM encoded data. |
-| `sslmode` | | Whether or not to use SSL (disable\|require\|verify-ca\|verify-full) |
+| URL Query                  | WithInstance Config   | Description                                                                                                                                                    |
+|----------------------------|-----------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `x-migrations-table`       | `MigrationsTable`     | Name of the migrations table                                                                                                                                   |
+| `x-lock-table`             | `LockTable`           | Name of the table which maintains the migration lock                                                                                                           |
+| `x-force-lock`             | `ForceLock`           | Force lock acquisition to fix faulty migrations which may not have released the schema lock (Boolean, default is `false`)                                      |
+| `x-max-retries`            | `MaxRetries`          | How many times retry queries on retryable errors (40001, 40P01, 08006, XX000). Default is 0 to keep the backward compatibility with the initial implementation |
+| `x-max-retry-interval`     | `MaxRetryInterval`    | Interval between retries increases exponentially. This option specifies maximum duration between retries. Default is `15s`                                     |
+| `x-max-retry-elapsed-time` | `MaxRetryElapsedTime` | Total retries timeout. Default is `30s`                                                                                                                        |
+| `dbname`                   | `DatabaseName`        | The name of the database to connect to                                                                                                                         |
+| `user`                     |                       | The user to sign in as                                                                                                                                         |
+| `password`                 |                       | The user's password                                                                                                                                            |
+| `host`                     |                       | The host to connect to. Values that start with / are for unix domain sockets. (default is localhost)                                                           |
+| `port`                     |                       | The port to bind to. (default is 5432)                                                                                                                         |
+| `connect_timeout`          |                       | Maximum wait for connection, in seconds. Zero or not specified means wait indefinitely.                                                                        |
+| `sslcert`                  |                       | Cert file location. The file must contain PEM encoded data.                                                                                                    |
+| `sslkey`                   |                       | Key file location. The file must contain PEM encoded data.                                                                                                     |
+| `sslrootcert`              |                       | The location of the root certificate file. The file must contain PEM encoded data.                                                                             |
+| `sslmode`                  |                       | Whether or not to use SSL (disable\|require\|verify-ca\|verify-full)                                                                                           |

database/cockroachdb/cockroachdb.go

@@ -3,18 +3,22 @@ package cockroachdb
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
 	"io"
 	nurl "net/url"
 	"regexp"
 	"strconv"
+	"time"
 
+	"github.com/cenkalti/backoff/v4"
 	"github.com/cockroachdb/cockroach-go/v2/crdb"
-	"github.com/golang-migrate/migrate/v4"
-	"github.com/golang-migrate/migrate/v4/database"
 	"github.com/hashicorp/go-multierror"
 	"github.com/lib/pq"
 	"go.uber.org/atomic"
+
+	"github.com/golang-migrate/migrate/v4"
+	"github.com/golang-migrate/migrate/v4/database"
 )
 
 func init() {
@@ -24,6 +28,11 @@ func init() {
 	database.Register("crdb-postgres", &db)
 }
 
+const (
+	DefaultMaxRetryInterval    = time.Second * 15
+	DefaultMaxRetryElapsedTime = time.Second * 30
+)
+
 var DefaultMigrationsTable = "schema_migrations"
 var DefaultLockTable = "schema_lock"
 
@@ -37,6 +46,10 @@ type Config struct {
 	LockTable       string
 	ForceLock       bool
 	DatabaseName    string
+
+	MaxRetryInterval    time.Duration
+	MaxRetryElapsedTime time.Duration
+	MaxRetries          int
 }
 
 type CockroachDb struct {
@@ -78,6 +91,14 @@ func WithInstance(instance *sql.DB, config *Config) (database.Driver, error) {
 		config.LockTable = DefaultLockTable
 	}
 
+	if config.MaxRetryInterval == 0 {
+		config.MaxRetryInterval = DefaultMaxRetryInterval
+	}
+
+	if config.MaxRetryElapsedTime == 0 {
+		config.MaxRetryElapsedTime = DefaultMaxRetryElapsedTime
+	}
+
 	px := &CockroachDb{
 		db:     instance,
 		config: config,
@@ -127,11 +148,33 @@ func (c *CockroachDb) Open(url string) (database.Driver, error) {
 		forceLock = false
 	}
 
+	maxIntervalStr := purl.Query().Get("x-max-retry-interval")
+	maxInterval, err := time.ParseDuration(maxIntervalStr)
+	if err != nil {
+		maxInterval = DefaultMaxRetryInterval
+	}
+
+	maxElapsedTimeStr := purl.Query().Get("x-max-retry-elapsed-time")
+	maxElapsedTime, err := time.ParseDuration(maxElapsedTimeStr)
+	if err != nil {
+		maxElapsedTime = DefaultMaxRetryElapsedTime
+	}
+
+	maxRetriesStr := purl.Query().Get("x-max-retries")
+	maxRetries, err := strconv.Atoi(maxRetriesStr)
+	if err != nil {
+		maxRetries = 0
+	}
+
 	px, err := WithInstance(db, &Config{
 		DatabaseName:    purl.Path,
 		MigrationsTable: migrationsTable,
 		LockTable:       lockTable,
 		ForceLock:       forceLock,
+
+		MaxRetryInterval:    maxInterval,
+		MaxRetryElapsedTime: maxElapsedTime,
+		MaxRetries:          maxRetries,
 	})
 	if err != nil {
 		return nil, err
@@ -147,6 +190,19 @@ func (c *CockroachDb) Close() error {
 // Locking is done manually with a separate lock table.  Implementing advisory locks in CRDB is being discussed
 // See: https://github.com/cockroachdb/cockroach/issues/13546
 func (c *CockroachDb) Lock() error {
+	// CRDB is using SERIALIZABLE isolation level by default, that means we can not run a loop inside the transaction,
+	// because transaction started when the lock is acquired does not see it being released
+	return backoff.Retry(func() error {
+		err := c.lock()
+		if err != nil && !errors.Is(err, database.ErrLocked) {
+			return backoff.Permanent(err)
+		}
+
+		return err
+	}, c.newBackoff())
+}
+
+func (c *CockroachDb) lock() error {
 	return database.CasRestoreOnErr(&c.isLocked, false, true, database.ErrLocked, func() (err error) {
 		return crdb.ExecuteTx(context.Background(), c.db, nil, func(tx *sql.Tx) (err error) {
 			aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
@@ -249,11 +305,12 @@ func (c *CockroachDb) Version() (version int, dirty bool, err error) {
 	err = c.db.QueryRow(query).Scan(&version, &dirty)
 
 	switch {
-	case err == sql.ErrNoRows:
+	case errors.Is(err, sql.ErrNoRows):
 		return database.NilVersion, false, nil
 
 	case err != nil:
-		if e, ok := err.(*pq.Error); ok {
+		var e *pq.Error
+		if errors.As(err, &e) {
 			// 42P01 is "UndefinedTableError" in CockroachDB
 			// https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/pgerror/codes.go
 			if e.Code == "42P01" {
@@ -363,3 +420,19 @@ func (c *CockroachDb) ensureLockTable() error {
 
 	return nil
 }
+
+func (c *CockroachDb) newBackoff() backoff.BackOff {
+	retrier := backoff.WithMaxRetries(backoff.WithContext(&backoff.ExponentialBackOff{
+		InitialInterval:     backoff.DefaultInitialInterval,
+		RandomizationFactor: backoff.DefaultRandomizationFactor,
+		Multiplier:          backoff.DefaultMultiplier,
+		MaxInterval:         c.config.MaxRetryInterval,
+		MaxElapsedTime:      c.config.MaxRetryElapsedTime,
+		Stop:                backoff.Stop,
+		Clock:               backoff.SystemClock,
+	}, context.Background()), uint64(c.config.MaxRetries))
+
+	retrier.Reset()
+
+	return retrier
+}

database/cockroachdb/cockroachdb_test.go

@@ -5,19 +5,18 @@ package cockroachdb
 import (
 	"context"
 	"database/sql"
+	"errors"
 	"fmt"
-	"github.com/golang-migrate/migrate/v4"
 	"log"
 	"strings"
 	"testing"
-)
+	"time"
 
-import (
 	"github.com/dhui/dktest"
 	_ "github.com/lib/pq"
-)
 
-import (
+	"github.com/golang-migrate/migrate/v4"
+	"github.com/golang-migrate/migrate/v4/database"
 	dt "github.com/golang-migrate/migrate/v4/database/testing"
 	"github.com/golang-migrate/migrate/v4/dktesting"
 	_ "github.com/golang-migrate/migrate/v4/source/file"
@@ -172,3 +171,102 @@ func TestFilterCustomQuery(t *testing.T) {
 		}
 	})
 }
+
+func TestLockDefault(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) {
+		createDB(t, ci)
+
+		ip, port, err := ci.Port(26257)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := fmt.Sprintf("cockroach://root@%v:%v/migrate?sslmode=disable", ip, port)
+		c := &CockroachDb{}
+		d1, err := c.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		d2, err := c.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if err := d1.Lock(); err != nil {
+			t.Fatal(err)
+		}
+
+		if err := d2.Lock(); err == nil || !errors.Is(err, database.ErrLocked) {
+			t.Fatalf("expected error %v, got %v", database.ErrLocked, err)
+		}
+
+		if err := d1.Unlock(); err != nil {
+			t.Fatal(err)
+		}
+
+		if err := d2.Lock(); err != nil {
+			t.Fatal(err)
+		}
+
+		if err := d2.Unlock(); err != nil {
+			t.Fatal(err)
+		}
+	})
+}
+
+func TestLockWait(t *testing.T) {
+	dktesting.ParallelTest(t, specs, func(t *testing.T, ci dktest.ContainerInfo) {
+		createDB(t, ci)
+
+		ip, port, err := ci.Port(26257)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		addr := fmt.Sprintf("cockroach://root@%v:%v/migrate?sslmode=disable&x-max-retries=10", ip, port)
+		c := &CockroachDb{}
+		d1, err := c.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		d2, err := c.Open(addr)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if err := d1.Lock(); err != nil {
+			t.Fatal(err)
+		}
+
+		// lock d2 in a goroutine to make it wait a bit
+		done := make(chan struct{})
+		waiting := make(chan struct{})
+		var d2LockErr error
+		go func() {
+			defer close(done)
+
+			close(waiting) // signal that the goroutine started and is waiting
+
+			// we can not call t.Fatal(err) in the goroutine, so remember it and check it in the main thread
+			d2LockErr = d2.Lock()
+		}()
+
+		<-waiting // ensure the goroutine started and is waiting
+
+		if err := d1.Unlock(); err != nil {
+			t.Fatal(err)
+		}
+
+		select {
+		case <-done: // we should get here once the d2 lock is acquired
+			if d2LockErr != nil {
+				t.Fatal(d2LockErr)
+			}
+			break
+		case <-time.After(DefaultMaxRetryInterval): // wait for at least one poll
+			t.Fatal("expected lock to be acquired by d2")
+		}
+	})
+}

dktesting/dktesting.go

@@ -4,9 +4,7 @@ import (
 	"context"
 	"fmt"
 	"testing"
-)
 
-import (
 	"github.com/dhui/dktest"
 	"github.com/docker/docker/api/types"
 	"github.com/docker/docker/client"
@@ -53,11 +51,12 @@ func ParallelTest(t *testing.T, specs []ContainerSpec,
 		// TODO: order is random, maybe always pick first version instead?
 		if i > 0 && testing.Short() {
 			t.Logf("Skipping %v in short mode", spec.ImageName)
-		} else {
-			t.Run(spec.ImageName, func(t *testing.T) {
-				t.Parallel()
-				dktest.Run(t, spec.ImageName, spec.Options, testFunc)
-			})
+			continue
 		}
+
+		t.Run(spec.ImageName, func(t *testing.T) {
+			t.Parallel()
+			dktest.Run(t, spec.ImageName, spec.Options, testFunc)
+		})
 	}
 }