placement: fix the issue cannot create dict follower constraints (#47750

)

close #47741

pkg/ddl/placement/bundle.go

@@ -67,8 +67,8 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle,
 	leaderConst := options.LeaderConstraints
 	learnerConstraints := options.LearnerConstraints
 	followerConstraints := options.FollowerConstraints
-	followerCount := options.Followers
-	learnerCount := options.Learners
+	explicitFollowerCount := options.Followers
+	explicitLearnerCount := options.Learners
 
 	rules := []*Rule{}
 	commonConstraints, err := NewConstraintsFromYaml([]byte(constraints))
@@ -77,7 +77,10 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle,
 		// The dictionary format specifies details for each replica. Constraints are used to define normal
 		// replicas that should act as voters.
 		// For example: CONSTRAINTS='{ "+region=us-east-1":2, "+region=us-east-2": 2, "+region=us-west-1": 1}'
-		normalReplicasRules, err := NewRulesWithDictConstraints(Voter, constraints)
+		normalReplicasRules, err := NewRuleBuilder().
+			SetRole(Voter).
+			SetConstraintStr(constraints).
+			BuildRulesWithDictConstraintsOnly()
 		if err != nil {
 			return nil, err
 		}
@@ -94,15 +97,15 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle,
 		}
 	}
 	leaderReplicas, followerReplicas := uint64(1), uint64(2)
-	if followerCount > 0 {
-		followerReplicas = followerCount
+	if explicitFollowerCount > 0 {
+		followerReplicas = explicitFollowerCount
 	}
 	if !needCreateDefault {
 		if len(leaderConstraints) == 0 {
 			leaderReplicas = 0
 		}
 		if len(followerConstraints) == 0 {
-			if followerCount > 0 {
+			if explicitFollowerCount > 0 {
 				return nil, fmt.Errorf("%w: specify follower count without specify follower constraints when specify other constraints", ErrInvalidPlacementOptions)
 			}
 			followerReplicas = 0
@@ -119,7 +122,12 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle,
 	// create follower rules.
 	// if no constraints, we need create default follower rules.
 	if followerReplicas > 0 {
-		followerRules, err := NewRules(Voter, followerReplicas, followerConstraints)
+		builder := NewRuleBuilder().
+			SetRole(Voter).
+			SetReplicasNum(followerReplicas).
+			SetSkipCheckReplicasConsistent(needCreateDefault && (explicitFollowerCount == 0)).
+			SetConstraintStr(followerConstraints)
+		followerRules, err := builder.BuildRules()
 		if err != nil {
 			return nil, fmt.Errorf("%w: invalid FollowerConstraints", err)
 		}
@@ -134,7 +142,11 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle,
 	}
 
 	// create learner rules.
-	learnerRules, err := NewRules(Learner, learnerCount, learnerConstraints)
+	builder := NewRuleBuilder().
+		SetRole(Learner).
+		SetReplicasNum(explicitLearnerCount).
+		SetConstraintStr(learnerConstraints)
+	learnerRules, err := builder.BuildRules()
 	if err != nil {
 		return nil, fmt.Errorf("%w: invalid LearnerConstraints", err)
 	}

pkg/ddl/placement/bundle_test.go

@@ -342,9 +342,9 @@ func TestString(t *testing.T) {
 		ID: GroupID(1),
 	}
 
-	rules1, err := NewRules(Voter, 3, `["+zone=sh", "+zone=sh"]`)
+	rules1, err := newRules(Voter, 3, `["+zone=sh", "+zone=sh"]`)
 	require.NoError(t, err)
-	rules2, err := NewRules(Voter, 4, `["-zone=sh", "+zone=bj"]`)
+	rules2, err := newRules(Voter, 4, `["-zone=sh", "+zone=bj"]`)
 	require.NoError(t, err)
 	bundle.Rules = append(rules1, rules2...)
 
@@ -727,6 +727,26 @@ func TestNewBundleFromOptions(t *testing.T) {
 		err: ErrInvalidConstraintsFormat,
 	})
 
+	tests = append(tests, TestCase{
+		name: "direct syntax: follower dict constraints",
+		input: &model.PlacementSettings{
+			FollowerConstraints: "{+disk=ssd: 1}",
+		},
+		output: []*Rule{
+			NewRule(Leader, 1, NewConstraintsDirect()),
+			NewRule(Voter, 1, NewConstraintsDirect(NewConstraintDirect("disk", In, "ssd"))),
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name: "direct syntax: invalid follower dict constraints",
+		input: &model.PlacementSettings{
+			FollowerConstraints: "{+disk=ssd: 1}",
+			Followers:           2,
+		},
+		err: ErrInvalidConstraintsReplicas,
+	})
+
 	tests = append(tests, TestCase{
 		name: "direct syntax: learner dict constraints",
 		input: &model.PlacementSettings{
@@ -799,7 +819,7 @@ func TestResetBundleWithSingleRule(t *testing.T) {
 		ID: GroupID(1),
 	}
 
-	rules, err := NewRules(Voter, 3, `["+zone=sh", "+zone=sh"]`)
+	rules, err := newRules(Voter, 3, `["+zone=sh", "+zone=sh"]`)
 	require.NoError(t, err)
 	bundle.Rules = rules
 
@@ -916,15 +936,15 @@ func TestTidy(t *testing.T) {
 		ID: GroupID(1),
 	}
 
-	rules0, err := NewRules(Voter, 1, `["+zone=sh", "+zone=sh"]`)
+	rules0, err := newRules(Voter, 1, `["+zone=sh", "+zone=sh"]`)
 	require.NoError(t, err)
 	require.Len(t, rules0, 1)
 	rules0[0].Count = 0 // test prune useless rules
 
-	rules1, err := NewRules(Voter, 4, `["-zone=sh", "+zone=bj"]`)
+	rules1, err := newRules(Voter, 4, `["-zone=sh", "+zone=bj"]`)
 	require.NoError(t, err)
 	require.Len(t, rules1, 1)
-	rules2, err := NewRules(Voter, 0, `{"-zone=sh,+zone=bj": 4}}`)
+	rules2, err := newRules(Voter, 0, `{"-zone=sh,+zone=bj": 4}}`)
 	require.NoError(t, err)
 	bundle.Rules = append(bundle.Rules, rules0...)
 	bundle.Rules = append(bundle.Rules, rules1...)
@@ -943,11 +963,11 @@ func TestTidy(t *testing.T) {
 	}, bundle.Rules[0].Constraints[2])
 
 	// merge
-	rules3, err := NewRules(Follower, 4, "")
+	rules3, err := newRules(Follower, 4, "")
 	require.NoError(t, err)
 	require.Len(t, rules3, 1)
 
-	rules4, err := NewRules(Follower, 5, "")
+	rules4, err := newRules(Follower, 5, "")
 	require.NoError(t, err)
 	require.Len(t, rules4, 1)
 

pkg/ddl/placement/rule.go

@@ -153,6 +153,70 @@ func (r *TiFlashRule) UnmarshalJSON(bytes []byte) error {
 	return err
 }
 
+// RuleBuilder is used to build the Rules from a constraint string.
+type RuleBuilder struct {
+	role                        PeerRoleType
+	replicasNum                 uint64
+	skipCheckReplicasConsistent bool
+	constraintStr               string
+}
+
+// NewRuleBuilder creates a new RuleBuilder.
+func NewRuleBuilder() *RuleBuilder {
+	return &RuleBuilder{}
+}
+
+// SetRole sets the role of the rule.
+func (b *RuleBuilder) SetRole(role PeerRoleType) *RuleBuilder {
+	b.role = role
+	return b
+}
+
+// SetReplicasNum sets the replicas number in the rule.
+func (b *RuleBuilder) SetReplicasNum(num uint64) *RuleBuilder {
+	b.replicasNum = num
+	return b
+}
+
+// SetSkipCheckReplicasConsistent sets the skipCheckReplicasConsistent flag.
+func (b *RuleBuilder) SetSkipCheckReplicasConsistent(skip bool) *RuleBuilder {
+	b.skipCheckReplicasConsistent = skip
+	return b
+}
+
+// SetConstraintStr sets the constraint string.
+func (b *RuleBuilder) SetConstraintStr(constraintStr string) *RuleBuilder {
+	b.constraintStr = constraintStr
+	return b
+}
+
+// BuildRulesWithDictConstraintsOnly constructs []*Rule from a yaml-compatible representation of
+// 'dict' constraints.
+func (b *RuleBuilder) BuildRulesWithDictConstraintsOnly() ([]*Rule, error) {
+	return newRulesWithDictConstraints(b.role, b.constraintStr)
+}
+
+// BuildRules constructs []*Rule from a yaml-compatible representation of
+// 'array' or 'dict' constraints.
+// Refer to https://github.com/pingcap/tidb/blob/master/docs/design/2020-06-24-placement-rules-in-sql.md.
+func (b *RuleBuilder) BuildRules() ([]*Rule, error) {
+	rules, err := newRules(b.role, b.replicasNum, b.constraintStr)
+	// check if replicas is consistent
+	if err == nil {
+		if b.skipCheckReplicasConsistent {
+			return rules, err
+		}
+		totalCnt := 0
+		for _, rule := range rules {
+			totalCnt += rule.Count
+		}
+		if b.replicasNum != 0 && b.replicasNum != uint64(totalCnt) {
+			err = fmt.Errorf("%w: count of replicas in dict constrains is %d, but got %d", ErrInvalidConstraintsReplicas, totalCnt, b.replicasNum)
+		}
+	}
+	return rules, err
+}
+
 // NewRule constructs *Rule from role, count, and constraints. It is here to
 // consistent the behavior of creating new rules.
 func NewRule(role PeerRoleType, replicas uint64, cnst Constraints) *Rule {
@@ -175,10 +239,10 @@ func getYamlMapFormatError(str string) error {
 	return nil
 }
 
-// NewRules constructs []*Rule from a yaml-compatible representation of
+// newRules constructs []*Rule from a yaml-compatible representation of
 // 'array' or 'dict' constraints.
 // Refer to https://github.com/pingcap/tidb/blob/master/docs/design/2020-06-24-placement-rules-in-sql.md.
-func NewRules(role PeerRoleType, replicas uint64, cnstr string) (rules []*Rule, err error) {
+func newRules(role PeerRoleType, replicas uint64, cnstr string) (rules []*Rule, err error) {
 	cnstbytes := []byte(cnstr)
 	constraints1, err1 := NewConstraintsFromYaml(cnstbytes)
 	if err1 == nil {
@@ -199,23 +263,12 @@ func NewRules(role PeerRoleType, replicas uint64, cnstr string) (rules []*Rule,
 		return
 	}
 
-	rules, err = NewRulesWithDictConstraints(role, cnstr)
-	// check if replicas is consistent
-	if err == nil {
-		totalCnt := 0
-		for _, rule := range rules {
-			totalCnt += rule.Count
-		}
-		if replicas != 0 && replicas != uint64(totalCnt) {
-			err = fmt.Errorf("%w: count of replicas in dict constrains is %d, but got %d", ErrInvalidConstraintsReplicas, totalCnt, replicas)
-		}
-	}
-	return
+	return newRulesWithDictConstraints(role, cnstr)
 }
 
-// NewRulesWithDictConstraints constructs []*Rule from a yaml-compatible representation of
+// newRulesWithDictConstraints constructs []*Rule from a yaml-compatible representation of
 // 'dict' constraints.
-func NewRulesWithDictConstraints(role PeerRoleType, cnstr string) ([]*Rule, error) {
+func newRulesWithDictConstraints(role PeerRoleType, cnstr string) ([]*Rule, error) {
 	rules := []*Rule{}
 	cnstbytes := []byte(cnstr)
 	constraints2 := map[string]int{}

pkg/ddl/placement/rule_test.go

@@ -172,7 +172,7 @@ func TestNewRuleAndNewRules(t *testing.T) {
 
 	for _, tt := range tests {
 		comment := fmt.Sprintf("[%s]", tt.name)
-		output, err := NewRules(Voter, tt.replicas, tt.input)
+		output, err := newRules(Voter, tt.replicas, tt.input)
 		if tt.err == nil {
 			require.NoError(t, err, comment)
 			matchRules(tt.output, output, comment, t)