balance: fix potential data race when setting config by adding a lock (#808) (#809)

Co-authored-by: djshow832 <[email protected]>

Author: ti-chi-bot

pkg/balance/factor/factor_balance.go

@@ -6,6 +6,7 @@ package factor
 import (
 	"sort"
 	"strconv"
+	"sync"
 	"time"
 
 	"github.com/pingcap/tiproxy/lib/config"
@@ -24,8 +25,8 @@ const (
 var _ policy.BalancePolicy = (*FactorBasedBalance)(nil)
 
 // FactorBasedBalance is the default balance policy.
-// It's not concurrency-safe for now.
 type FactorBasedBalance struct {
+	sync.Mutex
 	factors []Factor
 	// to reduce memory allocation
 	cachedList      []scoredBackend
@@ -53,6 +54,8 @@ func NewFactorBasedBalance(lg *zap.Logger, mr metricsreader.MetricsReader) *Fact
 // Init creates factors at the first time.
 // TODO: create factors according to config and update policy when config changes.
 func (fbb *FactorBasedBalance) Init(cfg *config.Config) {
+	fbb.Lock()
+	defer fbb.Unlock()
 	fbb.factors = make([]Factor, 0, 7)
 	fbb.setFactors(cfg)
 }
@@ -178,6 +181,8 @@ func (fbb *FactorBasedBalance) BackendToRoute(backends []policy.BackendCtx) poli
 		return nil
 	}
 
+	fbb.Lock()
+	defer fbb.Unlock()
 	scoredBackends := fbb.updateScore(backends)
 	if !fbb.canBeRouted(scoredBackends[0].scoreBits) {
 		return nil
@@ -240,6 +245,9 @@ func (fbb *FactorBasedBalance) BackendsToBalance(backends []policy.BackendCtx) (
 	if len(backends) <= 1 {
 		return
 	}
+
+	fbb.Lock()
+	defer fbb.Unlock()
 	scoredBackends := fbb.updateScore(backends)
 	if !fbb.canBeRouted(scoredBackends[0].scoreBits) {
 		return
@@ -304,10 +312,14 @@ func (fbb *FactorBasedBalance) canBeRouted(score uint64) bool {
 }
 
 func (fbb *FactorBasedBalance) SetConfig(cfg *config.Config) {
+	fbb.Lock()
+	defer fbb.Unlock()
 	fbb.setFactors(cfg)
 }
 
 func (fbb *FactorBasedBalance) Close() {
+	fbb.Lock()
+	defer fbb.Unlock()
 	for _, factor := range fbb.factors {
 		factor.Close()
 	}

pkg/balance/factor/factor_balance_test.go

@@ -4,12 +4,15 @@
 package factor
 
 import (
+	"context"
 	"slices"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/pingcap/tiproxy/lib/config"
 	"github.com/pingcap/tiproxy/lib/util/logger"
+	"github.com/pingcap/tiproxy/lib/util/waitgroup"
 	"github.com/pingcap/tiproxy/pkg/balance/policy"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
@@ -471,3 +474,34 @@ func TestCanBeRouted(t *testing.T) {
 		require.Equal(t, test.routed, b != nil, "test index %d", tIdx)
 	}
 }
+
+func TestSetFactorConcurrently(t *testing.T) {
+	fbb := NewFactorBasedBalance(zap.NewNop(), newMockMetricsReader())
+	var wg waitgroup.WaitGroup
+	cfg := &config.Config{}
+	fbb.Init(cfg)
+	wg.Add(10)
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	wg.Run(func() {
+		defer wg.Done()
+		policies := []string{config.BalancePolicyConnection, config.BalancePolicyResource, config.BalancePolicyLocation}
+		for i := 0; ctx.Err() != nil; i++ {
+			cfg.Balance = config.Balance{
+				Policy: policies[i%len(policies)],
+			}
+			fbb.SetConfig(cfg)
+		}
+	})
+	for i := 0; i < 9; i++ {
+		wg.Run(func() {
+			defer wg.Done()
+			for ctx.Err() != nil {
+				backends := createBackends(5)
+				fbb.BackendToRoute(backends)
+			}
+		})
+	}
+	wg.Wait()
+	cancel()
+	fbb.Close()
+}