diff --git a/pkg/ddl/placement/bundle.go b/pkg/ddl/placement/bundle.go index 110a39f397ca3..e9331571c6d2c 100644 --- a/pkg/ddl/placement/bundle.go +++ b/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) } diff --git a/pkg/ddl/placement/bundle_test.go b/pkg/ddl/placement/bundle_test.go index 960a37b9236ec..c1b7c067c774b 100644 --- a/pkg/ddl/placement/bundle_test.go +++ b/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) diff --git a/pkg/ddl/placement/rule.go b/pkg/ddl/placement/rule.go index f21588b664682..c52839c37bb32 100644 --- a/pkg/ddl/placement/rule.go +++ b/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{} diff --git a/pkg/ddl/placement/rule_test.go b/pkg/ddl/placement/rule_test.go index ed1e032b31892..dd6eadf4d29c5 100644 --- a/pkg/ddl/placement/rule_test.go +++ b/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)