[8.4.x]: SQLStore: Prevent concurrent migrations (#44101) (#45447)

* SQLStore: Prevent concurrent migrations (#44101)

* SQLStore: Prevent concurrent migrations

* Hide behind a feature toggle

* Configurable locking attempt timeout

* Update docs/sources/administration/configuration.md

Co-authored-by: Igor Suleymanov <radiohead@users.noreply.github.com>
Co-authored-by: achatterjee-grafana <70489351+achatterjee-grafana@users.noreply.github.com>
(cherry picked from commit d718ee1918)

* Resolve dependency cycle (#45427)

(cherry picked from commit 6a38ce2307)

pkg/services/sqlstore/migrator/dialect.go

@@ -7,6 +7,11 @@ import (
 	"xorm.io/xorm"
 )
 
+var (
+	ErrLockDB        = fmt.Errorf("failed to obtain lock")
+	ErrReleaseLockDB = fmt.Errorf("failed to release lock")
+)
+
 type Dialect interface {
 	DriverName() string
 	Quote(string) string
@@ -53,6 +58,13 @@ type Dialect interface {
 	IsUniqueConstraintViolation(err error) bool
 	ErrorMessage(err error) string
 	IsDeadlock(err error) bool
+	Lock(LockCfg) error
+	Unlock(LockCfg) error
+}
+
+type LockCfg struct {
+	Session *xorm.Session
+	Timeout int
 }
 
 type dialectFunc func(*xorm.Engine) Dialect
@@ -288,3 +300,11 @@ func (b *BaseDialect) TruncateDBTables() error {
 func (b *BaseDialect) UpsertSQL(tableName string, keyCols, updateCols []string) string {
 	return ""
 }
+
+func (b *BaseDialect) Lock(_ LockCfg) error {
+	return nil
+}
+
+func (b *BaseDialect) Unlock(_ LockCfg) error {
+	return nil
+}

pkg/services/sqlstore/migrator/migrator.go

@@ -7,6 +7,7 @@ import (
 	_ "github.com/go-sql-driver/mysql"
 	_ "github.com/lib/pq"
 	_ "github.com/mattn/go-sqlite3"
+	"go.uber.org/atomic"
 	"xorm.io/xorm"
 
 	"github.com/grafana/grafana/pkg/infra/log"
@@ -14,12 +15,18 @@ import (
 	"github.com/grafana/grafana/pkg/util/errutil"
 )
 
+var (
+	ErrMigratorIsLocked   = fmt.Errorf("migrator is locked")
+	ErrMigratorIsUnlocked = fmt.Errorf("migrator is unlocked")
+)
+
 type Migrator struct {
 	DBEngine   *xorm.Engine
 	Dialect    Dialect
 	migrations []Migration
 	Logger     log.Logger
 	Cfg        *setting.Cfg
+	isLocked   atomic.Bool
 }
 
 type MigrationLog struct {
@@ -87,7 +94,32 @@ func (mg *Migrator) GetMigrationLog() (map[string]MigrationLog, error) {
 	return logMap, nil
 }
 
-func (mg *Migrator) Start() error {
+func (mg *Migrator) Start(isDatabaseLockingEnabled bool, lockAttemptTimeout int) (err error) {
+	if !isDatabaseLockingEnabled {
+		return mg.run()
+	}
+
+	return mg.InTransaction(func(sess *xorm.Session) error {
+		mg.Logger.Info("Locking database")
+		if err := casRestoreOnErr(&mg.isLocked, false, true, ErrMigratorIsLocked, mg.Dialect.Lock, LockCfg{Session: sess, Timeout: lockAttemptTimeout}); err != nil {
+			mg.Logger.Error("Failed to lock database", "error", err)
+			return err
+		}
+
+		defer func() {
+			mg.Logger.Info("Unlocking database")
+			unlockErr := casRestoreOnErr(&mg.isLocked, true, false, ErrMigratorIsUnlocked, mg.Dialect.Unlock, LockCfg{Session: sess})
+			if unlockErr != nil {
+				mg.Logger.Error("Failed to unlock database", "error", unlockErr)
+			}
+		}()
+
+		// migration will run inside a nested transaction
+		return mg.run()
+	})
+}
+
+func (mg *Migrator) run() (err error) {
 	mg.Logger.Info("Starting DB migrations")
 
 	logMap, err := mg.GetMigrationLog()
@@ -211,3 +243,15 @@ func (mg *Migrator) InTransaction(callback dbTransactionFunc) error {
 
 	return nil
 }
+
+func casRestoreOnErr(lock *atomic.Bool, o, n bool, casErr error, f func(LockCfg) error, lockCfg LockCfg) error {
+	if !lock.CAS(o, n) {
+		return casErr
+	}
+	if err := f(lockCfg); err != nil {
+		// Automatically unlock/lock on error
+		lock.Store(o)
+		return err
+	}
+	return nil
+}

pkg/services/sqlstore/migrator/mysql_dialect.go

@@ -1,6 +1,7 @@
 package migrator
 
 import (
+	"database/sql"
 	"errors"
 	"fmt"
 	"strconv"
@@ -8,6 +9,7 @@ import (
 
 	"github.com/VividCortex/mysqlerr"
 	"github.com/go-sql-driver/mysql"
+	"github.com/golang-migrate/migrate/v4/database"
 	"github.com/grafana/grafana/pkg/util/errutil"
 	"xorm.io/xorm"
 )
@@ -225,3 +227,66 @@ func (db *MySQLDialect) UpsertSQL(tableName string, keyCols, updateCols []string
 	)
 	return s
 }
+
+func (db *MySQLDialect) Lock(cfg LockCfg) error {
+	query := "SELECT GET_LOCK(?, ?)"
+	var success sql.NullBool
+
+	lockName, err := db.getLockName()
+	if err != nil {
+		return fmt.Errorf("failed to generate lock name: %w", err)
+	}
+
+	// trying to obtain the lock with the specific name
+	// the lock is exclusive per session and is released explicitly by executing RELEASE_LOCK() or implicitly when the session terminates
+	// it returns 1 if the lock was obtained successfully,
+	// 0 if the attempt timed out (for example, because another client has previously locked the name),
+	// or NULL if an error occurred
+	// starting from MySQL 5.7 it is even possible for a given session to acquire multiple locks for the same name
+	// however other sessions cannot acquire a lock with that name until the acquiring session releases all its locks for the name.
+	_, err = cfg.Session.SQL(query, lockName, cfg.Timeout).Get(&success)
+	if err != nil {
+		return err
+	}
+	if !success.Valid || !success.Bool {
+		return ErrLockDB
+	}
+	return nil
+}
+
+func (db *MySQLDialect) Unlock(cfg LockCfg) error {
+	query := "SELECT RELEASE_LOCK(?)"
+	var success sql.NullBool
+
+	lockName, err := db.getLockName()
+	if err != nil {
+		return fmt.Errorf("failed to generate lock name: %w", err)
+	}
+
+	// trying to release the lock with the specific name
+	// it returns 1 if the lock was released,
+	// 0 if the lock was not established by this thread (in which case the lock is not released),
+	// and NULL if the named lock did not exist (it was never obtained by a call to GET_LOCK() or if it has previously been released)
+	_, err = cfg.Session.SQL(query, lockName).Get(&success)
+	if err != nil {
+		return err
+	}
+	if !success.Valid || !success.Bool {
+		return ErrReleaseLockDB
+	}
+	return nil
+}
+
+func (db *MySQLDialect) getLockName() (string, error) {
+	cfg, err := mysql.ParseDSN(db.engine.DataSourceName())
+	if err != nil {
+		return "", err
+	}
+
+	s, err := database.GenerateAdvisoryLockId(cfg.DBName)
+	if err != nil {
+		return "", fmt.Errorf("failed to generate advisory lock key: %w", err)
+	}
+
+	return s, nil
+}

pkg/services/sqlstore/migrator/postgres_dialect.go

@@ -3,9 +3,11 @@ package migrator
 import (
 	"errors"
 	"fmt"
+	"regexp"
 	"strconv"
 	"strings"
 
+	"github.com/golang-migrate/migrate/v4/database"
 	"github.com/lib/pq"
 
 	"github.com/grafana/grafana/pkg/util/errutil"
@@ -257,3 +259,76 @@ func (db *PostgresDialect) UpsertSQL(tableName string, keyCols, updateCols []str
 	)
 	return s
 }
+
+func (db *PostgresDialect) Lock(cfg LockCfg) error {
+	// trying to obtain the lock for a resource identified by a 64-bit or 32-bit key value
+	// the lock is exclusive: multiple lock requests stack, so that if the same resource is locked three times
+	// it must then be unlocked three times to be released for other sessions' use.
+	// it will either obtain the lock immediately and return true,
+	// or return false if the lock cannot be acquired immediately.
+	query := "SELECT pg_try_advisory_lock(?)"
+	var success bool
+
+	key, err := db.getLockKey()
+	if err != nil {
+		return fmt.Errorf("failed to generate advisory lock key: %w", err)
+	}
+	_, err = cfg.Session.SQL(query, key).Get(&success)
+	if err != nil {
+		return err
+	}
+	if !success {
+		return ErrLockDB
+	}
+
+	return nil
+}
+
+func (db *PostgresDialect) Unlock(cfg LockCfg) error {
+	// trying to release a previously-acquired exclusive session level advisory lock.
+	// it will either return true if the lock is successfully released or
+	// false if the lock was not held (in addition an SQL warning will be reported by the server)
+	query := "SELECT pg_advisory_unlock(?)"
+	var success bool
+
+	key, err := db.getLockKey()
+	if err != nil {
+		return fmt.Errorf("failed to generate advisory lock key: %w", err)
+	}
+	_, err = cfg.Session.SQL(query, key).Get(&success)
+	if err != nil {
+		return err
+	}
+	if !success {
+		return ErrReleaseLockDB
+	}
+	return nil
+}
+
+func getDBName(dsn string) (string, error) {
+	if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") {
+		parsedDSN, err := pq.ParseURL(dsn)
+		if err != nil {
+			return "", err
+		}
+		dsn = parsedDSN
+	}
+	re := regexp.MustCompile(`dbname=(\w+)`)
+	submatch := re.FindSubmatch([]byte(dsn))
+	if len(submatch) < 2 {
+		return "", fmt.Errorf("failed to get database name")
+	}
+	return string(submatch[1]), nil
+}
+
+func (db *PostgresDialect) getLockKey() (string, error) {
+	dbName, err := getDBName(db.engine.DataSourceName())
+	if err != nil {
+		return "", err
+	}
+	key, err := database.GenerateAdvisoryLockId(dbName)
+	if err != nil {
+		return "", err
+	}
+	return key, nil
+}