fixing local lock optimization implementation

Author: prinkov

database/pgx/pgx.go

@@ -221,44 +221,41 @@ 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.Load() {
+	if !p.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
 
 	aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
 	if err != nil {
+		p.isLocked.Store(false)
 		return 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 {
+		p.isLocked.Store(false)
 		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
 	}
-
-	if !p.isLocked.CAS(false, true) {
-		return database.ErrLocked
-	}
 	return nil
 }
 
 func (p *Postgres) Unlock() error {
-	if !p.isLocked.Load() {
-		return nil
+	if !p.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
 	}
 
 	aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
 	if err != nil {
+		p.isLocked.Store(true)
 		return err
 	}
 
 	query := `SELECT pg_advisory_unlock($1)`
 	if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
+		p.isLocked.Store(true)
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
-	if !p.isLocked.CAS(true, false) {
-		return database.ErrNotLocked
-	}
 	return nil
 }
 

database/postgres/postgres.go

@@ -215,39 +215,40 @@ 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.Load() {
+	if !p.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
 
 	aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
 	if err != nil {
+		p.isLocked.Store(false)
 		return 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 {
+		p.isLocked.Store(false)
 		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
 	}
 
-	if !p.isLocked.CAS(false, true) {
-		return database.ErrLocked
-	}
 	return nil
 }
 
 func (p *Postgres) Unlock() error {
-	if !p.isLocked.Load() {
-		return nil
+	if !p.isLocked.CAS(true, false) {
+		return database.ErrLocked
 	}
 
 	aid, err := database.GenerateAdvisoryLockId(p.config.DatabaseName, p.config.migrationsSchemaName, p.config.migrationsTableName)
 	if err != nil {
+		p.isLocked.Store(true)
 		return err
 	}
 
 	query := `SELECT pg_advisory_unlock($1)`
 	if _, err := p.conn.ExecContext(context.Background(), query, aid); err != nil {
+		p.isLocked.Store(true)
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
 	if !p.isLocked.CAS(true, false) {

database/sqlserver/sqlserver.go

@@ -155,7 +155,7 @@ func (ss *SQLServer) Close() error {
 
 // Lock creates an advisory local on the database to prevent multiple migrations from running at the same time.
 func (ss *SQLServer) Lock() error {
-	if ss.isLocked.Load() {
+	if !ss.isLocked.CAS(false, true) {
 		return database.ErrLocked
 	}
 
@@ -171,36 +171,35 @@ func (ss *SQLServer) Lock() error {
 
 	var status mssql.ReturnStatus
 	if _, err = ss.conn.ExecContext(context.Background(), query, aid, &status); err == nil && status > -1 {
-		if !ss.isLocked.CAS(false, true) {
-			return database.ErrLocked
-		}
 		return nil
 	} else if err != nil {
+		ss.isLocked.Store(false)
 		return &database.Error{OrigErr: err, Err: "try lock failed", Query: []byte(query)}
 	} else {
+		ss.isLocked.Store(false)
 		return &database.Error{Err: fmt.Sprintf("try lock failed with error %v: %v", status, lockErrorMap[status]), Query: []byte(query)}
 	}
 }
 
 // Unlock froms the migration lock from the database
 func (ss *SQLServer) Unlock() error {
-	if !ss.isLocked.Load() {
-		return nil
+	if !ss.isLocked.CAS(true, false) {
+		return database.ErrNotLocked
 	}
 
 	aid, err := database.GenerateAdvisoryLockId(ss.config.DatabaseName, ss.config.SchemaName)
 	if err != nil {
+		ss.isLocked.Store(true)
 		return err
 	}
 
 	// MS Docs: sp_releaseapplock: https://docs.microsoft.com/en-us/sql/relational-databases/system-stored-procedures/sp-releaseapplock-transact-sql?view=sql-server-2017
 	query := `EXEC sp_releaseapplock @Resource = @p1, @LockOwner = 'Session'`
 	if _, err := ss.conn.ExecContext(context.Background(), query, aid); err != nil {
+		ss.isLocked.Store(true)
 		return &database.Error{OrigErr: err, Query: []byte(query)}
 	}
-	if !ss.isLocked.CAS(true, false) {
-		return database.ErrNotLocked
-	}
+
 	return nil
 }