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

) (#47890)

close #47741

go.mod

@@ -186,7 +186,7 @@ require (
  github.com/eapache/go-resiliency v1.2.0 // indirect
  github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21 // indirect
  github.com/eapache/queue v1.1.0 // indirect
- github.com/fatih/structtag v1.2.0 // indirect
+ github.com/fatih/structtag v1.2.0
  github.com/felixge/httpsnoop v1.0.2 // indirect
  github.com/form3tech-oss/jwt-go v3.2.5+incompatible // indirect
  github.com/go-asn1-ber/asn1-ber v1.5.4 // indirect

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)