used uber atomic bool instead standard in lock/unlock db

golang-migrate · Jun 12, 2021 · 9814da9 · 9814da9
1 parent af256da
commit 9814da9
Show file tree

Hide file tree

Showing 14 changed files with 102 additions and 69 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
@@ -186,7 +187,9 @@ func (c *CockroachDb) Lock() error {
 	if err != nil {
 		return err
 	} else {
-		c.isLocked = true
+		if !c.isLocked.CAS(false, true) {
+			return database.ErrLocked
+		}
 		return nil
 	}
 }
@@ -208,14 +211,18 @@ func (c *CockroachDb) Unlock() error {
 			// 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
+				if !c.isLocked.CAS(true, false) {
+					return database.ErrNotLocked
+				}
 				return nil
 			}
 		}
 		return database.Error{OrigErr: err, Err: "failed to release migration lock", Query: []byte(query)}
 	}
 
-	c.isLocked = false
+	if !c.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }
 

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/mysql/mysql.go b/database/mysql/mysql.go
@@ -8,6 +8,7 @@ import (
 	"crypto/x509"
 	"database/sql"
 	"fmt"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -49,7 +50,7 @@ type Mysql struct {
 	// just do everything over a single conn anyway.
 	conn     *sql.Conn
 	db       *sql.DB
-	isLocked bool
+	isLocked atomic.Bool
 
 	config *Config
 }
@@ -251,12 +252,10 @@ func (m *Mysql) Close() error {
 }
 
 func (m *Mysql) Lock() error {
-	if m.isLocked {
-		return database.ErrLocked
-	}
-
 	if m.config.NoLock {
-		m.isLocked = true
+		if !m.isLocked.CAS(false, true) {
+			return database.ErrLocked
+		}
 		return nil
 	}
 
@@ -273,20 +272,24 @@ func (m *Mysql) Lock() error {
 	}
 
 	if success {
-		m.isLocked = true
+		if !m.isLocked.CAS(false, true) {
+			return database.ErrLocked
+		}
 		return nil
 	}
 
 	return database.ErrLocked
 }
 
 func (m *Mysql) Unlock() error {
-	if !m.isLocked {
+	if !m.isLocked.Load() {
 		return nil
 	}
 
 	if m.config.NoLock {
-		m.isLocked = false
+		if !m.isLocked.CAS(true, false) {
+			return database.ErrNotLocked
+		}
 		return nil
 	}
 
@@ -305,7 +308,9 @@ func (m *Mysql) Unlock() error {
 	// in which case isLocked should be true until the timeout expires -- synchronizing
 	// these states is likely not worth trying to do; reconsider the necessity of isLocked.
 
-	m.isLocked = false
+	if !m.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }
 

diff --git a/database/pgx/pgx.go b/database/pgx/pgx.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -58,7 +59,7 @@ type Postgres 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
@@ -220,7 +221,7 @@ func (p *Postgres) Close() error {
 
 // https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
 func (p *Postgres) Lock() error {
-	if p.isLocked {
+	if p.isLocked.Load() {
 		return database.ErrLocked
 	}
 
@@ -235,12 +236,14 @@ func (p *Postgres) Lock() error {
 		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
 	}
 
-	p.isLocked = true
+	if !p.isLocked.CAS(false, true) {
+		return database.ErrLocked
+	}
 	return nil
 }
 
 func (p *Postgres) Unlock() error {
-	if !p.isLocked {
+	if !p.isLocked.Load() {
 		return nil
 	}
 
@@ -253,7 +256,9 @@ func (p *Postgres) Unlock() error {
 	if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
-	p.isLocked = false
+	if !p.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }
 

diff --git a/database/postgres/postgres.go b/database/postgres/postgres.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -57,7 +58,7 @@ type Postgres 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
@@ -214,7 +215,7 @@ func (p *Postgres) Close() error {
 
 // https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
 func (p *Postgres) Lock() error {
-	if p.isLocked {
+	if p.isLocked.Load() {
 		return database.ErrLocked
 	}
 
@@ -229,12 +230,14 @@ func (p *Postgres) Lock() error {
 		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
 	}
 
-	p.isLocked = true
+	if !p.isLocked.CAS(false, true) {
+		return database.ErrLocked
+	}
 	return nil
 }
 
 func (p *Postgres) Unlock() error {
-	if !p.isLocked {
+	if !p.isLocked.Load() {
 		return nil
 	}
 
@@ -247,7 +250,9 @@ func (p *Postgres) Unlock() error {
 	if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
-	p.isLocked = false
+	if !p.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }
 

diff --git a/database/ql/ql.go b/database/ql/ql.go
@@ -4,6 +4,7 @@ import (
 	"database/sql"
 	"fmt"
 	"github.com/hashicorp/go-multierror"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	"strings"
@@ -34,7 +35,7 @@ type Config struct {
 
 type Ql struct {
 	db       *sql.DB
-	isLocked bool
+	isLocked atomic.Bool
 
 	config *Config
 }
@@ -166,17 +167,15 @@ func (m *Ql) Drop() (err error) {
 	return nil
 }
 func (m *Ql) Lock() error {
-	if m.isLocked {
+	if !m.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
-	m.isLocked = true
 	return nil
 }
 func (m *Ql) Unlock() error {
-	if !m.isLocked {
-		return nil
+	if !m.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
 	}
-	m.isLocked = false
 	return nil
 }
 func (m *Ql) Run(migration io.Reader) error {

diff --git a/database/redshift/redshift.go b/database/redshift/redshift.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"go.uber.org/atomic"
 	"io"
 	"io/ioutil"
 	nurl "net/url"
@@ -36,7 +37,7 @@ type Config struct {
 }
 
 type Redshift struct {
-	isLocked bool
+	isLocked atomic.Bool
 	conn     *sql.Conn
 	db       *sql.DB
 
@@ -126,15 +127,16 @@ func (p *Redshift) Close() error {
 
 // Redshift does not support advisory lock functions: https://docs.aws.amazon.com/redshift/latest/dg/c_unsupported-postgresql-functions.html
 func (p *Redshift) Lock() error {
-	if p.isLocked {
+	if !p.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
-	p.isLocked = true
 	return nil
 }
 
 func (p *Redshift) Unlock() error {
-	p.isLocked = false
+	if !p.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
+	}
 	return nil
 }