placement: make Constraints can use dict format (#47344)

close #47254

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))
+	leaderReplicas, followerReplicas := uint64(1), uint64(2)
+	if followerCount > 0 {
+		followerReplicas = followerCount
+	}
+	if !needCreateDefault {
+		if len(leaderConstraints) == 0 {
+			leaderReplicas = 0
+		}
+		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
+		}
+	}
 
-	followerRules, err := NewRules(Voter, followerCount, followerConstraints)
-	if err != nil {
-		return nil, fmt.Errorf("%w: invalid FollowerConstraints", err)
+	// create leader rule.
+	// if no constraints, we need create default leader rule.
+	if leaderReplicas > 0 {
+		leaderRule := NewRule(Leader, leaderReplicas, leaderConstraints)
+		rules = append(rules, leaderRule)
 	}
-	for _, rule := range followerRules {
-		// give a default of 2 followers
-		if rule.Count == 0 {
-			rule.Count = 2
+
+	// 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 _, cnst := range commonConstraints {
-			if err := rule.Constraints.Add(cnst); err != nil {
-				return nil, fmt.Errorf("%w: FollowerConstraints conflicts with Constraints", 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

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...)

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.

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.

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{

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)
 		}
 	}