MM-50427: Make MM survive DB replica outage (#22888)

We monitor the health of DB replicas, and on a fatal error,
take them out of the pool.

On a separate goroutine, we keep pinging the unhealthy replicas,
and on getting a good response back, we add them back to the pool.

https://mattermost.atlassian.net/browse/MM-50427

```release-note
Mattermost is now resilient against DB replica outages and will
dynamically choose a replica if it's alive.

Also added a config parameter ReplicaMonitorIntervalSeconds
whose default value is 5. This controls how frequently unhealthy
replicas will be monitored for liveness check.
```

Co-authored-by: Mattermost Build <build@mattermost.com>

server/channels/store/sqlstore/store.go

@@ -49,7 +49,7 @@ const (
 	MySQLForeignKeyViolationErrorCode = 1452
 	PGDuplicateObjectErrorCode        = "42710"
 	MySQLDuplicateObjectErrorCode     = 1022
-	DBPingAttempts                    = 18
+	DBPingAttempts                    = 5
 	DBPingTimeoutSecs                 = 10
 	// This is a numerical version string by postgres. The format is
 	// 2 characters for major, minor, and patch version prior to 10.
@@ -123,9 +123,9 @@ type SqlStore struct {
 
 	masterX *sqlxDBWrapper
 
-	ReplicaXs []*sqlxDBWrapper
+	ReplicaXs []*atomic.Pointer[sqlxDBWrapper]
 
-	searchReplicaXs []*sqlxDBWrapper
+	searchReplicaXs []*atomic.Pointer[sqlxDBWrapper]
 
 	replicaLagHandles []*dbsql.DB
 	stores            SqlStoreStores
@@ -138,17 +138,28 @@ type SqlStore struct {
 
 	isBinaryParam             bool
 	pgDefaultTextSearchConfig string
+
+	quitMonitor chan struct{}
+	wgMonitor   *sync.WaitGroup
 }
 
 func New(settings model.SqlSettings, metrics einterfaces.MetricsInterface) *SqlStore {
 	store := &SqlStore{
-		rrCounter: 0,
-		srCounter: 0,
-		settings:  &settings,
-		metrics:   metrics,
+		rrCounter:   0,
+		srCounter:   0,
+		settings:    &settings,
+		metrics:     metrics,
+		quitMonitor: make(chan struct{}),
+		wgMonitor:   &sync.WaitGroup{},
 	}
 
-	store.initConnection()
+	err := store.initConnection()
+	if err != nil {
+		mlog.Fatal("Error setting up connections", mlog.Err(err))
+	}
+
+	store.wgMonitor.Add(1)
+	go store.monitorReplicas()
 
 	ver, err := store.GetDbVersion(true)
 	if err != nil {
@@ -230,29 +241,28 @@ func New(settings model.SqlSettings, metrics einterfaces.MetricsInterface) *SqlS
 
 // SetupConnection sets up the connection to the database and pings it to make sure it's alive.
 // It also applies any database configuration settings that are required.
-func SetupConnection(connType string, dataSource string, settings *model.SqlSettings) *dbsql.DB {
+func SetupConnection(connType string, dataSource string, settings *model.SqlSettings, attempts int) (*dbsql.DB, error) {
 	db, err := dbsql.Open(*settings.DriverName, dataSource)
 	if err != nil {
-		mlog.Fatal("Failed to open SQL connection to err.", mlog.Err(err))
+		return nil, errors.Wrap(err, "failed to open SQL connection")
 	}
 
-	for i := 0; i < DBPingAttempts; i++ {
+	for i := 0; i < attempts; i++ {
 		// At this point, we have passed sql.Open, so we deliberately ignore any errors.
 		sanitized, _ := SanitizeDataSource(*settings.DriverName, dataSource)
 		mlog.Info("Pinging SQL", mlog.String("database", connType), mlog.String("dataSource", sanitized))
 		ctx, cancel := context.WithTimeout(context.Background(), DBPingTimeoutSecs*time.Second)
 		defer cancel()
 		err = db.PingContext(ctx)
-		if err == nil {
-			break
-		} else {
-			if i == DBPingAttempts-1 {
-				mlog.Fatal("Failed to ping DB, server will exit.", mlog.Err(err))
-			} else {
-				mlog.Error("Failed to ping DB", mlog.Err(err), mlog.Int("retrying in seconds", DBPingTimeoutSecs))
-				time.Sleep(DBPingTimeoutSecs * time.Second)
+		if err != nil {
+			if i == attempts-1 {
+				return nil, err
 			}
+			mlog.Error("Failed to ping DB", mlog.Err(err), mlog.Int("retrying in seconds", DBPingTimeoutSecs))
+			time.Sleep(DBPingTimeoutSecs * time.Second)
+			continue
 		}
+		break
 	}
 
 	if strings.HasPrefix(connType, replicaLagPrefix) {
@@ -272,7 +282,7 @@ func SetupConnection(connType string, dataSource string, settings *model.SqlSett
 	db.SetConnMaxLifetime(time.Duration(*settings.ConnMaxLifetimeMilliseconds) * time.Millisecond)
 	db.SetConnMaxIdleTime(time.Duration(*settings.ConnMaxIdleTimeMilliseconds) * time.Millisecond)
 
-	return db
+	return db, nil
 }
 
 func (ss *SqlStore) SetContext(context context.Context) {
@@ -285,7 +295,7 @@ func (ss *SqlStore) Context() context.Context {
 
 func noOpMapper(s string) string { return s }
 
-func (ss *SqlStore) initConnection() {
+func (ss *SqlStore) initConnection() error {
 	dataSource := *ss.settings.DataSource
 	if ss.DriverName() == model.DatabaseDriverMysql {
 		// TODO: We ignore the readTimeout datasource parameter for MySQL since QueryTimeout
@@ -294,11 +304,14 @@ func (ss *SqlStore) initConnection() {
 		var err error
 		dataSource, err = ResetReadTimeout(dataSource)
 		if err != nil {
-			mlog.Fatal("Failed to reset read timeout from datasource.", mlog.Err(err), mlog.String("src", dataSource))
+			return errors.Wrap(err, "failed to reset read timeout from datasource")
 		}
 	}
 
-	handle := SetupConnection("master", dataSource, ss.settings)
+	handle, err := SetupConnection("master", dataSource, ss.settings, DBPingAttempts)
+	if err != nil {
+		return err
+	}
 	ss.masterX = newSqlxDBWrapper(sqlx.NewDb(handle, ss.DriverName()),
 		time.Duration(*ss.settings.QueryTimeout)*time.Second,
 		*ss.settings.Trace)
@@ -310,34 +323,32 @@ func (ss *SqlStore) initConnection() {
 	}
 
 	if len(ss.settings.DataSourceReplicas) > 0 {
-		ss.ReplicaXs = make([]*sqlxDBWrapper, len(ss.settings.DataSourceReplicas))
+		ss.ReplicaXs = make([]*atomic.Pointer[sqlxDBWrapper], len(ss.settings.DataSourceReplicas))
 		for i, replica := range ss.settings.DataSourceReplicas {
-			handle := SetupConnection(fmt.Sprintf("replica-%v", i), replica, ss.settings)
-			ss.ReplicaXs[i] = newSqlxDBWrapper(sqlx.NewDb(handle, ss.DriverName()),
-				time.Duration(*ss.settings.QueryTimeout)*time.Second,
-				*ss.settings.Trace)
-			if ss.DriverName() == model.DatabaseDriverMysql {
-				ss.ReplicaXs[i].MapperFunc(noOpMapper)
-			}
-			if ss.metrics != nil {
-				ss.metrics.RegisterDBCollector(ss.ReplicaXs[i].DB.DB, "replica-"+strconv.Itoa(i))
+			ss.ReplicaXs[i] = &atomic.Pointer[sqlxDBWrapper]{}
+			handle, err = SetupConnection(fmt.Sprintf("replica-%v", i), replica, ss.settings, DBPingAttempts)
+			if err != nil {
+				// Initializing to be offline
+				ss.ReplicaXs[i].Store(&sqlxDBWrapper{isOnline: &atomic.Bool{}})
+				mlog.Warn("Failed to setup connection. Skipping..", mlog.String("db", fmt.Sprintf("replica-%v", i)), mlog.Err(err))
+				continue
 			}
+			ss.setDB(ss.ReplicaXs[i], handle, "replica-"+strconv.Itoa(i))
 		}
 	}
 
 	if len(ss.settings.DataSourceSearchReplicas) > 0 {
-		ss.searchReplicaXs = make([]*sqlxDBWrapper, len(ss.settings.DataSourceSearchReplicas))
+		ss.searchReplicaXs = make([]*atomic.Pointer[sqlxDBWrapper], len(ss.settings.DataSourceSearchReplicas))
 		for i, replica := range ss.settings.DataSourceSearchReplicas {
-			handle := SetupConnection(fmt.Sprintf("search-replica-%v", i), replica, ss.settings)
-			ss.searchReplicaXs[i] = newSqlxDBWrapper(sqlx.NewDb(handle, ss.DriverName()),
-				time.Duration(*ss.settings.QueryTimeout)*time.Second,
-				*ss.settings.Trace)
-			if ss.DriverName() == model.DatabaseDriverMysql {
-				ss.searchReplicaXs[i].MapperFunc(noOpMapper)
-			}
-			if ss.metrics != nil {
-				ss.metrics.RegisterDBCollector(ss.searchReplicaXs[i].DB.DB, "searchreplica-"+strconv.Itoa(i))
+			ss.searchReplicaXs[i] = &atomic.Pointer[sqlxDBWrapper]{}
+			handle, err = SetupConnection(fmt.Sprintf("search-replica-%v", i), replica, ss.settings, DBPingAttempts)
+			if err != nil {
+				// Initializing to be offline
+				ss.searchReplicaXs[i].Store(&sqlxDBWrapper{isOnline: &atomic.Bool{}})
+				mlog.Warn("Failed to setup connection. Skipping..", mlog.String("db", fmt.Sprintf("search-replica-%v", i)), mlog.Err(err))
+				continue
 			}
+			ss.setDB(ss.searchReplicaXs[i], handle, "searchreplica-"+strconv.Itoa(i))
 		}
 	}
 
@@ -347,9 +358,14 @@ func (ss *SqlStore) initConnection() {
 			if src.DataSource == nil {
 				continue
 			}
-			ss.replicaLagHandles[i] = SetupConnection(fmt.Sprintf(replicaLagPrefix+"-%d", i), *src.DataSource, ss.settings)
+			ss.replicaLagHandles[i], err = SetupConnection(fmt.Sprintf(replicaLagPrefix+"-%d", i), *src.DataSource, ss.settings, DBPingAttempts)
+			if err != nil {
+				mlog.Warn("Failed to setup replica lag handle. Skipping..", mlog.String("db", fmt.Sprintf(replicaLagPrefix+"-%d", i)), mlog.Err(err))
+				continue
+			}
 		}
 	}
+	return nil
 }
 
 func (ss *SqlStore) DriverName() string {
@@ -455,8 +471,15 @@ func (ss *SqlStore) GetSearchReplicaX() *sqlxDBWrapper {
 		return ss.GetReplicaX()
 	}
 
-	rrNum := atomic.AddInt64(&ss.srCounter, 1) % int64(len(ss.searchReplicaXs))
-	return ss.searchReplicaXs[rrNum]
+	for i := 0; i < len(ss.searchReplicaXs); i++ {
+		rrNum := atomic.AddInt64(&ss.srCounter, 1) % int64(len(ss.searchReplicaXs))
+		if ss.searchReplicaXs[rrNum].Load().Online() {
+			return ss.searchReplicaXs[rrNum].Load()
+		}
+	}
+
+	// If all search replicas are down, then go with replica.
+	return ss.GetReplicaX()
 }
 
 func (ss *SqlStore) GetReplicaX() *sqlxDBWrapper {
@@ -464,23 +487,64 @@ func (ss *SqlStore) GetReplicaX() *sqlxDBWrapper {
 		return ss.GetMasterX()
 	}
 
-	rrNum := atomic.AddInt64(&ss.rrCounter, 1) % int64(len(ss.ReplicaXs))
-	return ss.ReplicaXs[rrNum]
-}
-
-func (ss *SqlStore) GetInternalReplicaDBs() []*sql.DB {
-	if len(ss.settings.DataSourceReplicas) == 0 || ss.lockedToMaster || !ss.hasLicense() {
-		return []*sql.DB{
-			ss.GetMasterX().DB.DB,
+	for i := 0; i < len(ss.ReplicaXs); i++ {
+		rrNum := atomic.AddInt64(&ss.rrCounter, 1) % int64(len(ss.ReplicaXs))
+		if ss.ReplicaXs[rrNum].Load().Online() {
+			return ss.ReplicaXs[rrNum].Load()
 		}
 	}
 
-	dbs := make([]*sql.DB, len(ss.ReplicaXs))
-	for i, rx := range ss.ReplicaXs {
-		dbs[i] = rx.DB.DB
-	}
+	// If all replicas are down, then go with master.
+	return ss.GetMasterX()
+}
 
-	return dbs
+func (ss *SqlStore) monitorReplicas() {
+	t := time.NewTicker(time.Duration(*ss.settings.ReplicaMonitorIntervalSeconds) * time.Second)
+	defer func() {
+		t.Stop()
+		ss.wgMonitor.Done()
+	}()
+	for {
+		select {
+		case <-ss.quitMonitor:
+			return
+		case <-t.C:
+			setupReplica := func(r *atomic.Pointer[sqlxDBWrapper], dsn, name string) {
+				if r.Load().Online() {
+					return
+				}
+
+				handle, err := SetupConnection(name, dsn, ss.settings, 1)
+				if err != nil {
+					mlog.Warn("Failed to setup connection. Skipping..", mlog.String("db", name), mlog.Err(err))
+					return
+				}
+				if ss.metrics != nil && r.Load() != nil && r.Load().DB != nil {
+					ss.metrics.UnregisterDBCollector(r.Load().DB.DB, name)
+				}
+				ss.setDB(r, handle, name)
+			}
+			for i, replica := range ss.ReplicaXs {
+				setupReplica(replica, ss.settings.DataSourceReplicas[i], "replica-"+strconv.Itoa(i))
+			}
+
+			for i, replica := range ss.searchReplicaXs {
+				setupReplica(replica, ss.settings.DataSourceSearchReplicas[i], "search-replica-"+strconv.Itoa(i))
+			}
+		}
+	}
+}
+
+func (ss *SqlStore) setDB(replica *atomic.Pointer[sqlxDBWrapper], handle *dbsql.DB, name string) {
+	replica.Store(newSqlxDBWrapper(sqlx.NewDb(handle, ss.DriverName()),
+		time.Duration(*ss.settings.QueryTimeout)*time.Second,
+		*ss.settings.Trace))
+	if ss.DriverName() == model.DatabaseDriverMysql {
+		replica.Load().MapperFunc(noOpMapper)
+	}
+	if ss.metrics != nil {
+		ss.metrics.RegisterDBCollector(replica.Load().DB.DB, name)
+	}
 }
 
 func (ss *SqlStore) GetInternalReplicaDB() *sql.DB {
@@ -489,7 +553,7 @@ func (ss *SqlStore) GetInternalReplicaDB() *sql.DB {
 	}
 
 	rrNum := atomic.AddInt64(&ss.rrCounter, 1) % int64(len(ss.ReplicaXs))
-	return ss.ReplicaXs[rrNum].DB.DB
+	return ss.ReplicaXs[rrNum].Load().DB.DB
 }
 
 func (ss *SqlStore) TotalMasterDbConnections() int {
@@ -541,7 +605,10 @@ func (ss *SqlStore) TotalReadDbConnections() int {
 
 	count := 0
 	for _, db := range ss.ReplicaXs {
-		count = count + db.Stats().OpenConnections
+		if !db.Load().Online() {
+			continue
+		}
+		count = count + db.Load().Stats().OpenConnections
 	}
 
 	return count
@@ -554,7 +621,10 @@ func (ss *SqlStore) TotalSearchDbConnections() int {
 
 	count := 0
 	for _, db := range ss.searchReplicaXs {
-		count = count + db.Stats().OpenConnections
+		if !db.Load().Online() {
+			continue
+		}
+		count = count + db.Load().Stats().OpenConnections
 	}
 
 	return count
@@ -782,9 +852,14 @@ func IsUniqueConstraintError(err error, indexName []string) bool {
 }
 
 func (ss *SqlStore) GetAllConns() []*sqlxDBWrapper {
-	all := make([]*sqlxDBWrapper, len(ss.ReplicaXs)+1)
-	copy(all, ss.ReplicaXs)
-	all[len(ss.ReplicaXs)] = ss.masterX
+	all := make([]*sqlxDBWrapper, 0, len(ss.ReplicaXs)+1)
+	for i := range ss.ReplicaXs {
+		if !ss.ReplicaXs[i].Load().Online() {
+			continue
+		}
+		all = append(all, ss.ReplicaXs[i].Load())
+	}
+	all = append(all, ss.masterX)
 	return all
 }
 
@@ -807,11 +882,24 @@ func (ss *SqlStore) RecycleDBConnections(d time.Duration) {
 
 func (ss *SqlStore) Close() {
 	ss.masterX.Close()
+	// Closing monitor and waiting for it to be done.
+	// This needs to be done before closing the replica handles.
+	close(ss.quitMonitor)
+	ss.wgMonitor.Wait()
+
 	for _, replica := range ss.ReplicaXs {
-		replica.Close()
+		if replica.Load().Online() {
+			replica.Load().Close()
+		}
 	}
 
 	for _, replica := range ss.searchReplicaXs {
+		if replica.Load().Online() {
+			replica.Load().Close()
+		}
+	}
+
+	for _, replica := range ss.replicaLagHandles {
 		replica.Close()
 	}
 }
@@ -1132,7 +1220,10 @@ func (ss *SqlStore) migrate(direction migrationDirection) error {
 		if err != nil {
 			return err
 		}
-		db := SetupConnection("master", dataSource, ss.settings)
+		db, err2 := SetupConnection("master", dataSource, ss.settings, DBPingAttempts)
+		if err2 != nil {
+			return err2
+		}
 		driver, err = ms.WithInstance(db)
 		defer db.Close()
 	case model.DatabaseDriverPostgres: