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>

pkg/services/sqlstore/migrations/accesscontrol/test/ac_test.go

@@ -165,7 +165,7 @@ func TestMigrations(t *testing.T) {
 			acmigrator := migrator.NewMigrator(x, tc.config)
 			acmig.AddTeamMembershipMigrations(acmigrator)
 
-			errRunningMig := acmigrator.Start()
+			errRunningMig := acmigrator.Start(false, 0)
 			require.NoError(t, errRunningMig)
 
 			for _, user := range users {
@@ -221,7 +221,7 @@ func setupTestDB(t *testing.T) *xorm.Engine {
 	migrations := &migrations.OSSMigrations{}
 	migrations.AddMigration(mg)
 
-	err = mg.Start()
+	err = mg.Start(false, 0)
 	require.NoError(t, err)
 
 	return x

pkg/services/sqlstore/migrations/migrations_test.go

@@ -1,10 +1,16 @@
 package migrations
 
 import (
+	"errors"
 	"fmt"
+	"os"
 	"strings"
+	"sync"
+	"sync/atomic"
 	"testing"
 
+	"github.com/go-sql-driver/mysql"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"xorm.io/xorm"
 
@@ -32,7 +38,7 @@ func TestMigrations(t *testing.T) {
 	migrations.AddMigration(mg)
 	expectedMigrations := mg.GetMigrationIDs(true)
 
-	err = mg.Start()
+	err = mg.Start(false, 0)
 	require.NoError(t, err)
 
 	has, err := x.SQL(query).Get(&result)
@@ -44,7 +50,7 @@ func TestMigrations(t *testing.T) {
 	mg = NewMigrator(x, &setting.Cfg{})
 	migrations.AddMigration(mg)
 
-	err = mg.Start()
+	err = mg.Start(false, 0)
 	require.NoError(t, err)
 
 	has, err = x.SQL(query).Get(&result)
@@ -53,6 +59,179 @@ func TestMigrations(t *testing.T) {
 	checkStepsAndDatabaseMatch(t, mg, expectedMigrations)
 }
 
+func TestMigrationLock(t *testing.T) {
+	dbType := getDBType()
+	if dbType == SQLite {
+		t.Skip()
+	}
+
+	testDB := getTestDB(dbType)
+
+	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
+	require.NoError(t, err)
+
+	dialect := NewDialect(x)
+
+	sess := x.NewSession()
+	t.Cleanup(func() {
+		sess.Close()
+	})
+
+	cfg := LockCfg{Session: sess}
+
+	t.Run("obtaining lock should succeed", func(t *testing.T) {
+		err := dialect.Lock(cfg)
+		require.NoError(t, err)
+
+		t.Run("releasing previously obtained lock should succeed", func(t *testing.T) {
+			err := dialect.Unlock(cfg)
+			require.NoError(t, err)
+
+			t.Run("releasing already released lock should fail", func(t *testing.T) {
+				err := dialect.Unlock(cfg)
+				require.Error(t, err)
+				assert.ErrorIs(t, err, ErrReleaseLockDB)
+			})
+		})
+	})
+
+	t.Run("obtaining lock twice should succeed", func(t *testing.T) {
+		err = dialect.Lock(cfg)
+		require.NoError(t, err)
+
+		err = dialect.Lock(cfg)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := dialect.Unlock(cfg)
+			require.NoError(t, err)
+
+			err = dialect.Unlock(cfg)
+			require.NoError(t, err)
+		})
+	})
+
+	t.Run("obtaining same lock from another session should fail", func(t *testing.T) {
+		x2, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
+		require.NoError(t, err)
+		sess2 := x2.NewSession()
+
+		d2 := NewDialect(x2)
+
+		err = dialect.Lock(cfg)
+		require.NoError(t, err)
+
+		err = d2.Lock(LockCfg{Session: sess2})
+		require.Error(t, err)
+		assert.ErrorIs(t, err, ErrLockDB)
+
+		t.Cleanup(func() {
+			err := dialect.Unlock(cfg)
+			require.NoError(t, err)
+		})
+	})
+
+	t.Run("obtaining lock for a another database should succeed", func(t *testing.T) {
+		err := dialect.Lock(cfg)
+		require.NoError(t, err)
+
+		x, err := xorm.NewEngine(testDB.DriverName, replaceDBName(t, testDB.ConnStr, dbType))
+		require.NoError(t, err)
+
+		d := NewDialect(x)
+		err = d.Lock(cfg)
+		require.NoError(t, err)
+
+		t.Cleanup(func() {
+			err := dialect.Unlock(cfg)
+			require.NoError(t, err)
+
+			err = d.Unlock(cfg)
+			require.NoError(t, err)
+		})
+	})
+}
+
+func TestMigratorLocking(t *testing.T) {
+	dbType := getDBType()
+	testDB := getTestDB(dbType)
+
+	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
+	require.NoError(t, err)
+
+	err = NewDialect(x).CleanDB()
+	require.NoError(t, err)
+
+	mg := NewMigrator(x, &setting.Cfg{})
+	migrations := &OSSMigrations{}
+	migrations.AddMigration(mg)
+
+	var errorNum int64
+	t.Run("when concurrent migrations for the same migrator occur, the second one should fail", func(t *testing.T) {
+		for i := 0; i < 2; i++ {
+			i := i // capture i variable
+			t.Run(fmt.Sprintf("run migration %d", i), func(t *testing.T) {
+				t.Parallel()
+				err = mg.Start(true, 0)
+				if err != nil {
+					if errors.Is(err, ErrMigratorIsLocked) {
+						atomic.AddInt64(&errorNum, 1)
+					}
+				}
+			})
+		}
+	})
+	assert.Equal(t, int64(1), errorNum)
+}
+
+func TestDatabaseLocking(t *testing.T) {
+	dbType := getDBType()
+	// skip for SQLite since there is no database locking (only migrator locking)
+	if dbType == SQLite {
+		t.Skip()
+	}
+
+	testDB := getTestDB(dbType)
+
+	x, err := xorm.NewEngine(testDB.DriverName, testDB.ConnStr)
+	require.NoError(t, err)
+
+	err = NewDialect(x).CleanDB()
+	require.NoError(t, err)
+
+	mg1 := NewMigrator(x, &setting.Cfg{})
+	migrations := &OSSMigrations{}
+	migrations.AddMigration(mg1)
+	reg := registry{
+		migrators: make(map[int]*Migrator, 2),
+	}
+	reg.set(0, mg1)
+
+	mg2 := NewMigrator(x, &setting.Cfg{})
+	migrations.AddMigration(mg2)
+	reg.set(1, mg2)
+
+	var errorNum int64
+	t.Run("when concurrent migrations occur for different migrators occur, the second one should fail", func(t *testing.T) {
+		for i := 0; i < 2; i++ {
+			i := i // capture i variable
+			t.Run(fmt.Sprintf("run migration %d", i), func(t *testing.T) {
+				mg, err := reg.get(i)
+				require.NoError(t, err)
+				t.Parallel()
+				err = mg.Start(true, 0)
+				if err != nil {
+					assert.ErrorIs(t, err, ErrLockDB)
+					if errors.Is(err, ErrLockDB) {
+						atomic.AddInt64(&errorNum, 1)
+					}
+				}
+			})
+		}
+	})
+	assert.Equal(t, int64(1), errorNum)
+}
+
 func checkStepsAndDatabaseMatch(t *testing.T, mg *Migrator, expected []string) {
 	t.Helper()
 	log, err := mg.GetMigrationLog()
@@ -91,3 +270,61 @@ func checkStepsAndDatabaseMatch(t *testing.T, mg *Migrator, expected []string) {
 	}
 	require.Failf(t, "the number of migrations does not match log in database", msg)
 }
+
+func getDBType() string {
+	dbType := SQLite
+
+	// environment variable present for test db?
+	if db, present := os.LookupEnv("GRAFANA_TEST_DB"); present {
+		dbType = db
+	}
+	return dbType
+}
+
+func getTestDB(dbType string) sqlutil.TestDB {
+	switch dbType {
+	case "mysql":
+		return sqlutil.MySQLTestDB()
+	case "postgres":
+		return sqlutil.PostgresTestDB()
+	default:
+		return sqlutil.SQLite3TestDB()
+	}
+}
+
+func replaceDBName(t *testing.T, connStr, dbType string) string {
+	switch dbType {
+	case "mysql":
+		cfg, err := mysql.ParseDSN(connStr)
+		require.NoError(t, err)
+		cfg.DBName = "grafana_ds_tests"
+		return cfg.FormatDSN()
+	case "postgres":
+		return strings.Replace(connStr, "dbname=grafanatest", "dbname=grafanadstest", 1)
+	default:
+		return connStr
+	}
+}
+
+type registry struct {
+	mu        sync.Mutex
+	migrators map[int]*Migrator
+}
+
+func (r *registry) get(i int) (*Migrator, error) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	m, ok := r.migrators[i]
+	if !ok {
+		return nil, fmt.Errorf("invalid index: %d", i)
+	}
+	return m, nil
+}
+
+func (r *registry) set(i int, mg *Migrator) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.migrators[i] = mg
+}

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
+}

pkg/services/sqlstore/mockstore/mockstore.go

@@ -546,7 +546,7 @@ func (m *SQLStoreMock) UpdateDataSource(ctx context.Context, cmd *models.UpdateD
 	return m.ExpectedError
 }
 
-func (m *SQLStoreMock) Migrate() error {
+func (m *SQLStoreMock) Migrate(_ bool) error {
 	return m.ExpectedError
 }
 

pkg/services/sqlstore/sqlstore.go

@@ -56,8 +56,7 @@ type SQLStore struct {
 	tracer                      tracing.Tracer
 }
 
-func ProvideService(cfg *setting.Cfg, cacheService *localcache.CacheService, bus bus.Bus, migrations registry.DatabaseMigrator, tracer tracing.Tracer,
-) (*SQLStore, error) {
+func ProvideService(cfg *setting.Cfg, cacheService *localcache.CacheService, bus bus.Bus, migrations registry.DatabaseMigrator, tracer tracing.Tracer, features *featuremgmt.FeatureManager) (*SQLStore, error) {
 	// This change will make xorm use an empty default schema for postgres and
 	// by that mimic the functionality of how it was functioning before
 	// xorm's changes above.
@@ -67,7 +66,7 @@ func ProvideService(cfg *setting.Cfg, cacheService *localcache.CacheService, bus
 		return nil, err
 	}
 
-	if err := s.Migrate(); err != nil {
+	if err := s.Migrate(features.IsEnabled(featuremgmt.FlagMigrationLocking)); err != nil {
 		return nil, err
 	}
 
@@ -153,7 +152,7 @@ func newSQLStore(cfg *setting.Cfg, cacheService *localcache.CacheService, b bus.
 // Migrate performs database migrations.
 // Has to be done in a second phase (after initialization), since other services can register migrations during
 // the initialization phase.
-func (ss *SQLStore) Migrate() error {
+func (ss *SQLStore) Migrate(isDatabaseLockingEnabled bool) error {
 	if ss.dbCfg.SkipMigrations {
 		return nil
 	}
@@ -161,7 +160,7 @@ func (ss *SQLStore) Migrate() error {
 	migrator := migrator.NewMigrator(ss.engine, ss.Cfg)
 	ss.migrations.AddMigration(migrator)
 
-	return migrator.Start()
+	return migrator.Start(isDatabaseLockingEnabled, ss.dbCfg.MigrationLockAttemptTimeout)
 }
 
 // Sync syncs changes to the database.
@@ -443,6 +442,7 @@ func (ss *SQLStore) readConfig() error {
 
 	ss.dbCfg.CacheMode = sec.Key("cache_mode").MustString("private")
 	ss.dbCfg.SkipMigrations = sec.Key("skip_migrations").MustBool()
+	ss.dbCfg.MigrationLockAttemptTimeout = sec.Key("locking_attempt_timeout_sec").MustInt()
 	return nil
 }
 
@@ -562,7 +562,7 @@ func initTestDB(migration registry.DatabaseMigrator, opts ...InitTestDBOpt) (*SQ
 			return nil, err
 		}
 
-		if err := testSQLStore.Migrate(); err != nil {
+		if err := testSQLStore.Migrate(false); err != nil {
 			return nil, err
 		}
 
@@ -622,23 +622,24 @@ func IsTestDBMSSQL() bool {
 }
 
 type DatabaseConfig struct {
-	Type             string
-	Host             string
-	Name             string
-	User             string
-	Pwd              string
-	Path             string
-	SslMode          string
-	CaCertPath       string
-	ClientKeyPath    string
-	ClientCertPath   string
-	ServerCertName   string
-	ConnectionString string
-	IsolationLevel   string
-	MaxOpenConn      int
-	MaxIdleConn      int
-	ConnMaxLifetime  int
-	CacheMode        string
-	UrlQueryParams   map[string][]string
-	SkipMigrations   bool
+	Type                        string
+	Host                        string
+	Name                        string
+	User                        string
+	Pwd                         string
+	Path                        string
+	SslMode                     string
+	CaCertPath                  string
+	ClientKeyPath               string
+	ClientCertPath              string
+	ServerCertName              string
+	ConnectionString            string
+	IsolationLevel              string
+	MaxOpenConn                 int
+	MaxIdleConn                 int
+	ConnMaxLifetime             int
+	CacheMode                   string
+	UrlQueryParams              map[string][]string
+	SkipMigrations              bool
+	MigrationLockAttemptTimeout int
 }

pkg/services/sqlstore/store.go

@@ -125,7 +125,7 @@ type Store interface {
 	DeleteDataSource(ctx context.Context, cmd *models.DeleteDataSourceCommand) error
 	AddDataSource(ctx context.Context, cmd *models.AddDataSourceCommand) error
 	UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSourceCommand) error
-	Migrate() error
+	Migrate(bool) error
 	Sync() error
 	Reset() error
 	Quote(value string) string