diff --git a/ddl/placement/bundle.go b/ddl/placement/bundle.go index 28b3a0999d..666fa7ece6 100644 --- a/ddl/placement/bundle.go +++ b/ddl/placement/bundle.go @@ -70,13 +70,20 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle, followerCount := options.Followers learnerCount := options.Learners + rules := []*Rule{} commonConstraints, err := NewConstraintsFromYaml([]byte(constraints)) if err != nil { - return nil, fmt.Errorf("%w: 'Constraints' should be [constraint1, ...] or any yaml compatible array representation", err) + // If it's not in array format, attempt to parse it as a dictionary for more detailed definitions. + // 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) + if err != nil { + return nil, err + } + rules = append(rules, normalReplicasRules...) } - - rules := []*Rule{} - + needCreateDefault := len(rules) == 0 leaderConstraints, err := NewConstraintsFromYaml([]byte(leaderConst)) if err != nil { return nil, fmt.Errorf("%w: 'LeaderConstraints' should be [constraint1, ...] or any yaml compatible array representation", err) @@ -86,44 +93,59 @@ func NewBundleFromConstraintsOptions(options *model.PlacementSettings) (*Bundle, return nil, fmt.Errorf("%w: LeaderConstraints conflicts with Constraints", err) } } - rules = append(rules, NewRule(Leader, 1, leaderConstraints)) - - followerRules, err := NewRules(Voter, followerCount, followerConstraints) - if err != nil { - return nil, fmt.Errorf("%w: invalid FollowerConstraints", err) + leaderReplicas, followerReplicas := uint64(1), uint64(2) + if followerCount > 0 { + followerReplicas = followerCount } - for _, rule := range followerRules { - // give a default of 2 followers - if rule.Count == 0 { - rule.Count = 2 + if !needCreateDefault { + if len(leaderConstraints) == 0 { + leaderReplicas = 0 } - for _, cnst := range commonConstraints { - if err := rule.Constraints.Add(cnst); err != nil { - return nil, fmt.Errorf("%w: FollowerConstraints conflicts with Constraints", err) + if len(followerConstraints) == 0 { + if followerCount > 0 { + return nil, fmt.Errorf("%w: specify follower count without specify follower constraints when specify other constraints", ErrInvalidPlacementOptions) + } + followerReplicas = 0 + } + } + + // create leader rule. + // if no constraints, we need create default leader rule. + if leaderReplicas > 0 { + leaderRule := NewRule(Leader, leaderReplicas, leaderConstraints) + rules = append(rules, leaderRule) + } + + // create follower rules. + // if no constraints, we need create default follower rules. + if followerReplicas > 0 { + followerRules, err := NewRules(Voter, followerReplicas, followerConstraints) + if err != nil { + return nil, fmt.Errorf("%w: invalid FollowerConstraints", err) + } + for _, followerRule := range followerRules { + for _, cnst := range commonConstraints { + if err := followerRule.Constraints.Add(cnst); err != nil { + return nil, fmt.Errorf("%w: FollowerConstraints conflicts with Constraints", err) + } } } + rules = append(rules, followerRules...) } - rules = append(rules, followerRules...) + // create learner rules. learnerRules, err := NewRules(Learner, learnerCount, learnerConstraints) if err != nil { return nil, fmt.Errorf("%w: invalid LearnerConstraints", err) } for _, rule := range learnerRules { - if rule.Count == 0 { - if len(rule.Constraints) > 0 { - return nil, fmt.Errorf("%w: specify learner constraints without specify how many learners to be placed", ErrInvalidPlacementOptions) - } - } for _, cnst := range commonConstraints { if err := rule.Constraints.Add(cnst); err != nil { return nil, fmt.Errorf("%w: LearnerConstraints conflicts with Constraints", err) } } - if rule.Count > 0 { - rules = append(rules, rule) - } } + rules = append(rules, learnerRules...) labels, err := newLocationLabelsFromSurvivalPreferences(options.SurvivalPreferences) if err != nil { return nil, err diff --git a/ddl/placement/bundle_test.go b/ddl/placement/bundle_test.go index 918c2cf18d..01f8600830 100644 --- a/ddl/placement/bundle_test.go +++ b/ddl/placement/bundle_test.go @@ -396,7 +396,7 @@ func TestNewBundleFromOptions(t *testing.T) { input: &model.PlacementSettings{ LearnerConstraints: "[+region=us]", }, - err: ErrInvalidPlacementOptions, + err: ErrInvalidConstraintsReplicas, }) tests = append(tests, TestCase{ @@ -588,6 +588,40 @@ func TestNewBundleFromOptions(t *testing.T) { }, }) + tests = append(tests, TestCase{ + name: "direct syntax: only leader constraints", + input: &model.PlacementSettings{ + LeaderConstraints: "[+region=as]", + }, + output: []*Rule{ + NewRule(Leader, 1, NewConstraintsDirect(NewConstraintDirect("region", In, "as"))), + NewRule(Voter, 2, NewConstraintsDirect()), + }, + }) + + tests = append(tests, TestCase{ + name: "direct syntax: only leader constraints", + input: &model.PlacementSettings{ + LeaderConstraints: "[+region=as]", + Followers: 4, + }, + output: []*Rule{ + NewRule(Leader, 1, NewConstraintsDirect(NewConstraintDirect("region", In, "as"))), + NewRule(Voter, 4, NewConstraintsDirect()), + }, + }) + tests = append(tests, TestCase{ + name: "direct syntax: leader and follower constraints", + input: &model.PlacementSettings{ + LeaderConstraints: "[+region=as]", + FollowerConstraints: `{"+region=us": 2}`, + }, + output: []*Rule{ + NewRule(Leader, 1, NewConstraintsDirect(NewConstraintDirect("region", In, "as"))), + NewRule(Voter, 2, NewConstraintsDirect(NewConstraintDirect("region", In, "us"))), + }, + }) + tests = append(tests, TestCase{ name: "direct syntax: lack count 1", input: &model.PlacementSettings{ @@ -606,7 +640,7 @@ func TestNewBundleFromOptions(t *testing.T) { LeaderConstraints: "[+region=as]", LearnerConstraints: "[-region=us]", }, - err: ErrInvalidPlacementOptions, + err: ErrInvalidConstraintsReplicas, }) tests = append(tests, TestCase{ @@ -711,7 +745,41 @@ func TestNewBundleFromOptions(t *testing.T) { LearnerConstraints: `{"+region=us": 2}`, Learners: 4, }, - err: ErrInvalidConstraintsRelicas, + err: ErrInvalidConstraintsReplicas, + }) + + tests = append(tests, TestCase{ + name: "direct syntax: dict constraints", + input: &model.PlacementSettings{ + Constraints: `{"+region=us": 3}`, + }, + output: []*Rule{ + NewRule(Voter, 3, NewConstraintsDirect(NewConstraintDirect("region", In, "us"))), + }, + }) + + tests = append(tests, TestCase{ + name: "direct syntax: dict constraints, 2:2:1", + input: &model.PlacementSettings{ + Constraints: `{ "+region=us-east-1":2, "+region=us-east-2": 2, "+region=us-west-1": 1}`, + }, + output: []*Rule{ + NewRule(Voter, 2, NewConstraintsDirect(NewConstraintDirect("region", In, "us-east-1"))), + NewRule(Voter, 2, NewConstraintsDirect(NewConstraintDirect("region", In, "us-east-2"))), + NewRule(Voter, 1, NewConstraintsDirect(NewConstraintDirect("region", In, "us-west-1"))), + }, + }) + + tests = append(tests, TestCase{ + name: "direct syntax: dict constraints", + input: &model.PlacementSettings{ + Constraints: `{"+region=us-east": 3}`, + LearnerConstraints: `{"+region=us-west": 1}`, + }, + output: []*Rule{ + NewRule(Voter, 3, NewConstraintsDirect(NewConstraintDirect("region", In, "us-east"))), + NewRule(Learner, 1, NewConstraintsDirect(NewConstraintDirect("region", In, "us-west"))), + }, }) for _, test := range tests { @@ -856,7 +924,7 @@ func TestTidy(t *testing.T) { rules1, err := NewRules(Voter, 4, `["-zone=sh", "+zone=bj"]`) require.NoError(t, err) require.Len(t, rules1, 1) - rules2, err := NewRules(Voter, 4, `["-zone=sh", "+zone=bj"]`) + 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...) diff --git a/ddl/placement/errors.go b/ddl/placement/errors.go index a93d0e2bd4..99842ba9a3 100644 --- a/ddl/placement/errors.go +++ b/ddl/placement/errors.go @@ -31,8 +31,8 @@ var ( ErrInvalidConstraintsFormat = errors.New("invalid label constraints format") // ErrInvalidSurvivalPreferenceFormat is from rule.go. ErrInvalidSurvivalPreferenceFormat = errors.New("survival preference format should be in format [xxx=yyy, ...]") - // ErrInvalidConstraintsRelicas is from rule.go. - ErrInvalidConstraintsRelicas = errors.New("label constraints with invalid REPLICAS") + // ErrInvalidConstraintsReplicas is from rule.go. + ErrInvalidConstraintsReplicas = errors.New("label constraints with invalid REPLICAS") // ErrInvalidBundleID is from bundle.go. ErrInvalidBundleID = errors.New("invalid bundle ID") // ErrInvalidBundleIDFormat is from bundle.go. diff --git a/ddl/placement/rule.go b/ddl/placement/rule.go index 4a1b759b83..711fc57c6d 100644 --- a/ddl/placement/rule.go +++ b/ddl/placement/rule.go @@ -178,24 +178,49 @@ func getYamlMapFormatError(str string) error { // 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) ([]*Rule, error) { - rules := []*Rule{} - +func NewRules(role PeerRoleType, replicas uint64, cnstr string) (rules []*Rule, err error) { cnstbytes := []byte(cnstr) - constraints1, err1 := NewConstraintsFromYaml(cnstbytes) if err1 == nil { + if replicas == 0 { + if len(cnstr) > 0 { + return nil, fmt.Errorf("%w: count of replicas should be positive, but got %d, constraint %s", ErrInvalidConstraintsReplicas, replicas, cnstr) + } + return nil, nil + } rules = append(rules, NewRule(role, replicas, constraints1)) - return rules, nil + err = err1 + return + } + // check if is dict constraints + constraints2 := map[string]int{} + if err2 := yaml.UnmarshalStrict(cnstbytes, &constraints2); err2 != nil { + err = fmt.Errorf("%w: should be [constraint1, ...] (error %s), {constraint1: cnt1, ...} (error %s), or any yaml compatible representation", ErrInvalidConstraintsFormat, err1, err2) + 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 +} + +// NewRulesWithDictConstraints constructs []*Rule from a yaml-compatible representation of +// 'dict' constraints. +func NewRulesWithDictConstraints(role PeerRoleType, cnstr string) ([]*Rule, error) { + rules := []*Rule{} + cnstbytes := []byte(cnstr) constraints2 := map[string]int{} err2 := yaml.UnmarshalStrict(cnstbytes, &constraints2) if err2 == nil { - if replicas != 0 { - return rules, fmt.Errorf("%w: should not specify replicas=%d when using dict syntax", ErrInvalidConstraintsRelicas, replicas) - } - for labels, cnt := range constraints2 { if cnt <= 0 { if err := getYamlMapFormatError(string(cnstbytes)); err != nil { @@ -214,13 +239,15 @@ func NewRules(role PeerRoleType, replicas uint64, cnstr string) ([]*Rule, error) if err != nil { return rules, err } - + if cnt == 0 { + return nil, fmt.Errorf("%w: count of replicas should be positive, but got %d", ErrInvalidConstraintsReplicas, cnt) + } rules = append(rules, NewRule(overrideRole, uint64(cnt), labelConstraints)) } return rules, nil } - return nil, fmt.Errorf("%w: should be [constraint1, ...] (error %s), {constraint1: cnt1, ...} (error %s), or any yaml compatible representation", ErrInvalidConstraintsFormat, err1, err2) + return nil, fmt.Errorf("%w: should be [constraint1, ...] or {constraint1: cnt1, ...}, error %s, or any yaml compatible representation", ErrInvalidConstraintsFormat, err2) } // Clone is used to duplicate a RuleOp for safe modification. diff --git a/ddl/placement/rule_test.go b/ddl/placement/rule_test.go index c2c0e5313f..ed1e032b31 100644 --- a/ddl/placement/rule_test.go +++ b/ddl/placement/rule_test.go @@ -70,9 +70,7 @@ func TestNewRuleAndNewRules(t *testing.T) { name: "zero replicas", input: "", replicas: 0, - output: []*Rule{ - NewRule(Voter, 0, NewConstraintsDirect()), - }, + output: nil, }) tests = append(tests, TestCase{ @@ -102,10 +100,17 @@ func TestNewRuleAndNewRules(t *testing.T) { }) tests = append(tests, TestCase{ - name: "normal dict constraints, with count", - input: "{'+zone=sh,-zone=bj':2, '+zone=sh': 1}", - replicas: 4, - err: ErrInvalidConstraintsRelicas, + name: "normal dict constraints, with count", + input: "{'+zone=sh,-zone=bj':2, '+zone=sh': 1}", + output: []*Rule{ + NewRule(Voter, 2, NewConstraintsDirect( + NewConstraintDirect("zone", In, "sh"), + NewConstraintDirect("zone", NotIn, "bj"), + )), + NewRule(Voter, 1, NewConstraintsDirect( + NewConstraintDirect("zone", In, "sh"), + )), + }, }) tests = append(tests, TestCase{ diff --git a/ddl/placement_policy_test.go b/ddl/placement_policy_test.go index 91c8ffb0df..d0761812fc 100644 --- a/ddl/placement_policy_test.go +++ b/ddl/placement_policy_test.go @@ -404,7 +404,7 @@ func TestPlacementValidation(t *testing.T) { settings: "LEARNERS=1 " + "LEARNER_CONSTRAINTS=\"[+zone=cn-west-1]\" " + "CONSTRAINTS=\"{'+disk=ssd':2}\"", - errmsg: "invalid label constraints format: 'Constraints' should be [constraint1, ...] or any yaml compatible array representation", + success: true, }, { name: "constraints may be incompatible with itself", @@ -428,7 +428,7 @@ func TestPlacementValidation(t *testing.T) { tk.MustExec("drop placement policy if exists x") } else { err := tk.ExecToErr(sql) - require.NotNil(t, err) + require.NotNil(t, err, sql) require.EqualErrorf(t, err, ca.errmsg, ca.name) } }