factor, config: add configs for factors (#552)

conf/proxy.toml

@@ -113,3 +113,14 @@ graceful-close-conn-timeout = 15
 [advance]
 
 # ignore-wrong-namespace = true
+
+[balance]
+[balance.error]
+# enable = true
+[balance.memory]
+# enable = true
+[balance.cpu]
+# enable = true
+[balance.location]
+# enable = true
+# location-frist = true

lib/config/balance.go

@@ -0,0 +1,44 @@
+// Copyright 2024 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package config
+
+type Balance struct {
+	Label    LabelBalance    `yaml:"label,omitempty" toml:"label,omitempty" json:"label,omitempty"`
+	Error    ErrorBalance    `yaml:"error,omitempty" toml:"error,omitempty" json:"error,omitempty"`
+	Memory   MemoryBalance   `yaml:"memory,omitempty" toml:"memory,omitempty" json:"memory,omitempty"`
+	CPU      CPUBalance      `yaml:"cpu,omitempty" toml:"cpu,omitempty" json:"cpu,omitempty"`
+	Location LocationBalance `yaml:"location,omitempty" toml:"location,omitempty" json:"location,omitempty"`
+}
+
+type LabelBalance struct {
+	Enable    bool   `yaml:"enable,omitempty" toml:"enable,omitempty" json:"enable,omitempty"`
+	LabelName string `yaml:"label-name,omitempty" toml:"label-name,omitempty" json:"label-name,omitempty"`
+}
+
+type ErrorBalance struct {
+	Enable bool `yaml:"enable,omitempty" toml:"enable,omitempty" json:"enable,omitempty"`
+}
+
+type MemoryBalance struct {
+	Enable bool `yaml:"enable,omitempty" toml:"enable,omitempty" json:"enable,omitempty"`
+}
+
+type CPUBalance struct {
+	Enable bool `yaml:"enable,omitempty" toml:"enable,omitempty" json:"enable,omitempty"`
+}
+
+type LocationBalance struct {
+	Enable        bool `yaml:"enable,omitempty" toml:"enable,omitempty" json:"enable,omitempty"`
+	LocationFirst bool `yaml:"location-first,omitempty" toml:"location-first,omitempty" json:"location-first,omitempty"`
+}
+
+func DefaultBalance() Balance {
+	return Balance{
+		Label:    LabelBalance{Enable: false},
+		Error:    ErrorBalance{Enable: true},
+		Memory:   MemoryBalance{Enable: true},
+		CPU:      CPUBalance{Enable: true},
+		Location: LocationBalance{Enable: true, LocationFirst: true},
+	}
+}

lib/config/proxy.go

@@ -116,14 +116,6 @@ type Security struct {
 	RequireBackendTLS bool      `yaml:"require-backend-tls,omitempty" toml:"require-backend-tls,omitempty" json:"require-backend-tls,omitempty"`
 }
 
-type Balance struct {
-	Label LabelBalance `yaml:"label,omitempty" toml:"label,omitempty" json:"label,omitempty"`
-}
-
-type LabelBalance struct {
-	LabelName string `yaml:"label-name,omitempty" toml:"label-name,omitempty" json:"label-name,omitempty"`
-}
-
 func DefaultKeepAlive() (frontend, backendHealthy, backendUnhealthy KeepAlive) {
 	frontend.Enabled = true
 	backendHealthy.Enabled = true
@@ -161,6 +153,8 @@ func NewConfig() *Config {
 	cfg.Security.ServerHTTPTLS.MinTLSVersion = "1.2"
 	cfg.Security.ClusterTLS.MinTLSVersion = "1.2"
 
+	cfg.Balance = DefaultBalance()
+
 	return &cfg
 }
 

pkg/balance/factor/factor.go

@@ -16,4 +16,5 @@ type Factor interface {
 	// 0 indicates the factor is already balanced.
 	BalanceCount(from, to scoredBackend) int
 	SetConfig(cfg *config.Config)
+	Close()
 }

pkg/balance/factor/factor_balance.go

@@ -22,10 +22,17 @@ var _ policy.BalancePolicy = (*FactorBasedBalance)(nil)
 type FactorBasedBalance struct {
 	factors []Factor
 	// to reduce memory allocation
-	cachedList  []scoredBackend
-	mr          metricsreader.MetricsReader
-	lg          *zap.Logger
-	totalBitNum int
+	cachedList      []scoredBackend
+	mr              metricsreader.MetricsReader
+	lg              *zap.Logger
+	factorHealth    *FactorHealth
+	factorLabel     *FactorLabel
+	factorError     *FactorError
+	factorMemory    *FactorMemory
+	factorCPU       *FactorCPU
+	factorLocation  *FactorLocation
+	factorConnCount *FactorConnCount
+	totalBitNum     int
 }
 
 func NewFactorBasedBalance(lg *zap.Logger, mr metricsreader.MetricsReader) *FactorBasedBalance {
@@ -39,20 +46,86 @@ 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.factors = []Factor{
-		NewFactorHealth(),
-		NewFactorLabel(),
-		NewFactorError(fbb.mr),
-		NewFactorMemory(fbb.mr),
-		NewFactorCPU(fbb.mr),
-		NewFactorLocation(),
-		NewFactorConnCount(),
+	fbb.factors = make([]Factor, 0, 7)
+	fbb.setFactors(cfg)
+}
+
+func (fbb *FactorBasedBalance) setFactors(cfg *config.Config) {
+	fbb.factors = fbb.factors[:0]
+
+	if fbb.factorHealth == nil {
+		fbb.factorHealth = NewFactorHealth()
+	}
+	fbb.factors = append(fbb.factors, fbb.factorHealth)
+
+	if cfg.Balance.Label.Enable && cfg.Balance.Label.LabelName != "" {
+		if fbb.factorLabel == nil {
+			fbb.factorLabel = NewFactorLabel()
+		}
+		fbb.factors = append(fbb.factors, fbb.factorLabel)
+	} else if fbb.factorLabel != nil {
+		fbb.factorLabel.Close()
+		fbb.factorLabel = nil
+	}
+
+	if cfg.Balance.Location.Enable {
+		if fbb.factorLocation == nil {
+			fbb.factorLocation = NewFactorLocation()
+		}
+	} else if fbb.factorLocation != nil {
+		fbb.factorLocation.Close()
+		fbb.factorLocation = nil
+	}
+	if cfg.Balance.Location.Enable && cfg.Balance.Location.LocationFirst {
+		fbb.factors = append(fbb.factors, fbb.factorLocation)
+	}
+
+	if cfg.Balance.Error.Enable {
+		if fbb.factorError == nil {
+			fbb.factorError = NewFactorError(fbb.mr)
+		}
+		fbb.factors = append(fbb.factors, fbb.factorError)
+	} else if fbb.factorError != nil {
+		fbb.factorError.Close()
+		fbb.factorError = nil
+	}
+
+	if cfg.Balance.Memory.Enable {
+		if fbb.factorMemory == nil {
+			fbb.factorMemory = NewFactorMemory(fbb.mr)
+		}
+		fbb.factors = append(fbb.factors, fbb.factorMemory)
+	} else if fbb.factorMemory != nil {
+		fbb.factorMemory.Close()
+		fbb.factorMemory = nil
+	}
+
+	if cfg.Balance.CPU.Enable {
+		if fbb.factorCPU == nil {
+			fbb.factorCPU = NewFactorCPU(fbb.mr)
+		}
+		fbb.factors = append(fbb.factors, fbb.factorCPU)
+	} else if fbb.factorCPU != nil {
+		fbb.factorCPU.Close()
+		fbb.factorCPU = nil
+	}
+
+	if cfg.Balance.Location.Enable && !cfg.Balance.Location.LocationFirst {
+		fbb.factors = append(fbb.factors, fbb.factorLocation)
+	}
+
+	if fbb.factorConnCount == nil {
+		fbb.factorConnCount = NewFactorConnCount()
 	}
+	fbb.factors = append(fbb.factors, fbb.factorConnCount)
+
 	err := fbb.updateBitNum()
 	if err != nil {
 		panic(err.Error())
 	}
-	fbb.SetConfig(cfg)
+	for _, factor := range fbb.factors {
+		factor.SetConfig(cfg)
+	}
 }
 
 func (fbb *FactorBasedBalance) updateBitNum() error {
@@ -172,7 +245,11 @@ func (fbb *FactorBasedBalance) BackendsToBalance(backends []policy.BackendCtx) (
 }
 
 func (fbb *FactorBasedBalance) SetConfig(cfg *config.Config) {
+	fbb.setFactors(cfg)
+}
+
+func (fbb *FactorBasedBalance) Close() {
 	for _, factor := range fbb.factors {
-		factor.SetConfig(cfg)
+		factor.Close()
 	}
 }

pkg/balance/factor/factor_balance_test.go

@@ -278,20 +278,6 @@ func TestBalanceWith3Factors(t *testing.T) {
 	}
 }
 
-func TestSetFactorConfig(t *testing.T) {
-	lg, _ := logger.CreateLoggerForTest(t)
-	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	factors := []*mockFactor{{bitNum: 1}, {bitNum: 2}, {bitNum: 2}}
-	fm.factors = []Factor{factors[0], factors[1], factors[2]}
-	cfg := &config.Config{
-		Labels: map[string]string{"k1": "v1"},
-	}
-	fm.SetConfig(cfg)
-	for _, factor := range factors {
-		require.Equal(t, "v1", factor.cfg.Labels["k1"])
-	}
-}
-
 // Even if the factor doesn't add a score, the score is still updated so that we can find the unbalanced factor
 // by locating the unbalanced bits.
 func TestScoreAlwaysUpdated(t *testing.T) {
@@ -317,3 +303,70 @@ func createBackends(num int) []policy.BackendCtx {
 	}
 	return backends
 }
+
+func TestSetFactors(t *testing.T) {
+	tests := []struct {
+		setFunc       func(balance *config.Balance)
+		expectedNames []string
+	}{
+		{
+			setFunc:       func(balance *config.Balance) {},
+			expectedNames: []string{"health", "label", "error", "memory", "cpu", "location", "conn"},
+		},
+		{
+			setFunc: func(balance *config.Balance) {
+				balance.Location.LocationFirst = true
+			},
+			expectedNames: []string{"health", "label", "location", "error", "memory", "cpu", "conn"},
+		},
+		{
+			setFunc: func(balance *config.Balance) {
+				balance.Label.Enable = false
+			},
+			expectedNames: []string{"health", "error", "memory", "cpu", "location", "conn"},
+		},
+		{
+			setFunc: func(balance *config.Balance) {
+				balance.Label.Enable = false
+				balance.Error.Enable = false
+				balance.Memory.Enable = false
+				balance.CPU.Enable = false
+				balance.Location.Enable = false
+			},
+			expectedNames: []string{"health", "conn"},
+		},
+	}
+
+	lg, _ := logger.CreateLoggerForTest(t)
+	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
+	for i, test := range tests {
+		cfg := &config.Config{
+			Balance: config.Balance{
+				Label: config.LabelBalance{
+					Enable:    true,
+					LabelName: "group",
+				},
+				Error: config.ErrorBalance{
+					Enable: true,
+				},
+				Memory: config.MemoryBalance{
+					Enable: true,
+				},
+				CPU: config.CPUBalance{
+					Enable: true,
+				},
+				Location: config.LocationBalance{
+					Enable: true,
+				},
+			},
+		}
+		fm.Init(cfg)
+		test.setFunc(&cfg.Balance)
+		fm.SetConfig(cfg)
+		require.Len(t, fm.factors, len(test.expectedNames), "test index %d", i)
+		for j := 0; j < len(fm.factors); j++ {
+			require.Equal(t, test.expectedNames[j], fm.factors[j].Name(), "test index %d", i)
+		}
+	}
+	fm.Close()
+}

pkg/balance/factor/factor_conn.go

@@ -49,3 +49,6 @@ func (fcc *FactorConnCount) BalanceCount(from, to scoredBackend) int {
 
 func (fcc *FactorConnCount) SetConfig(cfg *config.Config) {
 }
+
+func (fcc *FactorConnCount) Close() {
+}

pkg/balance/factor/factor_cpu.go

@@ -217,3 +217,7 @@ func (fc *FactorCPU) BalanceCount(from, to scoredBackend) int {
 
 func (fc *FactorCPU) SetConfig(cfg *config.Config) {
 }
+
+func (fc *FactorCPU) Close() {
+	fc.mr.RemoveQueryExpr(fc.queryID)
+}

pkg/balance/factor/factor_error.go

@@ -257,3 +257,9 @@ func (fe *FactorError) BalanceCount(from, to scoredBackend) int {
 
 func (fe *FactorError) SetConfig(cfg *config.Config) {
 }
+
+func (fe *FactorError) Close() {
+	for _, indicator := range fe.indicators {
+		fe.mr.RemoveQueryExpr(indicator.queryID)
+	}
+}

pkg/balance/factor/factor_health.go

@@ -52,5 +52,8 @@ func (fh *FactorHealth) BalanceCount(from, to scoredBackend) int {
 	return 1
 }
 
-func (fcc *FactorHealth) SetConfig(cfg *config.Config) {
+func (fh *FactorHealth) SetConfig(cfg *config.Config) {
+}
+
+func (fh *FactorHealth) Close() {
 }

pkg/balance/factor/factor_label.go

@@ -59,3 +59,6 @@ func (fl *FactorLabel) SetConfig(cfg *config.Config) {
 		fl.selfLabelVal = cfg.Labels[fl.labelName]
 	}
 }
+
+func (fl *FactorLabel) Close() {
+}

pkg/balance/factor/factor_location.go

@@ -58,3 +58,6 @@ func (fl *FactorLocation) SetConfig(cfg *config.Config) {
 		fl.selfLocation = cfg.Labels[locationLabelName]
 	}
 }
+
+func (fl *FactorLocation) Close() {
+}

pkg/balance/factor/factor_memory.go

@@ -217,3 +217,7 @@ func (fm *FactorMemory) BalanceCount(from, to scoredBackend) int {
 
 func (fm *FactorMemory) SetConfig(cfg *config.Config) {
 }
+
+func (fm *FactorMemory) Close() {
+	fm.mr.RemoveQueryExpr(fm.queryID)
+}

pkg/balance/factor/mock_test.go

@@ -82,6 +82,9 @@ func (mf *mockFactor) SetConfig(cfg *config.Config) {
 	mf.cfg = cfg
 }
 
+func (mf *mockFactor) Close() {
+}
+
 var _ metricsreader.MetricsReader = (*mockMetricsReader)(nil)
 
 type mockMetricsReader struct {