Fixing behavior of Replicas and SearchReplicas in canary environments (#14576)

* Fixing behavior of Replicas and SearchReplicas in canary environments

* Trying to fix tests

* Revert "Trying to fix tests"

This reverts commit 3531da961844f5cb8557efcd3b570b06f362c6df.

* Revert "Fixing behavior of Replicas and SearchReplicas in canary environments"

This reverts commit 0c05901c843e4ccd60c8320fb4e0123b1bacf430.

* Revert "Disable read/search db replicas in TE/E0 (#14400)"

This reverts commit ef5ac519d9.

* Making the store aware of the license

* Readding the unit tests

* Fixing sqlstor supplier tests

* Adding mutex to ensure license write consistency and fixing tests

* Fixing tests

* Fixing tests

* Shuting down server properly during tests

* Trying to fix tests

* Trying to fix the tests

* Skipping flaky tests

Co-authored-by: mattermod <mattermod@users.noreply.github.com>

store/sqlstore/supplier.go

@@ -11,6 +11,7 @@ import (
 	"fmt"
 	"os"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -113,6 +114,8 @@ type SqlSupplier struct {
 	settings       *model.SqlSettings
 	lockedToMaster bool
 	context        context.Context
+	license        *model.License
+	licenseMutex   sync.Mutex
 }
 
 type TraceOnAdapter struct{}
@@ -327,6 +330,10 @@ func (ss *SqlSupplier) GetMaster() *gorp.DbMap {
 }
 
 func (ss *SqlSupplier) GetSearchReplica() *gorp.DbMap {
+	if ss.license == nil {
+		return ss.GetMaster()
+	}
+
 	if len(ss.settings.DataSourceSearchReplicas) == 0 {
 		return ss.GetReplica()
 	}
@@ -336,7 +343,7 @@ func (ss *SqlSupplier) GetSearchReplica() *gorp.DbMap {
 }
 
 func (ss *SqlSupplier) GetReplica() *gorp.DbMap {
-	if len(ss.settings.DataSourceReplicas) == 0 || ss.lockedToMaster {
+	if len(ss.settings.DataSourceReplicas) == 0 || ss.lockedToMaster || ss.license == nil {
 		return ss.GetMaster()
 	}
 
@@ -1178,6 +1185,12 @@ func (ss *SqlSupplier) CheckIntegrity() <-chan store.IntegrityCheckResult {
 	return results
 }
 
+func (ss *SqlSupplier) UpdateLicense(license *model.License) {
+	ss.licenseMutex.Lock()
+	defer ss.licenseMutex.Unlock()
+	ss.license = license
+}
+
 type mattermConverter struct{}
 
 func (me mattermConverter) ToDb(val interface{}) (interface{}, error) {

store/sqlstore/supplier_test.go

@@ -75,13 +75,14 @@ func TestGetReplica(t *testing.T) {
 
 	for _, testCase := range testCases {
 		testCase := testCase
-		t.Run(testCase.Description, func(t *testing.T) {
+		t.Run(testCase.Description+" with license", func(t *testing.T) {
 			t.Parallel()
 
 			settings := makeSqlSettings(model.DATABASE_DRIVER_SQLITE)
 			settings.DataSourceReplicas = testCase.DataSourceReplicas
 			settings.DataSourceSearchReplicas = testCase.DataSourceSearchReplicas
 			supplier := sqlstore.NewSqlSupplier(*settings, nil)
+			supplier.UpdateLicense(&model.License{})
 
 			replicas := make(map[*gorp.DbMap]bool)
 			for i := 0; i < 5; i++ {
@@ -130,6 +131,59 @@ func TestGetReplica(t *testing.T) {
 				}
 			}
 		})
+
+		t.Run(testCase.Description+" without license", func(t *testing.T) {
+			t.Parallel()
+
+			settings := makeSqlSettings(model.DATABASE_DRIVER_SQLITE)
+			settings.DataSourceReplicas = testCase.DataSourceReplicas
+			settings.DataSourceSearchReplicas = testCase.DataSourceSearchReplicas
+			supplier := sqlstore.NewSqlSupplier(*settings, nil)
+
+			replicas := make(map[*gorp.DbMap]bool)
+			for i := 0; i < 5; i++ {
+				replicas[supplier.GetReplica()] = true
+			}
+
+			searchReplicas := make(map[*gorp.DbMap]bool)
+			for i := 0; i < 5; i++ {
+				searchReplicas[supplier.GetSearchReplica()] = true
+			}
+
+			if len(testCase.DataSourceReplicas) > 0 {
+				// If replicas were defined, ensure none are the master.
+				assert.Len(t, replicas, 1)
+
+				for replica := range replicas {
+					assert.Same(t, supplier.GetMaster(), replica)
+				}
+
+			} else if assert.Len(t, replicas, 1) {
+				// Otherwise ensure the replicas contains only the master.
+				for replica := range replicas {
+					assert.Equal(t, supplier.GetMaster(), replica)
+				}
+			}
+
+			if len(testCase.DataSourceSearchReplicas) > 0 {
+				// If search replicas were defined, ensure none are the master nor the replicas.
+				assert.Len(t, searchReplicas, 1)
+
+				for searchReplica := range searchReplicas {
+					assert.Same(t, supplier.GetMaster(), searchReplica)
+				}
+
+			} else if len(testCase.DataSourceReplicas) > 0 {
+				// If no search replicas were defined, but replicas were, ensure they are equal.
+				assert.Equal(t, replicas, searchReplicas)
+
+			} else if assert.Len(t, searchReplicas, 1) {
+				// Otherwise ensure the search replicas contains the master.
+				for searchReplica := range searchReplicas {
+					assert.Equal(t, supplier.GetMaster(), searchReplica)
+				}
+			}
+		})
 	}
 }