balance, config: add configs to tune the migration speed (#1032)

Author: djshow832

lib/config/balance.go

@@ -10,15 +10,26 @@ const (
 	BalancePolicyLocation   = "location"
 	BalancePolicyConnection = "connection"
 
-	RoutePolicyPreferIdle = "prefer-idle"
-	RoutePolicyRandom     = "random"
+	RoutingPolicyPreferIdle = "prefer-idle"
+	RoutingPolicyRandom     = "random"
+	RoutingPolicyIdlest     = "idlest"
 )
 
 type Balance struct {
-	LabelName   string `yaml:"label-name,omitempty" toml:"label-name,omitempty" json:"label-name,omitempty" reloadable:"true"`
-	RoutingRule string `yaml:"routing-rule,omitempty" toml:"routing-rule,omitempty" json:"routing-rule,omitempty" reloadable:"false"`
-	Policy      string `yaml:"policy,omitempty" toml:"policy,omitempty" json:"policy,omitempty" reloadable:"true"`
-	RoutePolicy string `yaml:"route-policy,omitempty" toml:"route-policy,omitempty" json:"route-policy,omitempty" reloadable:"true"`
+	LabelName     string `yaml:"label-name,omitempty" toml:"label-name,omitempty" json:"label-name,omitempty" reloadable:"true"`
+	RoutingRule   string `yaml:"routing-rule,omitempty" toml:"routing-rule,omitempty" json:"routing-rule,omitempty" reloadable:"false"`
+	Policy        string `yaml:"policy,omitempty" toml:"policy,omitempty" json:"policy,omitempty" reloadable:"true"`
+	RoutingPolicy string `yaml:"routing-policy,omitempty" toml:"routing-policy,omitempty" json:"routing-policy,omitempty" reloadable:"true"`
+	Status        Factor `yaml:"status,omitempty" toml:"status,omitempty" json:"status,omitempty" reloadable:"true"`
+	Health        Factor `yaml:"health,omitempty" toml:"health,omitempty" json:"health,omitempty" reloadable:"true"`
+	Memory        Factor `yaml:"memory,omitempty" toml:"memory,omitempty" json:"memory,omitempty" reloadable:"true"`
+	CPU           Factor `yaml:"cpu,omitempty" toml:"cpu,omitempty" json:"cpu,omitempty" reloadable:"true"`
+	Location      Factor `yaml:"location,omitempty" toml:"location,omitempty" json:"location,omitempty" reloadable:"true"`
+	ConnCount     Factor `yaml:"conn-count,omitempty" toml:"conn-count,omitempty" json:"conn-count,omitempty" reloadable:"true"`
+}
+
+type Factor struct {
+	MigrationsPerSecond float64 `yaml:"migrations-per-second,omitempty" toml:"migrations-per-second,omitempty" json:"migrations-per-second,omitempty" reloadable:"true"`
 }
 
 func (b *Balance) Check() error {
@@ -30,20 +41,38 @@ func (b *Balance) Check() error {
 	default:
 		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.policy")
 	}
-	switch b.RoutePolicy {
-	case RoutePolicyPreferIdle, RoutePolicyRandom:
+	switch b.RoutingPolicy {
+	case RoutingPolicyPreferIdle, RoutingPolicyRandom:
 		return nil
 	case "":
-		b.RoutePolicy = RoutePolicyPreferIdle
+		b.RoutingPolicy = RoutingPolicyPreferIdle
 	default:
 		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.route-policy")
 	}
+	if b.Status.MigrationsPerSecond < 0 {
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.status.migrations-per-second")
+	}
+	if b.Health.MigrationsPerSecond < 0 {
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.health.migrations-per-second")
+	}
+	if b.Memory.MigrationsPerSecond < 0 {
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.memory.migrations-per-second")
+	}
+	if b.CPU.MigrationsPerSecond < 0 {
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.cpu.migrations-per-second")
+	}
+	if b.Location.MigrationsPerSecond < 0 {
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.location.migrations-per-second")
+	}
+	if b.ConnCount.MigrationsPerSecond < 0 {
+		return errors.Wrapf(ErrInvalidConfigValue, "invalid balance.conn-count.migrations-per-second")
+	}
 	return nil
 }
 
 func DefaultBalance() Balance {
 	return Balance{
-		Policy:      BalancePolicyResource,
-		RoutePolicy: RoutePolicyPreferIdle,
+		Policy:        BalancePolicyResource,
+		RoutingPolicy: RoutingPolicyPreferIdle,
 	}
 }

pkg/balance/factor/factor_balance.go

@@ -193,13 +193,32 @@ func (fbb *FactorBasedBalance) BackendToRoute(backends []policy.BackendCtx) poli
 	}
 
 	switch fbb.routePolicy {
-	case config.RoutePolicyRandom:
+	case config.RoutingPolicyRandom:
 		return fbb.routeRandom(scoredBackends, &fields)
+	case config.RoutingPolicyIdlest:
+		return fbb.routeIdlest(scoredBackends, &fields)
 	default:
 		return fbb.routePreferIdle(scoredBackends, &fields)
 	}
 }
 
+func (fbb *FactorBasedBalance) routeIdlest(scoredBackends []scoredBackend, fields *[]zap.Field) policy.BackendCtx {
+	// Evict the backends that are can't be routed to, and then choose the idlest one.
+	// It's like least-connection algorithm.
+	idx := -1
+	for i := range scoredBackends {
+		if fbb.canBeRouted(scoredBackends[i].scoreBits) {
+			idx = i
+			break
+		}
+	}
+	if idx < 0 {
+		return nil
+	}
+	*fields = append(*fields, zap.String("target", scoredBackends[idx].Addr()))
+	return scoredBackends[idx].BackendCtx
+}
+
 func (fbb *FactorBasedBalance) routeRandom(scoredBackends []scoredBackend, fields *[]zap.Field) policy.BackendCtx {
 	// Evict the backends that are can't be routed to, and then choose one randomly, with the idlest one having 10% higher priority.
 	// It works when new connections are more intensive. If new connections are all routed to the newly scaled backend,
@@ -370,7 +389,7 @@ func (fbb *FactorBasedBalance) SetConfig(cfg *config.Config) {
 	fbb.Lock()
 	defer fbb.Unlock()
 	fbb.setFactors(cfg)
-	fbb.routePolicy = cfg.Balance.RoutePolicy
+	fbb.routePolicy = cfg.Balance.RoutingPolicy
 }
 
 func (fbb *FactorBasedBalance) Close() {

pkg/balance/factor/factor_balance_test.go

@@ -33,62 +33,67 @@ func TestRouteWithOneFactor(t *testing.T) {
 		{
 			scores:   []int{10},
 			idxRange: []int{0},
-			policy:   config.RoutePolicyRandom,
+			policy:   config.RoutingPolicyRandom,
 		},
 		{
 			scores:   []int{10, 20},
 			idxRange: []int{0, 1},
-			policy:   config.RoutePolicyRandom,
+			policy:   config.RoutingPolicyRandom,
 		},
 		{
 			scores:   []int{10, 20, 30},
 			idxRange: []int{0, 1, 2},
-			policy:   config.RoutePolicyRandom,
+			policy:   config.RoutingPolicyRandom,
 		},
 		{
 			scores:   []int{30, 20, 10},
 			idxRange: []int{0, 1, 2},
-			policy:   config.RoutePolicyRandom,
+			policy:   config.RoutingPolicyRandom,
 		},
 		{
 			scores:   []int{30, 11, 10},
 			idxRange: []int{0, 1, 2},
-			policy:   config.RoutePolicyRandom,
+			policy:   config.RoutingPolicyRandom,
 		},
 		{
 			scores:   []int{11, 11, 10},
 			idxRange: []int{0, 1, 2},
-			policy:   config.RoutePolicyRandom,
+			policy:   config.RoutingPolicyRandom,
 		},
 		{
 			scores:   []int{10},
 			idxRange: []int{0},
-			policy:   config.RoutePolicyPreferIdle,
+			policy:   config.RoutingPolicyPreferIdle,
 		},
 		{
 			scores:   []int{10, 20},
 			idxRange: []int{0},
-			policy:   config.RoutePolicyPreferIdle,
+			policy:   config.RoutingPolicyPreferIdle,
 		},
 		{
 			scores:   []int{10, 20, 30},
 			idxRange: []int{0},
-			policy:   config.RoutePolicyPreferIdle,
+			policy:   config.RoutingPolicyPreferIdle,
 		},
 		{
 			scores:   []int{30, 20, 10},
 			idxRange: []int{2},
-			policy:   config.RoutePolicyPreferIdle,
+			policy:   config.RoutingPolicyPreferIdle,
 		},
 		{
 			scores:   []int{30, 11, 10},
 			idxRange: []int{1, 2},
-			policy:   config.RoutePolicyPreferIdle,
+			policy:   config.RoutingPolicyPreferIdle,
 		},
 		{
 			scores:   []int{11, 11, 10},
 			idxRange: []int{0, 1, 2},
-			policy:   config.RoutePolicyPreferIdle,
+			policy:   config.RoutingPolicyPreferIdle,
+		},
+		{
+			scores:   []int{10, 10, 9},
+			idxRange: []int{2},
+			policy:   config.RoutingPolicyIdlest,
 		},
 		{
 			scores:   []int{10, 20},
@@ -114,10 +119,10 @@ func TestRouteWithOneFactor(t *testing.T) {
 	}
 }
 
-func TestRouteIdlestWith2Factors(t *testing.T) {
+func TestRoutePreferIdleWith2Factors(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	fm.routePolicy = config.RoutePolicyPreferIdle
+	fm.routePolicy = config.RoutingPolicyPreferIdle
 	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}
@@ -189,7 +194,7 @@ func TestRouteIdlestWith2Factors(t *testing.T) {
 func TestRouteRandomWith2Factors(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
-	fm.routePolicy = config.RoutePolicyRandom
+	fm.routePolicy = config.RoutingPolicyRandom
 	factor1 := &mockFactor{bitNum: 4, balanceCount: 1, threshold: 1, canBeRouted: false}
 	factor2 := &mockFactor{bitNum: 12, balanceCount: 1, threshold: 1, canBeRouted: true}
 	fm.factors = []Factor{factor1, factor2}
@@ -253,6 +258,58 @@ func TestRouteRandomWith2Factors(t *testing.T) {
 	}
 }
 
+func TestRouteIdlestWith2Factors(t *testing.T) {
+	lg, _ := logger.CreateLoggerForTest(t)
+	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
+	fm.routePolicy = config.RoutingPolicyIdlest
+	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())
+	tests := []struct {
+		scores1  []int
+		scores2  []int
+		idxRange []int
+	}{
+		{
+			scores1:  []int{10, 0, 0},
+			scores2:  []int{0, 100, 200},
+			idxRange: []int{1},
+		},
+		{
+			scores1:  []int{10, 10, 0},
+			scores2:  []int{0, 100, 200},
+			idxRange: []int{2},
+		},
+		{
+			scores1:  []int{10, 10, 10},
+			scores2:  []int{0, 100, 200},
+			idxRange: []int{0},
+		},
+	}
+	for tIdx, test := range tests {
+		factor1.updateScore = func(backends []scoredBackend) {
+			for i := range backends {
+				backends[i].addScore(test.scores1[i], factor1.bitNum)
+			}
+		}
+		factor2.updateScore = func(backends []scoredBackend) {
+			for i := range backends {
+				backends[i].addScore(test.scores2[i], factor2.bitNum)
+			}
+		}
+		backends := createBackends(len(test.scores1))
+		targets := make(map[int]struct{}, len(test.scores1))
+		for i := 0; len(targets) < len(test.idxRange) || i < 100; i++ {
+			require.Less(t, i, 100000, "test index %d", tIdx)
+			backend := fm.BackendToRoute(backends)
+			idx := slices.Index(backends, backend)
+			require.Contains(t, test.idxRange, idx, "test index %d", tIdx)
+			targets[idx] = struct{}{}
+		}
+	}
+}
+
 func TestBalanceWithOneFactor(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	fm := NewFactorBasedBalance(lg, newMockMetricsReader())
@@ -639,11 +696,11 @@ func TestSetConfigsConcurrently(t *testing.T) {
 	wg.Run(func() {
 		defer wg.Done()
 		policies := []string{config.BalancePolicyConnection, config.BalancePolicyResource, config.BalancePolicyLocation}
-		routePolicies := []string{config.RoutePolicyRandom, config.RoutePolicyPreferIdle}
+		routePolicies := []string{config.RoutingPolicyRandom, config.RoutingPolicyPreferIdle, config.RoutingPolicyIdlest}
 		for i := 0; ctx.Err() != nil; i++ {
 			cfg.Balance = config.Balance{
-				Policy:      policies[i%len(policies)],
-				RoutePolicy: routePolicies[i%len(routePolicies)],
+				Policy:        policies[i%len(policies)],
+				RoutingPolicy: routePolicies[i%len(routePolicies)],
 			}
 			fbb.SetConfig(cfg)
 		}

pkg/balance/factor/factor_conn.go

@@ -25,7 +25,8 @@ var _ Factor = (*FactorConnCount)(nil)
 // - The backend CPU usages are unavailable
 // - The workload just starts and the backend CPU usages are low
 type FactorConnCount struct {
-	bitNum int
+	bitNum              int
+	migrationsPerSecond float64
 }
 
 func NewFactorConnCount() *FactorConnCount {
@@ -58,6 +59,9 @@ func (fcc *FactorConnCount) BalanceCount(from, to scoredBackend) (BalanceAdvice,
 	if float64(from.ConnScore()) <= float64(to.ConnScore()+1)*connBalancedRatio {
 		return AdviceNeutral, 0, nil
 	}
+	if fcc.migrationsPerSecond > 0 {
+		return AdvicePositive, fcc.migrationsPerSecond, nil
+	}
 	targetTo := float64(from.ConnScore()+to.ConnScore()+1) / (1 + connBalancedRatio)
 	count := (targetTo - float64(to.ConnScore()+1)) / balanceSeconds4Conn
 	if count < 0 {
@@ -67,6 +71,7 @@ func (fcc *FactorConnCount) BalanceCount(from, to scoredBackend) (BalanceAdvice,
 }
 
 func (fcc *FactorConnCount) SetConfig(cfg *config.Config) {
+	fcc.migrationsPerSecond = cfg.Balance.ConnCount.MigrationsPerSecond
 }
 
 func (fcc *FactorConnCount) CanBeRouted(_ uint64) bool {

pkg/balance/factor/factor_conn_test.go

@@ -6,6 +6,7 @@ package factor
 import (
 	"testing"
 
+	"github.com/pingcap/tiproxy/lib/config"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 )
@@ -94,3 +95,36 @@ func TestFactorConnSpeed(t *testing.T) {
 		require.LessOrEqual(t, backend2.connScore, test.targetRange[1], "case id: %d", i)
 	}
 }
+
+func TestFactorConnConfig(t *testing.T) {
+	tests := []struct {
+		score1 int
+		score2 int
+		speed  float64
+	}{
+		{
+			score1: 120,
+			score2: 100,
+			speed:  0,
+		},
+		{
+			score1: 150,
+			score2: 100,
+			speed:  10,
+		},
+	}
+
+	factor := NewFactorConnCount()
+	factor.SetConfig(&config.Config{Balance: config.Balance{ConnCount: config.Factor{MigrationsPerSecond: 10}}})
+	require.EqualValues(t, 10, factor.migrationsPerSecond)
+	backend1 := newMockBackend(true, 0)
+	backend2 := newMockBackend(true, 0)
+	scoredBackend1 := newScoredBackend(backend1, zap.NewNop())
+	scoredBackend2 := newScoredBackend(backend2, zap.NewNop())
+	for i, test := range tests {
+		backend1.connScore = test.score1
+		backend2.connScore = test.score2
+		_, balanceCount, _ := factor.BalanceCount(scoredBackend1, scoredBackend2)
+		require.EqualValues(t, balanceCount, test.speed, "case id: %d", i)
+	}
+}

pkg/balance/factor/factor_cpu.go

@@ -94,10 +94,11 @@ type FactorCPU struct {
 	// The updated time of the metric that we've read last time.
 	lastMetricTime time.Time
 	// The estimated average CPU usage used by one connection.
-	usagePerConn float64
-	mr           metricsreader.MetricsReader
-	bitNum       int
-	lg           *zap.Logger
+	usagePerConn        float64
+	mr                  metricsreader.MetricsReader
+	bitNum              int
+	migrationsPerSecond float64
+	lg                  *zap.Logger
 }
 
 func NewFactorCPU(mr metricsreader.MetricsReader, lg *zap.Logger) *FactorCPU {
@@ -278,10 +279,14 @@ func (fc *FactorCPU) BalanceCount(from, to scoredBackend) (BalanceAdvice, float6
 	if 1.3-toAvgUsage < (1.3-fromAvgUsage)*cpuBalancedRatio || 1.3-toLatestUsage < (1.3-fromLatestUsage)*cpuBalancedRatio {
 		return AdviceNeutral, 0, fields
 	}
+	if fc.migrationsPerSecond > 0 {
+		return AdvicePositive, fc.migrationsPerSecond, fields
+	}
 	return AdvicePositive, 1 / fc.usagePerConn / balanceRatio4Cpu, fields
 }
 
 func (fc *FactorCPU) SetConfig(cfg *config.Config) {
+	fc.migrationsPerSecond = cfg.Balance.CPU.MigrationsPerSecond
 }
 
 func (fc *FactorCPU) CanBeRouted(_ uint64) bool {