factor, config: update load balance configs (#554)

pingcap · Jun 11, 2024 · de4a59b · de4a59b
1 parent 00fb1b3
commit de4a59b
Show file tree

Hide file tree

Showing 14 changed files with 695 additions and 728 deletions.
diff --git a/conf/proxy.toml b/conf/proxy.toml
@@ -115,12 +115,4 @@ graceful-close-conn-timeout = 15
 # ignore-wrong-namespace = true
 
 [balance]
-[balance.error]
-# enable = true
-[balance.memory]
-# enable = true
-[balance.cpu]
-# enable = true
-[balance.location]
-# enable = true
-# location-frist = true
+# policy = "resource"
diff --git a/lib/config/balance.go b/lib/config/balance.go
@@ -3,42 +3,33 @@
 
 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"`
-}
+import "github.com/pingcap/tiproxy/lib/util/errors"
 
-type ErrorBalance struct {
-	Enable bool `yaml:"enable,omitempty" toml:"enable,omitempty" json:"enable,omitempty"`
-}
+const (
+	BalancePolicyResource   = "resource"
+	BalancePolicyLocation   = "location"
+	BalancePolicyConnection = "connection"
+)
 
-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 Balance struct {
+	LabelName string `yaml:"label-name,omitempty" toml:"label-name,omitempty" json:"label-name,omitempty"`
+	Policy    string `yaml:"policy,omitempty" toml:"policy,omitempty" json:"policy,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 (b *Balance) Check() error {
+	switch b.Policy {
+	case BalancePolicyResource, BalancePolicyLocation, BalancePolicyConnection:
+		return nil
+	case "":
+		b.Policy = BalancePolicyResource
+	default:
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.policy")
+	}
+	return nil
 }
 
 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},
+		Policy: BalancePolicyResource,
 	}
 }
diff --git a/lib/config/proxy.go b/lib/config/proxy.go
@@ -183,6 +183,10 @@ func (cfg *Config) Check() error {
 		return errors.Wrapf(ErrInvalidConfigValue, "conn-buffer-size must be between 1K and 16M")
 	}
 
+	if err := cfg.Balance.Check(); err != nil {
+		return err
+	}
+
 	return nil
 }
 

diff --git a/pkg/balance/factor/factor_balance.go b/pkg/balance/factor/factor_balance.go
@@ -25,9 +25,9 @@ type FactorBasedBalance struct {
 	cachedList      []scoredBackend
 	mr              metricsreader.MetricsReader
 	lg              *zap.Logger
-	factorHealth    *FactorHealth
+	factorStatus    *FactorStatus
 	factorLabel     *FactorLabel
-	factorError     *FactorError
+	factorHealth    *FactorHealth
 	factorMemory    *FactorMemory
 	factorCPU       *FactorCPU
 	factorLocation  *FactorLocation
@@ -53,12 +53,12 @@ func (fbb *FactorBasedBalance) Init(cfg *config.Config) {
 func (fbb *FactorBasedBalance) setFactors(cfg *config.Config) {
 	fbb.factors = fbb.factors[:0]
 
-	if fbb.factorHealth == nil {
-		fbb.factorHealth = NewFactorHealth()
+	if fbb.factorStatus == nil {
+		fbb.factorStatus = NewFactorStatus()
 	}
-	fbb.factors = append(fbb.factors, fbb.factorHealth)
+	fbb.factors = append(fbb.factors, fbb.factorStatus)
 
-	if cfg.Balance.Label.Enable && cfg.Balance.Label.LabelName != "" {
+	if cfg.Balance.LabelName != "" {
 		if fbb.factorLabel == nil {
 			fbb.factorLabel = NewFactorLabel()
 		}
@@ -68,50 +68,44 @@ func (fbb *FactorBasedBalance) setFactors(cfg *config.Config) {
 		fbb.factorLabel = nil
 	}
 
-	if cfg.Balance.Location.Enable {
+	switch cfg.Balance.Policy {
+	case config.BalancePolicyResource, config.BalancePolicyLocation:
 		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)
+		if fbb.factorHealth == nil {
+			fbb.factorHealth = NewFactorHealth(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
+	default:
+		if fbb.factorLocation != nil {
+			fbb.factorLocation.Close()
+			fbb.factorLocation = nil
+		}
+		if fbb.factorHealth != nil {
+			fbb.factorHealth.Close()
+			fbb.factorHealth = nil
+		}
+		if fbb.factorMemory != nil {
+			fbb.factorMemory.Close()
+			fbb.factorMemory = nil
+		}
+		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)
+	switch cfg.Balance.Policy {
+	case config.BalancePolicyResource:
+		fbb.factors = append(fbb.factors, fbb.factorHealth, fbb.factorMemory, fbb.factorCPU, fbb.factorLocation)
+	case config.BalancePolicyLocation:
+		fbb.factors = append(fbb.factors, fbb.factorLocation, fbb.factorHealth, fbb.factorMemory, fbb.factorCPU)
 	}
 
 	if fbb.factorConnCount == nil {

diff --git a/pkg/balance/factor/factor_balance_test.go b/pkg/balance/factor/factor_balance_test.go
@@ -311,29 +311,39 @@ func TestSetFactors(t *testing.T) {
 	}{
 		{
 			setFunc:       func(balance *config.Balance) {},
-			expectedNames: []string{"health", "label", "error", "memory", "cpu", "location", "conn"},
+			expectedNames: []string{"status", "health", "memory", "cpu", "location", "conn"},
 		},
 		{
 			setFunc: func(balance *config.Balance) {
-				balance.Location.LocationFirst = true
+				balance.Policy = config.BalancePolicyLocation
 			},
-			expectedNames: []string{"health", "label", "location", "error", "memory", "cpu", "conn"},
+			expectedNames: []string{"status", "location", "health", "memory", "cpu", "conn"},
 		},
 		{
 			setFunc: func(balance *config.Balance) {
-				balance.Label.Enable = false
+				balance.LabelName = "group"
 			},
-			expectedNames: []string{"health", "error", "memory", "cpu", "location", "conn"},
+			expectedNames: []string{"status", "label", "health", "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
+				balance.Policy = config.BalancePolicyLocation
+				balance.LabelName = "group"
 			},
-			expectedNames: []string{"health", "conn"},
+			expectedNames: []string{"status", "label", "location", "health", "memory", "cpu", "conn"},
+		},
+		{
+			setFunc: func(balance *config.Balance) {
+				balance.Policy = config.BalancePolicyConnection
+			},
+			expectedNames: []string{"status", "conn"},
+		},
+		{
+			setFunc: func(balance *config.Balance) {
+				balance.Policy = config.BalancePolicyConnection
+				balance.LabelName = "group"
+			},
+			expectedNames: []string{"status", "label", "conn"},
 		},
 	}
 
@@ -342,22 +352,7 @@ func TestSetFactors(t *testing.T) {
 	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,
-				},
+				Policy: config.BalancePolicyResource,
 			},
 		}
 		fm.Init(cfg)