balance: never route to backends with a different label (#805) (#807)

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

Author: ti-chi-bot

pkg/balance/factor/factor.go

@@ -19,5 +19,7 @@ type Factor interface {
 	// 0 indicates the factor is already balanced.
 	BalanceCount(from, to scoredBackend) (float64, []zap.Field)
 	SetConfig(cfg *config.Config)
+	// CanBeRouted returns whether a connection can be routed or migrated to the backend with the score.
+	CanBeRouted(score uint64) bool
 	Close()
 }

pkg/balance/factor/factor_balance.go

@@ -60,11 +60,6 @@ func (fbb *FactorBasedBalance) Init(cfg *config.Config) {
 func (fbb *FactorBasedBalance) setFactors(cfg *config.Config) {
 	fbb.factors = fbb.factors[:0]
 
-	if fbb.factorStatus == nil {
-		fbb.factorStatus = NewFactorStatus(fbb.lg.Named("status"))
-	}
-	fbb.factors = append(fbb.factors, fbb.factorStatus)
-
 	if cfg.Balance.LabelName != "" {
 		if fbb.factorLabel == nil {
 			fbb.factorLabel = NewFactorLabel()
@@ -75,6 +70,11 @@ func (fbb *FactorBasedBalance) setFactors(cfg *config.Config) {
 		fbb.factorLabel = nil
 	}
 
+	if fbb.factorStatus == nil {
+		fbb.factorStatus = NewFactorStatus(fbb.lg.Named("status"))
+	}
+	fbb.factors = append(fbb.factors, fbb.factorStatus)
+
 	switch cfg.Balance.Policy {
 	case config.BalancePolicyResource, config.BalancePolicyLocation:
 		if fbb.factorLocation == nil {
@@ -177,12 +177,16 @@ func (fbb *FactorBasedBalance) BackendToRoute(backends []policy.BackendCtx) poli
 	if len(backends) == 0 {
 		return nil
 	}
+
+	scoredBackends := fbb.updateScore(backends)
+	if !fbb.canBeRouted(scoredBackends[0].scoreBits) {
+		return nil
+	}
 	if len(backends) == 1 {
 		return backends[0]
 	}
 
 	var fields []zap.Field
-	scoredBackends := fbb.updateScore(backends)
 	for _, backend := range scoredBackends {
 		fields = append(fields, zap.String(backend.Addr(), strconv.FormatUint(backend.scoreBits, 16)))
 	}
@@ -237,8 +241,7 @@ func (fbb *FactorBasedBalance) BackendsToBalance(backends []policy.BackendCtx) (
 		return
 	}
 	scoredBackends := fbb.updateScore(backends)
-	// No need to migrate to unhealthy backends.
-	if !scoredBackends[0].Healthy() {
+	if !fbb.canBeRouted(scoredBackends[0].scoreBits) {
 		return
 	}
 	if scoredBackends[0].scoreBits == scoredBackends[len(scoredBackends)-1].scoreBits {
@@ -287,6 +290,19 @@ func (fbb *FactorBasedBalance) BackendsToBalance(backends []policy.BackendCtx) (
 	return
 }
 
+func (fbb *FactorBasedBalance) canBeRouted(score uint64) bool {
+	leftBitNum := fbb.totalBitNum
+	for _, factor := range fbb.factors {
+		bitNum := factor.ScoreBitNum()
+		score := score << (maxBitNum - leftBitNum) >> (maxBitNum - bitNum)
+		if !factor.CanBeRouted(score) {
+			return false
+		}
+		leftBitNum -= bitNum
+	}
+	return true
+}
+
 func (fbb *FactorBasedBalance) SetConfig(cfg *config.Config) {
 	fbb.setFactors(cfg)
 }

pkg/balance/factor/factor_balance_test.go

@@ -12,12 +12,13 @@ import (
 	"github.com/pingcap/tiproxy/lib/util/logger"
 	"github.com/pingcap/tiproxy/pkg/balance/policy"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 )
 
 func TestRouteWithOneFactor(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factor := &mockFactor{bitNum: 8, balanceCount: 1, threshold: 1}
+	factor := &mockFactor{bitNum: 8, balanceCount: 1, threshold: 1, canBeRouted: true}
 	fm.factors = []Factor{factor}
 	require.NoError(t, fm.updateBitNum())
 
@@ -71,8 +72,8 @@ func TestRouteWithOneFactor(t *testing.T) {
 func TestRouteWith2Factors(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factor1 := &mockFactor{bitNum: 4, balanceCount: 1, threshold: 1}
-	factor2 := &mockFactor{bitNum: 12, balanceCount: 1, threshold: 1}
+	factor1 := &mockFactor{bitNum: 4, balanceCount: 1, threshold: 1, canBeRouted: true}
+	factor2 := &mockFactor{bitNum: 12, balanceCount: 1, threshold: 1, canBeRouted: true}
 	fm.factors = []Factor{factor1, factor2}
 	require.NoError(t, fm.updateBitNum())
 
@@ -143,7 +144,7 @@ func TestRouteWith2Factors(t *testing.T) {
 func TestBalanceWithOneFactor(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factor := &mockFactor{bitNum: 2, balanceCount: 1}
+	factor := &mockFactor{bitNum: 2, balanceCount: 1, canBeRouted: true}
 	fm.factors = []Factor{factor}
 	require.NoError(t, fm.updateBitNum())
 
@@ -193,8 +194,8 @@ func TestBalanceWithOneFactor(t *testing.T) {
 func TestBalanceWith2Factors(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factor1 := &mockFactor{bitNum: 2, balanceCount: 2, threshold: 1}
-	factor2 := &mockFactor{bitNum: 12, balanceCount: 1}
+	factor1 := &mockFactor{bitNum: 2, balanceCount: 2, threshold: 1, canBeRouted: true}
+	factor2 := &mockFactor{bitNum: 12, balanceCount: 1, canBeRouted: true}
 	fm.factors = []Factor{factor1, factor2}
 	require.NoError(t, fm.updateBitNum())
 
@@ -280,7 +281,7 @@ func TestBalanceWith2Factors(t *testing.T) {
 func TestBalanceWith3Factors(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factors := []*mockFactor{{bitNum: 1}, {bitNum: 2}, {bitNum: 2}}
+	factors := []*mockFactor{{bitNum: 1, canBeRouted: true}, {bitNum: 2, canBeRouted: true}, {bitNum: 2, canBeRouted: true}}
 	fm.factors = []Factor{factors[0], factors[1], factors[2]}
 	require.NoError(t, fm.updateBitNum())
 
@@ -333,7 +334,7 @@ func TestBalanceWith3Factors(t *testing.T) {
 func TestScoreAlwaysUpdated(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factors := []*mockFactor{{bitNum: 1}, {bitNum: 2}}
+	factors := []*mockFactor{{bitNum: 1, canBeRouted: true}, {bitNum: 2, canBeRouted: true}}
 	fm.factors = []Factor{factors[0], factors[1]}
 	factors[0].updateScore = func(backends []scoredBackend) {
 		for i := 0; i < len(backends); i++ {
@@ -375,14 +376,14 @@ func TestSetFactors(t *testing.T) {
 			setFunc: func(balance *config.Balance) {
 				balance.LabelName = "group"
 			},
-			expectedNames: []string{"status", "label", "health", "memory", "cpu", "location", "conn"},
+			expectedNames: []string{"label", "status", "health", "memory", "cpu", "location", "conn"},
 		},
 		{
 			setFunc: func(balance *config.Balance) {
 				balance.Policy = config.BalancePolicyLocation
 				balance.LabelName = "group"
 			},
-			expectedNames: []string{"status", "label", "location", "health", "memory", "cpu", "conn"},
+			expectedNames: []string{"label", "status", "location", "health", "memory", "cpu", "conn"},
 		},
 		{
 			setFunc: func(balance *config.Balance) {
@@ -395,7 +396,7 @@ func TestSetFactors(t *testing.T) {
 				balance.Policy = config.BalancePolicyConnection
 				balance.LabelName = "group"
 			},
-			expectedNames: []string{"status", "label", "conn"},
+			expectedNames: []string{"label", "status", "conn"},
 		},
 	}
 
@@ -417,3 +418,56 @@ func TestSetFactors(t *testing.T) {
 	}
 	fm.Close()
 }
+
+func TestCanBeRouted(t *testing.T) {
+	fm := NewFactorBasedBalance(zap.NewNop(), newMockMetricsReader())
+	factor1 := &mockFactor{bitNum: 2, balanceCount: 1, canBeRouted: false}
+	factor2 := &mockFactor{bitNum: 2, balanceCount: 1, canBeRouted: true}
+	fm.factors = []Factor{factor1, factor2}
+	require.NoError(t, fm.updateBitNum())
+
+	tests := []struct {
+		scores1 []int
+		scores2 []int
+		routed  bool
+	}{
+		{
+			scores1: []int{1, 0},
+			scores2: []int{0, 0},
+			routed:  true,
+		},
+		{
+			scores1: []int{1, 1},
+			scores2: []int{1, 0},
+			routed:  false,
+		},
+		{
+			scores1: []int{0, 0},
+			scores2: []int{2, 1},
+			routed:  true,
+		},
+		{
+			scores1: []int{2, 1},
+			scores2: []int{0, 0},
+			routed:  false,
+		},
+	}
+	for tIdx, test := range tests {
+		factor1.updateScore = func(backends []scoredBackend) {
+			for i := 0; i < len(backends); i++ {
+				backends[i].addScore(test.scores1[i], factor1.bitNum)
+			}
+		}
+		factor2.updateScore = func(backends []scoredBackend) {
+			for i := 0; i < len(backends); i++ {
+				backends[i].addScore(test.scores2[i], factor2.bitNum)
+			}
+		}
+		backends := createBackends(len(test.scores1))
+		_, _, count, _, _ := fm.BackendsToBalance(backends)
+		require.Equal(t, test.routed, count > 0, "test index %d", tIdx)
+		backends = createBackends(len(test.scores1))
+		b := fm.BackendToRoute(backends)
+		require.Equal(t, test.routed, b != nil, "test index %d", tIdx)
+	}
+}

pkg/balance/factor/factor_conn.go

@@ -69,5 +69,9 @@ func (fcc *FactorConnCount) BalanceCount(from, to scoredBackend) (float64, []zap
 func (fcc *FactorConnCount) SetConfig(cfg *config.Config) {
 }
 
+func (fcc *FactorConnCount) CanBeRouted(_ uint64) bool {
+	return true
+}
+
 func (fcc *FactorConnCount) Close() {
 }

pkg/balance/factor/factor_cpu.go

@@ -272,6 +272,10 @@ func (fc *FactorCPU) BalanceCount(from, to scoredBackend) (float64, []zap.Field)
 func (fc *FactorCPU) SetConfig(cfg *config.Config) {
 }
 
+func (fc *FactorCPU) CanBeRouted(_ uint64) bool {
+	return true
+}
+
 func (fc *FactorCPU) Close() {
 	fc.mr.RemoveQueryExpr(fc.Name())
 }

pkg/balance/factor/factor_health.go

@@ -366,6 +366,10 @@ func (fh *FactorHealth) BalanceCount(from, to scoredBackend) (float64, []zap.Fie
 func (fh *FactorHealth) SetConfig(cfg *config.Config) {
 }
 
+func (fh *FactorHealth) CanBeRouted(_ uint64) bool {
+	return true
+}
+
 func (fh *FactorHealth) Close() {
 	for _, indicator := range fh.indicators {
 		fh.mr.RemoveQueryExpr(indicator.key)

pkg/balance/factor/factor_label.go

@@ -63,5 +63,10 @@ func (fl *FactorLabel) SetConfig(cfg *config.Config) {
 	}
 }
 
+func (fl *FactorLabel) CanBeRouted(score uint64) bool {
+	// Label is used to isolate business resources. Never route to the backends of another business.
+	return score == 0
+}
+
 func (fl *FactorLabel) Close() {
 }

pkg/balance/factor/factor_location.go

@@ -53,5 +53,9 @@ func (fl *FactorLocation) BalanceCount(from, to scoredBackend) (float64, []zap.F
 func (fl *FactorLocation) SetConfig(cfg *config.Config) {
 }
 
+func (fl *FactorLocation) CanBeRouted(_ uint64) bool {
+	return true
+}
+
 func (fl *FactorLocation) Close() {
 }

pkg/balance/factor/factor_memory.go

@@ -288,6 +288,10 @@ func (fm *FactorMemory) BalanceCount(from, to scoredBackend) (float64, []zap.Fie
 func (fm *FactorMemory) SetConfig(cfg *config.Config) {
 }
 
+func (fm *FactorMemory) CanBeRouted(_ uint64) bool {
+	return true
+}
+
 func (fm *FactorMemory) Close() {
 	fm.mr.RemoveQueryExpr(fm.Name())
 }

pkg/balance/factor/factor_status.go

@@ -107,5 +107,9 @@ func (fs *FactorStatus) BalanceCount(from, to scoredBackend) (float64, []zap.Fie
 func (fs *FactorStatus) SetConfig(cfg *config.Config) {
 }
 
+func (fl *FactorStatus) CanBeRouted(score uint64) bool {
+	return score == 0
+}
+
 func (fs *FactorStatus) Close() {
 }