placement: add rule definitions (#1834)

Signed-off-by: disksing <i@disksing.com>

pkg/slice/slice.go

@@ -0,0 +1,37 @@
+// Copyright 2019 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 slice
+
+import "reflect"
+
+// AnyOf returns true if any element in the slice matches the predict func.
+func AnyOf(s interface{}, p func(int) bool) bool {
+	l := reflect.ValueOf(s).Len()
+	for i := 0; i < l; i++ {
+		if p(i) {
+			return true
+		}
+	}
+	return false
+}
+
+// NoneOf returns true if no element in the slice matches the predict func.
+func NoneOf(s interface{}, p func(int) bool) bool {
+	return !AnyOf(s, p)
+}
+
+// AllOf returns true if all elements in the slice match the predict func.
+func AllOf(s interface{}, p func(int) bool) bool {
+	return NoneOf(s, func(i int) bool { return !p(i) })
+}

pkg/slice/slice_test.go

@@ -0,0 +1,51 @@
+// Copyright 2019 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 slice_test
+
+import (
+	"testing"
+
+	. "github.com/pingcap/check"
+	"github.com/pingcap/pd/pkg/slice"
+)
+
+func Test(t *testing.T) {
+	TestingT(t)
+}
+
+var _ = Suite(&testSliceSuite{})
+
+type testSliceSuite struct {
+}
+
+func (s *testSliceSuite) Test(c *C) {
+	tests := []struct {
+		a      []int
+		anyOf  bool
+		noneOf bool
+		allOf  bool
+	}{
+		{[]int{}, false, true, true},
+		{[]int{1, 2, 3}, true, false, false},
+		{[]int{1, 3}, false, true, false},
+		{[]int{2, 2, 4}, true, false, true},
+	}
+
+	for _, t := range tests {
+		even := func(i int) bool { return t.a[i]%2 == 0 }
+		c.Assert(slice.AnyOf(t.a, even), Equals, t.anyOf)
+		c.Assert(slice.NoneOf(t.a, even), Equals, t.noneOf)
+		c.Assert(slice.AllOf(t.a, even), Equals, t.allOf)
+	}
+}

server/schedule/placement/label_constraint.go

@@ -0,0 +1,86 @@
+// Copyright 2019 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 (
+	"github.com/pingcap/pd/pkg/slice"
+	"github.com/pingcap/pd/server/core"
+)
+
+// LabelConstraintOp defines how a LabelConstraint matches a store. It can be one of
+// 'in', 'notIn', 'exists', or 'notExists'.
+type LabelConstraintOp string
+
+const (
+	// In restricts the store label value should in the value list.
+	// If label does not exist, `in` is always false.
+	In LabelConstraintOp = "in"
+	// NotIn restricts the store label value should not in the value list.
+	// If label does not exist, `notIn` is always true.
+	NotIn LabelConstraintOp = "notIn"
+	// Exists restricts the store should have the label.
+	Exists LabelConstraintOp = "exists"
+	// NotExists restricts the store should not have the label.
+	NotExists LabelConstraintOp = "notExists"
+)
+
+func validateOp(op LabelConstraintOp) bool {
+	return op == In || op == NotIn || op == Exists || op == NotExists
+}
+
+// LabelConstraint is used to filter store when trying to place peer of a region.
+type LabelConstraint struct {
+	Key    string            `json:"key,omitempty"`
+	Op     LabelConstraintOp `json:"op,omitempty"`
+	Values []string          `json:"values,omitempty"`
+}
+
+// MatchStore checks if a store matches the constraint.
+func (c *LabelConstraint) MatchStore(store *core.StoreInfo) bool {
+	switch c.Op {
+	case In:
+		label := store.GetLabelValue(c.Key)
+		return label != "" && slice.AnyOf(c.Values, func(i int) bool { return c.Values[i] == label })
+	case NotIn:
+		label := store.GetLabelValue(c.Key)
+		return label == "" || slice.NoneOf(c.Values, func(i int) bool { return c.Values[i] == label })
+	case Exists:
+		return store.GetLabelValue(c.Key) != ""
+	case NotExists:
+		return store.GetLabelValue(c.Key) == ""
+	}
+	return false
+}
+
+// If a store has exclusiveLabels, it can only be selected when the label is
+// exciplitly specified in constraints.
+// TODO: move it to config.
+var exclusiveLabels = []string{"engine"}
+
+// MatchLabelConstraints checks if a store matches label constraints list.
+func MatchLabelConstraints(store *core.StoreInfo, constraints []LabelConstraint) bool {
+	if store == nil {
+		return false
+	}
+
+	if slice.AnyOf(exclusiveLabels, func(i int) bool { // if there is any exclusive label that
+		label := exclusiveLabels[i]
+		return store.GetLabelValue(label) != "" && // ... the store has the exclusive label
+			slice.NoneOf(constraints, func(i int) bool { return constraints[i].Key == label }) // ... but the exclusive label is not in constraints
+	}) {
+		return false // ... then the store should be ignored
+	}
+
+	return slice.AllOf(constraints, func(i int) bool { return constraints[i].MatchStore(store) })
+}

server/schedule/placement/label_constraint_test.go

@@ -0,0 +1,101 @@
+// Copyright 2019 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 (
+	"testing"
+
+	. "github.com/pingcap/check"
+	"github.com/pingcap/pd/server/core"
+)
+
+func TestPlacement(t *testing.T) {
+	TestingT(t)
+}
+
+var _ = Suite(&testLabelConstraintsSuite{})
+
+type testLabelConstraintsSuite struct{}
+
+func (s *testLabelConstraintsSuite) TestLabelConstraint(c *C) {
+	stores := []map[string]string{
+		{"zone": "zone1", "rack": "rack1"},                    // 1
+		{"zone": "zone1", "rack": "rack2"},                    // 2
+		{"zone": "zone2"},                                     // 3
+		{"zone": "zone1", "engine": "rocksdb", "disk": "hdd"}, // 4
+		{"rack": "rack1", "disk": "ssd"},                      // 5
+		{"zone": "zone3"},                                     // 6
+	}
+	constraints := []LabelConstraint{
+		{Key: "zone", Op: "in", Values: []string{"zone1"}},
+		{Key: "zone", Op: "in", Values: []string{"zone1", "zone2"}},
+		{Key: "zone", Op: "notIn", Values: []string{"zone1"}},
+		{Key: "zone", Op: "notIn", Values: []string{"zone1", "zone2"}},
+		{Key: "zone", Op: "exists"},
+		{Key: "disk", Op: "notExists"},
+	}
+	expect := [][]int{
+		{1, 2, 4},
+		{1, 2, 3, 4},
+		{3, 5, 6},
+		{5, 6},
+		{1, 2, 3, 4, 6},
+		{1, 2, 3, 6},
+	}
+	for i, constraint := range constraints {
+		var matched []int
+		for j, store := range stores {
+			if constraint.MatchStore(core.NewStoreInfoWithLabel(uint64(j), 0, store)) {
+				matched = append(matched, j+1)
+			}
+		}
+		c.Assert(matched, DeepEquals, expect[i])
+	}
+}
+
+func (s *testLabelConstraintsSuite) TestLabelConstraints(c *C) {
+	stores := []map[string]string{
+		{},                                       // 1
+		{"k1": "v1"},                             // 2
+		{"k1": "v2"},                             // 3
+		{"k2": "v1"},                             // 4
+		{"k2": "v2"},                             // 5
+		{"engine": "e1"},                         // 6
+		{"engine": "e2"},                         // 7
+		{"k1": "v1", "k2": "v1"},                 // 8
+		{"k2": "v2", "engine": "e1"},             // 9
+		{"k1": "v1", "k2": "v2", "engine": "e2"}, // 10
+	}
+	constraints := [][]LabelConstraint{
+		{},
+		{{Key: "engine", Op: "in", Values: []string{"e1", "e2"}}},
+		{{Key: "k1", Op: "notExists"}, {Key: "k2", Op: "exists"}},
+		{{Key: "engine", Op: "exists"}, {Key: "k1", Op: "in", Values: []string{"v1", "v2"}}, {Key: "k2", Op: "notIn", Values: []string{"v3"}}},
+	}
+	expect := [][]int{
+		{1, 2, 3, 4, 5, 8},
+		{6, 7, 9, 10},
+		{4, 5},
+		{10},
+	}
+	for i, cs := range constraints {
+		var matched []int
+		for j, store := range stores {
+			if MatchLabelConstraints(core.NewStoreInfoWithLabel(uint64(j), 0, store), cs) {
+				matched = append(matched, j+1)
+			}
+		}
+		c.Assert(matched, DeepEquals, expect[i])
+	}
+}

server/schedule/placement/rule.go

@@ -0,0 +1,98 @@
+// Copyright 2019 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 "sort"
+
+// 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"
+)
+
+func validateRole(s PeerRoleType) bool {
+	return s == Voter || s == Leader || s == Follower || s == Learner
+}
+
+// Rule is the placement rule that can be checked against a region. When
+// applying rules (apply means schedule regions to match selected rules), the
+// apply order is defined by the tuple [GroupID, Index, ID].
+type Rule struct {
+	GroupID          string            `json:"group_id"`                    // mark the source that add the rule
+	ID               string            `json:"id"`                          // unique ID within a group
+	Index            int               `json:"index,omitempty"`             // rule apply order in a group, rule with less ID is applied first when indexes are equal
+	Override         bool              `json:"override,omitempty"`          // when it is true, all rules with less indexes are disabled
+	StartKey         []byte            `json:"-"`                           // range start key
+	StartKeyHex      string            `json:"start_key"`                   // hex format start key, for marshal/unmarshal
+	EndKey           []byte            `json:"-"`                           // range end key
+	EndKeyHex        string            `json:"end_key"`                     // hex format end key, for marshal/unmarshal
+	Role             PeerRoleType      `json:"role"`                        // expected role of the peers
+	Count            int               `json:"count"`                       // expected count of the peers
+	LabelConstraints []LabelConstraint `json:"label_constraints,omitempty"` // used to select stores to place peers
+	LocationLabels   []string          `json:"location_labels,omitempty"`   // used to make peers isolated physically
+}
+
+// Key returns (groupID, ID) as the global unique key of a rule.
+func (r Rule) Key() [2]string {
+	return [2]string{r.GroupID, r.ID}
+}
+
+// Rules are ordered by (GroupID, Index, ID).
+func compareRule(a, b *Rule) int {
+	switch {
+	case a.GroupID < b.GroupID:
+		return -1
+	case a.GroupID > b.GroupID:
+		return 1
+	case a.Index < b.Index:
+		return -1
+	case a.Index > b.Index:
+		return 1
+	case a.ID < b.ID:
+		return -1
+	case a.ID > b.ID:
+		return 1
+	default:
+		return 0
+	}
+}
+
+func sortRules(rules []*Rule) {
+	sort.Slice(rules, func(i, j int) bool { return compareRule(rules[i], rules[j]) < 0 })
+}
+
+// Sort Rules, trim concealed rules.
+func prepareRulesForApply(rules []*Rule) []*Rule {
+	sortRules(rules)
+	res := rules[:0]
+	var i, j int
+	for i = 1; i < len(rules); i++ {
+		if rules[j].GroupID != rules[i].GroupID {
+			res = append(res, rules[j:i]...)
+			j = i
+		}
+		if rules[i].Override {
+			j = i
+		}
+	}
+	return append(res, rules[j:]...)
+}