used uber atomic bool instead standard in lock/unlock db (#580)

* used uber atomic bool instead standard in lock/unlock db * fixing local lock optimization implementation * fixing mysql local lock optimization implementation * fixing cockroachdb local lock optimization implementation * CAS wrapper to automatically restore * added test for cas with restore * added atomic lock to sqlite & mongo * added atomic lock to sqlite & mongo * fix/refactor test cas restore function * disable lock for mongo * mongo local lock before acquiring the db lock Co-authored-by: a.prinkov <[email protected]>
golang-migrate · Jun 23, 2021 · 3dfb0ff · 3dfb0ff
1 parent 3b3c1b6
commit 3dfb0ff
Show file tree

Hide file tree

Showing 19 changed files with 347 additions and 299 deletions.
diff --git a/database/cassandra/cassandra.go b/database/cassandra/cassandra.go
@@ -3,6 +3,7 @@ package cassandra
 import (
 	"errors"
 	"fmt"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -45,7 +46,7 @@ type Config struct {
 
 type Cassandra struct {
 	session  *gocql.Session
-	isLocked bool
+	isLocked atomic.Bool
 
 	// Open and WithInstance need to guarantee that config is never nil
 	config *Config
@@ -182,15 +183,16 @@ func (c *Cassandra) Close() error {
 }
 
 func (c *Cassandra) Lock() error {
-	if c.isLocked {
+	if !c.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
-	c.isLocked = true
 	return nil
 }
 
 func (c *Cassandra) Unlock() error {
-	c.isLocked = false
+	if !c.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }
 

diff --git a/database/clickhouse/clickhouse.go b/database/clickhouse/clickhouse.go
@@ -285,7 +285,7 @@ func (ch *ClickHouse) Lock() error {
 }
 func (ch *ClickHouse) Unlock() error {
 	if !ch.isLocked.CAS(true, false) {
-		return database.ErrLocked
+		return database.ErrNotLocked
 	}
 
 	return nil

diff --git a/database/cockroachdb/cockroachdb.go b/database/cockroachdb/cockroachdb.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -46,7 +47,7 @@ type Config struct {
 
 type CockroachDb struct {
 	db       *sql.DB
-	isLocked bool
+	isLocked atomic.Bool
 
 	// Open and WithInstance need to guarantee that config is never nil
 	config *Config
@@ -152,71 +153,67 @@ 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 {
-	err := crdb.ExecuteTx(context.Background(), c.db, nil, func(tx *sql.Tx) (err error) {
-		aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
-		if err != nil {
-			return err
-		}
-
-		query := "SELECT * FROM " + c.config.LockTable + " WHERE lock_id = $1"
-		rows, err := tx.Query(query, aid)
-		if err != nil {
-			return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)}
-		}
-		defer func() {
-			if errClose := rows.Close(); errClose != nil {
-				err = multierror.Append(err, errClose)
+	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)
+			if err != nil {
+				return err
 			}
-		}()
 
-		// If row exists at all, lock is present
-		locked := rows.Next()
-		if locked && !c.config.ForceLock {
-			return database.ErrLocked
-		}
+			query := "SELECT * FROM " + c.config.LockTable + " WHERE lock_id = $1"
+			rows, err := tx.Query(query, aid)
+			if err != nil {
+				return database.Error{OrigErr: err, Err: "failed to fetch migration lock", Query: []byte(query)}
+			}
+			defer func() {
+				if errClose := rows.Close(); errClose != nil {
+					err = multierror.Append(err, errClose)
+				}
+			}()
+
+			// If row exists at all, lock is present
+			locked := rows.Next()
+			if locked && !c.config.ForceLock {
+				return database.ErrLocked
+			}
 
-		query = "INSERT INTO " + c.config.LockTable + " (lock_id) VALUES ($1)"
-		if _, err := tx.Exec(query, aid); err != nil {
-			return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)}
-		}
+			query = "INSERT INTO " + c.config.LockTable + " (lock_id) VALUES ($1)"
+			if _, err := tx.Exec(query, aid); err != nil {
+				return database.Error{OrigErr: err, Err: "failed to set migration lock", Query: []byte(query)}
+			}
 
-		return nil
+			return nil
+		})
 	})
-
-	if err != nil {
-		return err
-	} else {
-		c.isLocked = true
-		return nil
-	}
 }
 
 // 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) Unlock() error {
-	aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
-	if err != nil {
-		return err
-	}
+	return database.CasRestoreOnErr(&c.isLocked, true, false, database.ErrNotLocked, func() (err error) {
+		aid, err := database.GenerateAdvisoryLockId(c.config.DatabaseName)
+		if err != nil {
+			return err
+		}
 
-	// In the event of an implementation (non-migration) error, it is possible for the lock to not be released.  Until
-	// a better locking mechanism is added, a manual purging of the lock table may be required in such circumstances
-	query := "DELETE FROM " + c.config.LockTable + " WHERE lock_id = $1"
-	if _, err := c.db.Exec(query, aid); err != nil {
-		if e, ok := err.(*pq.Error); ok {
-			// 42P01 is "UndefinedTableError" in CockroachDB
-			// https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/pgerror/codes.go
-			if e.Code == "42P01" {
-				// On drops, the lock table is fully removed;  This is fine, and is a valid "unlocked" state for the schema
-				c.isLocked = false
-				return nil
+		// In the event of an implementation (non-migration) error, it is possible for the lock to not be released.  Until
+		// a better locking mechanism is added, a manual purging of the lock table may be required in such circumstances
+		query := "DELETE FROM " + c.config.LockTable + " WHERE lock_id = $1"
+		if _, err := c.db.Exec(query, aid); err != nil {
+			if e, ok := err.(*pq.Error); ok {
+				// 42P01 is "UndefinedTableError" in CockroachDB
+				// https://github.com/cockroachdb/cockroach/blob/master/pkg/sql/pgwire/pgerror/codes.go
+				if e.Code == "42P01" {
+					// On drops, the lock table is fully removed;  This is fine, and is a valid "unlocked" state for the schema
+					return nil
+				}
 			}
+
+			return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
 		}
-		return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
-	}
 
-	c.isLocked = false
-	return nil
+		return nil
+	})
 }
 
 func (c *CockroachDb) Run(migration io.Reader) error {

diff --git a/database/firebird/firebird.go b/database/firebird/firebird.go
@@ -10,6 +10,7 @@ import (
 	"github.com/golang-migrate/migrate/v4/database"
 	"github.com/hashicorp/go-multierror"
 	_ "github.com/nakagami/firebirdsql"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -36,7 +37,7 @@ type Firebird struct {
 	// Locking and unlocking need to use the same connection
 	conn     *sql.Conn
 	db       *sql.DB
-	isLocked bool
+	isLocked atomic.Bool
 
 	// Open and WithInstance need to guarantee that config is never nil
 	config *Config
@@ -106,15 +107,16 @@ func (f *Firebird) Close() error {
 }
 
 func (f *Firebird) Lock() error {
-	if f.isLocked {
+	if !f.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
-	f.isLocked = true
 	return nil
 }
 
 func (f *Firebird) Unlock() error {
-	f.isLocked = false
+	if !f.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }
 

diff --git a/database/mongodb/mongodb.go b/database/mongodb/mongodb.go
@@ -10,6 +10,7 @@ import (
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
 	"go.mongodb.org/mongo-driver/x/mongo/driver/connstring"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	"net/url"
@@ -40,9 +41,10 @@ var (
 )
 
 type Mongo struct {
-	client *mongo.Client
-	db     *mongo.Database
-	config *Config
+	client   *mongo.Client
+	db       *mongo.Database
+	config   *Config
+	isLocked atomic.Bool
 }
 
 type Locking struct {
@@ -327,55 +329,60 @@ func (m *Mongo) ensureVersionTable() (err error) {
 // Utilizes advisory locking on the config.LockingCollection collection
 // This uses a unique index on the `locking_key` field.
 func (m *Mongo) Lock() error {
-	if !m.config.Locking.Enabled {
-		return nil
-	}
-	pid := os.Getpid()
-	hostname, err := os.Hostname()
-	if err != nil {
-		hostname = fmt.Sprintf("Could not determine hostname. Error: %s", err.Error())
-	}
+	return database.CasRestoreOnErr(&m.isLocked, false, true, database.ErrLocked, func() error {
+		if !m.config.Locking.Enabled {
+			return nil
+		}
 
-	newLockObj := lockObj{
-		Key:       lockKeyUniqueValue,
-		Pid:       pid,
-		Hostname:  hostname,
-		CreatedAt: time.Now(),
-	}
-	operation := func() error {
-		timeout, cancelFunc := context.WithTimeout(context.Background(), contextWaitTimeout)
-		_, err := m.db.Collection(m.config.Locking.CollectionName).InsertOne(timeout, newLockObj)
-		defer cancelFunc()
-		return err
-	}
-	exponentialBackOff := backoff.NewExponentialBackOff()
-	duration := time.Duration(m.config.Locking.Timeout) * time.Second
-	exponentialBackOff.MaxElapsedTime = duration
-	exponentialBackOff.MaxInterval = time.Duration(m.config.Locking.Interval) * time.Second
+		pid := os.Getpid()
+		hostname, err := os.Hostname()
+		if err != nil {
+			hostname = fmt.Sprintf("Could not determine hostname. Error: %s", err.Error())
+		}
 
-	err = backoff.Retry(operation, exponentialBackOff)
-	if err != nil {
-		return database.ErrLocked
-	}
+		newLockObj := lockObj{
+			Key:       lockKeyUniqueValue,
+			Pid:       pid,
+			Hostname:  hostname,
+			CreatedAt: time.Now(),
+		}
+		operation := func() error {
+			timeout, cancelFunc := context.WithTimeout(context.Background(), contextWaitTimeout)
+			_, err := m.db.Collection(m.config.Locking.CollectionName).InsertOne(timeout, newLockObj)
+			defer cancelFunc()
+			return err
+		}
+		exponentialBackOff := backoff.NewExponentialBackOff()
+		duration := time.Duration(m.config.Locking.Timeout) * time.Second
+		exponentialBackOff.MaxElapsedTime = duration
+		exponentialBackOff.MaxInterval = time.Duration(m.config.Locking.Interval) * time.Second
 
-	return nil
+		err = backoff.Retry(operation, exponentialBackOff)
+		if err != nil {
+			return database.ErrLocked
+		}
 
+		return nil
+	})
 }
+
 func (m *Mongo) Unlock() error {
-	if !m.config.Locking.Enabled {
-		return nil
-	}
+	return database.CasRestoreOnErr(&m.isLocked, true, false, database.ErrNotLocked, func() error {
+		if !m.config.Locking.Enabled {
+			return nil
+		}
 
-	filter := findFilter{
-		Key: lockKeyUniqueValue,
-	}
+		filter := findFilter{
+			Key: lockKeyUniqueValue,
+		}
 
-	ctx, cancel := context.WithTimeout(context.Background(), contextWaitTimeout)
-	_, err := m.db.Collection(m.config.Locking.CollectionName).DeleteMany(ctx, filter)
-	defer cancel()
+		ctx, cancel := context.WithTimeout(context.Background(), contextWaitTimeout)
+		_, err := m.db.Collection(m.config.Locking.CollectionName).DeleteMany(ctx, filter)
+		defer cancel()
 
-	if err != nil {
-		return err
-	}
-	return nil
+		if err != nil {
+			return err
+		}
+		return nil
+	})
 }
diff --git a/database/mongodb/mongodb_test.go b/database/mongodb/mongodb_test.go
@@ -221,18 +221,7 @@ func TestLockWorks(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		// disable locking, validate wer can lock twice
-		mc.config.Locking.Enabled = false
-		err = mc.Lock()
-		if err != nil {
-			t.Fatal(err)
-		}
-		err = mc.Lock()
-		if err != nil {
-			t.Fatal(err)
-		}
-
-		// re-enable locking,
+		// enable locking,
 		//try to hit a lock conflict
 		mc.config.Locking.Enabled = true
 		mc.config.Locking.Timeout = 1