ddl: refactor rule [4/6]

ddl/ddl_api.go

@@ -5900,8 +5900,8 @@ func buildPlacementSpecReplicasAndConstraint(replicas uint64, cnstr string) ([]*
 		}
 
 		rules = append(rules, &placement.Rule{
-			Count:            int(replicas),
-			LabelConstraints: labelConstraints,
+			Count:       int(replicas),
+			Constraints: labelConstraints,
 		})
 
 		return rules, nil
@@ -5930,8 +5930,8 @@ func buildPlacementSpecReplicasAndConstraint(replicas uint64, cnstr string) ([]*
 			}
 
 			rules = append(rules, &placement.Rule{
-				Count:            cnt,
-				LabelConstraints: labelConstraints,
+				Count:       cnt,
+				Constraints: labelConstraints,
 			})
 		}
 
@@ -6056,14 +6056,14 @@ func (d *ddl) AlterTableAlterPartition(ctx sessionctx.Context, ident ast.Ident,
 	newRules := bundle.Rules[:0]
 	for i, rule := range bundle.Rules {
 		// merge all empty constraints
-		if len(rule.LabelConstraints) == 0 {
+		if len(rule.Constraints) == 0 {
 			extraCnt[rule.Role] += rule.Count
 			continue
 		}
 		// refer to tidb#22065.
 		// add -engine=tiflash to every rule to avoid schedules to tiflash instances.
 		// placement rules in SQL is not compatible with `set tiflash replica` yet
-		if err := rule.LabelConstraints.Add(placement.Constraint{
+		if err := rule.Constraints.Add(placement.Constraint{
 			Op:     placement.NotIn,
 			Key:    placement.EngineLabelKey,
 			Values: []string{placement.EngineLabelTiFlash},
@@ -6088,7 +6088,7 @@ func (d *ddl) AlterTableAlterPartition(ctx sessionctx.Context, ident ast.Ident,
 			Count:       cnt,
 			StartKeyHex: startKey,
 			EndKeyHex:   endKey,
-			LabelConstraints: []placement.Constraint{{
+			Constraints: []placement.Constraint{{
 				Op:     placement.NotIn,
 				Key:    placement.EngineLabelKey,
 				Values: []string{placement.EngineLabelTiFlash},

ddl/placement/errors.go

@@ -24,4 +24,10 @@ var (
 	ErrUnsupportedConstraint = errors.New("unsupported label constraint")
 	// ErrConflictingConstraints is from constraints.go.
 	ErrConflictingConstraints = errors.New("conflicting label constraints")
+	// ErrInvalidConstraintsMapcnt is from rule.go.
+	ErrInvalidConstraintsMapcnt = errors.New("label constraints in map syntax have invalid replicas")
+	// ErrInvalidConstraintsFormat is from rule.go.
+	ErrInvalidConstraintsFormat = errors.New("invalid label constraints format")
+	// ErrInvalidConstraintsRelicas is from rule.go.
+	ErrInvalidConstraintsRelicas = errors.New("label constraints with invalid REPLICAS")
 )

ddl/placement/rule.go

@@ -0,0 +1,132 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package placement
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/go-yaml/yaml"
+)
+
+// PeerRoleType is the expected peer type of the placement rule.
+type PeerRoleType string
+
+const (
+	// Voter can either match a leader peer or follower peer.
+	Voter PeerRoleType = "voter"
+	// Leader matches a leader.
+	Leader PeerRoleType = "leader"
+	// Follower matches a follower.
+	Follower PeerRoleType = "follower"
+	// Learner matches a learner.
+	Learner PeerRoleType = "learner"
+)
+
+// Rule is the core placement rule struct. Check https://github.com/tikv/pd/blob/master/server/schedule/placement/rule.go.
+type Rule struct {
+	GroupID        string       `json:"group_id"`
+	ID             string       `json:"id"`
+	Index          int          `json:"index,omitempty"`
+	Override       bool         `json:"override,omitempty"`
+	StartKeyHex    string       `json:"start_key"`
+	EndKeyHex      string       `json:"end_key"`
+	Role           PeerRoleType `json:"role"`
+	Count          int          `json:"count"`
+	Constraints    Constraints  `json:"label_constraints,omitempty"`
+	LocationLabels []string     `json:"location_labels,omitempty"`
+	IsolationLevel string       `json:"isolation_level,omitempty"`
+}
+
+// NewRules constructs []*Rule from a yaml-compatible representation of
+// array or map of constraints. It converts 'CONSTRAINTS' field in RFC
+// https://github.com/pingcap/tidb/blob/master/docs/design/2020-06-24-placement-rules-in-sql.md to structs.
+func NewRules(replicas uint64, cnstr string) ([]*Rule, error) {
+	rules := []*Rule{}
+
+	cnstbytes := []byte(cnstr)
+
+	constraints1 := []string{}
+	err1 := yaml.UnmarshalStrict(cnstbytes, &constraints1)
+	if err1 == nil {
+		// can not emit REPLICAS with an array or empty label
+		if replicas == 0 {
+			return rules, fmt.Errorf("%w: should be positive", ErrInvalidConstraintsRelicas)
+		}
+
+		labelConstraints, err := NewConstraints(constraints1)
+		if err != nil {
+			return rules, err
+		}
+
+		rules = append(rules, &Rule{
+			Count:       int(replicas),
+			Constraints: labelConstraints,
+		})
+
+		return rules, nil
+	}
+
+	constraints2 := map[string]int{}
+	err2 := yaml.UnmarshalStrict(cnstbytes, &constraints2)
+	if err2 == nil {
+		ruleCnt := 0
+		for labels, cnt := range constraints2 {
+			if cnt <= 0 {
+				return rules, fmt.Errorf("%w: count of labels '%s' should be positive, but got %d", ErrInvalidConstraintsMapcnt, labels, cnt)
+			}
+			ruleCnt += cnt
+		}
+
+		if replicas == 0 {
+			replicas = uint64(ruleCnt)
+		}
+
+		if int(replicas) < ruleCnt {
+			return rules, fmt.Errorf("%w: should be larger or equal to the number of total replicas, but REPLICAS=%d < total=%d", ErrInvalidConstraintsRelicas, replicas, ruleCnt)
+		}
+
+		for labels, cnt := range constraints2 {
+			labelConstraints, err := NewConstraints(strings.Split(labels, ","))
+			if err != nil {
+				return rules, err
+			}
+
+			rules = append(rules, &Rule{
+				Count:       cnt,
+				Constraints: labelConstraints,
+			})
+		}
+
+		remain := int(replicas) - ruleCnt
+		if remain > 0 {
+			rules = append(rules, &Rule{
+				Count: remain,
+			})
+		}
+
+		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)
+}
+
+// Clone is used to duplicate a RuleOp for safe modification.
+// Note that it is a shallow copy: LocationLabels and Constraints
+// is not cloned.
+func (r *Rule) Clone() *Rule {
+	n := &Rule{}
+	*n = *r
+	return n
+}

ddl/placement/rule_test.go

@@ -0,0 +1,206 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package placement
+
+import (
+	"encoding/json"
+	"errors"
+
+	. "github.com/pingcap/check"
+)
+
+var _ = Suite(&testRuleSuite{})
+
+type testRuleSuite struct{}
+
+func (t *testRuleSuite) TestClone(c *C) {
+	rule := &Rule{ID: "434"}
+	newRule := rule.Clone()
+	newRule.ID = "121"
+
+	c.Assert(rule, DeepEquals, &Rule{ID: "434"})
+	c.Assert(newRule, DeepEquals, &Rule{ID: "121"})
+}
+
+func matchRule(r1 *Rule, t2 []*Rule) bool {
+	for _, r2 := range t2 {
+		if ok, _ := DeepEquals.Check([]interface{}{r1, r2}, nil); ok {
+			return true
+		}
+	}
+	return false
+}
+
+func matchRules(t1, t2 []*Rule, prefix string, c *C) {
+	expected, err := json.Marshal(t1)
+	c.Assert(err, IsNil)
+	got, err := json.Marshal(t2)
+	c.Assert(err, IsNil)
+	comment := Commentf("%s, expected %s\nbut got %s", prefix, expected, got)
+	c.Assert(len(t1), Equals, len(t2), comment)
+	for _, r1 := range t1 {
+		c.Assert(matchRule(r1, t2), IsTrue, comment)
+	}
+}
+
+func (t *testRuleSuite) TestNewRules(c *C) {
+	type TestCase struct {
+		name     string
+		input    string
+		replicas uint64
+		output   []*Rule
+		err      error
+	}
+	tests := []TestCase{}
+
+	tests = append(tests, TestCase{
+		name:     "empty constraints",
+		input:    "",
+		replicas: 3,
+		output: []*Rule{
+			{
+				Count:       3,
+				Constraints: Constraints{},
+			},
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name:     "zero replicas",
+		input:    "",
+		replicas: 0,
+		err:      ErrInvalidConstraintsRelicas,
+	})
+
+	labels, err := NewConstraints([]string{"+zone=sh", "+zone=sh"})
+	c.Assert(err, IsNil)
+	tests = append(tests, TestCase{
+		name:     "normal array constraints",
+		input:    `["+zone=sh", "+zone=sh"]`,
+		replicas: 3,
+		output: []*Rule{
+			{
+				Count:       3,
+				Constraints: labels,
+			},
+		},
+	})
+
+	labels1, err := NewConstraints([]string{"+zone=sh", "-zone=bj"})
+	c.Assert(err, IsNil)
+	labels2, err := NewConstraints([]string{"+zone=sh"})
+	c.Assert(err, IsNil)
+	tests = append(tests, TestCase{
+		name:     "normal object constraints",
+		input:    `{"+zone=sh,-zone=bj":2, "+zone=sh": 1}`,
+		replicas: 3,
+		output: []*Rule{
+			{
+				Count:       2,
+				Constraints: labels1,
+			},
+			{
+				Count:       1,
+				Constraints: labels2,
+			},
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name:     "normal object constraints, with extra count",
+		input:    "{'+zone=sh,-zone=bj':2, '+zone=sh': 1}",
+		replicas: 4,
+		output: []*Rule{
+			{
+				Count:       2,
+				Constraints: labels1,
+			},
+			{
+				Count:       1,
+				Constraints: labels2,
+			},
+			{
+				Count: 1,
+			},
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name:  "normal object constraints, without count",
+		input: "{'+zone=sh,-zone=bj':2, '+zone=sh': 1}",
+		output: []*Rule{
+			{
+				Count:       2,
+				Constraints: labels1,
+			},
+			{
+				Count:       1,
+				Constraints: labels2,
+			},
+		},
+	})
+
+	tests = append(tests, TestCase{
+		name:     "zero count in object constraints",
+		input:    `{"+zone=sh,-zone=bj":0, "+zone=sh": 1}`,
+		replicas: 3,
+		err:      ErrInvalidConstraintsMapcnt,
+	})
+
+	tests = append(tests, TestCase{
+		name:     "overlarge total count in object constraints",
+		input:    `{"+ne=sh,-zone=bj":1, "+zone=sh": 4}`,
+		replicas: 3,
+		err:      ErrInvalidConstraintsRelicas,
+	})
+
+	tests = append(tests, TestCase{
+		name:     "invalid array",
+		input:    `["+ne=sh", "+zone=sh"`,
+		replicas: 3,
+		err:      ErrInvalidConstraintsFormat,
+	})
+
+	tests = append(tests, TestCase{
+		name:     "invalid array constraints",
+		input:    `["ne=sh", "+zone=sh"]`,
+		replicas: 3,
+		err:      ErrInvalidConstraintFormat,
+	})
+
+	tests = append(tests, TestCase{
+		name:     "invalid map",
+		input:    `{+ne=sh,-zone=bj:1, "+zone=sh": 4`,
+		replicas: 5,
+		err:      ErrInvalidConstraintsFormat,
+	})
+
+	tests = append(tests, TestCase{
+		name:     "invalid map constraints",
+		input:    `{"nesh,-zone=bj":1, "+zone=sh": 4}`,
+		replicas: 6,
+		err:      ErrInvalidConstraintFormat,
+	})
+
+	for _, t := range tests {
+		comment := Commentf("%s", t.name)
+		output, err := NewRules(t.replicas, t.input)
+		if t.err == nil {
+			c.Assert(err, IsNil, comment)
+			matchRules(t.output, output, comment.CheckCommentString(), c)
+		} else {
+			c.Assert(errors.Is(err, t.err), IsTrue, comment)
+		}
+	}
+}