scheduler: add new api to add a new job for the special scheduler

pkg/core/basic_cluster.go

@@ -16,6 +16,7 @@
 
 import (
 	"bytes"
+	"encoding/hex"
 	"encoding/json"
 
 	"github.com/tikv/pd/pkg/core/constant"
@@ -157,15 +158,38 @@
 	EndKey   []byte `json:"end-key"`
 }
 
+var _ json.Marshaler = &KeyRange{}
+var _ json.Unmarshaler = &KeyRange{}
+
 // MarshalJSON marshals to json.
-func (kr KeyRange) MarshalJSON() ([]byte, error) {
+func (kr *KeyRange) MarshalJSON() ([]byte, error) {
 	m := map[string]string{
 		"start-key": HexRegionKeyStr(kr.StartKey),
 		"end-key":   HexRegionKeyStr(kr.EndKey),
 	}
 	return json.Marshal(m)
 }
 
+// UnmarshalJSON unmarshals from json.
+func (kr *KeyRange) UnmarshalJSON(data []byte) error {
+	m := make(map[string]string)
+	if err := json.Unmarshal(data, &m); err != nil {
+		return err
+	}
+
+	startKey, err := hex.DecodeString(m["start-key"])
+	if err != nil {
+		return err
+	}
+	endKey, err := hex.DecodeString(m["end-key"])
+	if err != nil {
+		return err
+	}
+	kr.StartKey = startKey
+	kr.EndKey = endKey
+	return nil
+}
+
 // NewKeyRange create a KeyRange with the given start key and end key.
 func NewKeyRange(startKey, endKey string) KeyRange {
 	return KeyRange{

pkg/core/basic_cluster_test.go

@@ -15,11 +15,35 @@
 package core
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
+func TestCodecKeyRange(t *testing.T) {
+	re := require.New(t)
+
+	testCases := []struct {
+		ks KeyRange
+	}{
+		{
+			NewKeyRange(fmt.Sprintf("%20d", 0), fmt.Sprintf("%20d", 5)),
+		},
+		{
+			NewKeyRange(fmt.Sprintf("%20d", 0), fmt.Sprintf("%20d", 10)),
+		},
+	}
+
+	for _, tc := range testCases {
+		data, err := tc.ks.MarshalJSON()
+		re.NoError(err)
+		var ks KeyRange
+		re.NoError(ks.UnmarshalJSON(data))
+		re.Equal(tc.ks, ks)
+	}
+}
+
 func TestMergeKeyRanges(t *testing.T) {
 	re := require.New(t)
 

pkg/core/region.go

@@ -360,68 +360,78 @@
 	return nil
 }
 
-// Role is the role of the region.
-type Role int
+// Rule is the rule for balance range scheduler
+type Rule int
 
 const (
-	// Leader is the leader of the region.
-	Leader Role = iota
-	// Follower is the follower of the region.
-	Follower
-	// Learner is the learner of the region.
-	Learner
-	// Unknown is the unknown role of the region include witness.
+	// LeaderScatter scatter the leader of the region.
+	LeaderScatter Rule = iota
+	// PeerScatter scatter all the peers of the region.
+	PeerScatter
+	// LearnerScatter is the learner of the region.
+	LearnerScatter
+	// Unknown is the unknown rule of the region include witness.
 	Unknown
 )
 
-// String returns the string value of the role.
-func (r Role) String() string {
-	switch r {
-	case Leader:
-		return "leader"
-	case Follower:
-		return "voter"
-	case Learner:
-		return "learner"
+// String returns the string value of the rule.
+func (r *Rule) String() string {
+	switch *r {
+	case LeaderScatter:
+		return "leader-scatter"
+	case PeerScatter:
+		return "peer-scatter"
+	case LearnerScatter:
+		return "learner-scatter"
 	default:
 		return "unknown"
 	}
 }
 
-// NewRole creates a new role.
-func NewRole(role string) Role {
-	switch role {
-	case "leader":
-		return Leader
-	case "follower":
-		return Follower
-	case "learner":
-		return Learner
+// NewRule creates a new rule.
+func NewRule(rule string) Rule {
+	switch rule {
+	case "leader-scatter":
+		return LeaderScatter
+	case "peer-scatter":
+		return PeerScatter
+	case "learner-scatter":
+		return LearnerScatter
 	default:
 		return Unknown
 	}
 }
 
-// MarshalJSON returns the JSON encoding of Role.
-func (r Role) MarshalJSON() ([]byte, error) {
+// MarshalJSON returns the JSON encoding of rule.
+func (r *Rule) MarshalJSON() ([]byte, error) {
 	return []byte(`"` + r.String() + `"`), nil
 }
 
-// GetPeersByRole returns the peers with specified role.
-func (r *RegionInfo) GetPeersByRole(role Role) []*metapb.Peer {
-	switch role {
-	case Leader:
+// UnmarshalJSON parses the JSON-encoded data and stores the result in rule.
+func (r *Rule) UnmarshalJSON(data []byte) error {
+	s := string(data)
+	switch s {
+	case `"leader-scatter"`:
+		*r = LeaderScatter
+	case `"peer-scatter"`:
+		*r = PeerScatter
+	case `"learner-scatter"`:
+		*r = LearnerScatter
+	default:
+		*r = Unknown
+	}
+	return nil
+}
+
+// GetPeersByRule returns the peers with specified rule.
+func (r *RegionInfo) GetPeersByRule(rule Rule) []*metapb.Peer {
+	switch rule {
+	case LeaderScatter:
 		return []*metapb.Peer{r.GetLeader()}
-	case Follower:
-		followers := r.GetFollowers()
-		ret := make([]*metapb.Peer, 0, len(followers))
-		for _, peer := range followers {
-			ret = append(ret, peer)
-		}
-		return ret
-	case Learner:
-		learners := r.GetLearners()
-		return learners
+	case PeerScatter:
+		return r.GetPeers()
+	case LearnerScatter:
+		return r.GetLearners()
 	default:
 		return nil
 	}

pkg/core/region_test.go

@@ -16,11 +16,13 @@ package core
 
 import (
 	"crypto/rand"
+	"encoding/json"
 	"fmt"
 	"math"
 	mrand "math/rand"
 	"sort"
 	"strconv"
+	"strings"
 	"testing"
 	"time"
 
@@ -1330,31 +1332,47 @@ func TestGetPeers(t *testing.T) {
 		leader, follower1, follower2, learner,
 	}}, leader, WithLearners([]*metapb.Peer{learner}))
 	for _, v := range []struct {
-		role  string
+		rule  string
 		peers []*metapb.Peer
 	}{
 		{
-			role:  "leader",
+			rule:  "leader-scatter",
 			peers: []*metapb.Peer{leader},
 		},
 		{
-			role:  "follower",
-			peers: []*metapb.Peer{follower1, follower2},
+			rule:  "peer-scatter",
+			peers: []*metapb.Peer{learner, leader, follower1, follower2},
 		},
 		{
-			role:  "learner",
+			rule:  "learner-scatter",
 			peers: []*metapb.Peer{learner},
 		},
 		{
-			role:  "witness",
+			rule:  "witness-scatter",
 			peers: nil,
 		},
 	} {
-		role := NewRole(v.role)
-		peers := region.GetPeersByRole(role)
+		role := NewRule(v.rule)
+		peers := region.GetPeersByRule(role)
 		sort.Slice(peers, func(i, j int) bool {
 			return peers[i].Id <= peers[j].Id
 		})
 		re.Equal(v.peers, peers, role)
 	}
 }
+
+func TestCodecRule(t *testing.T) {
+	re := require.New(t)
+	for _, v := range []string{"leader", "peer", "learner", "witness"} {
+		rule := NewRule(v)
+		if rule != Unknown {
+			re.Equal(rule.String(), v)
+		}
+		body, err := json.Marshal(&rule)
+		re.NoError(err)
+		re.Equal(strings.Join([]string{"\"", rule.String(), "\""}, ""), string(body))
+		var rule2 Rule
+		re.NoError(json.Unmarshal(body, &rule2))
+		re.Equal(rule.String(), rule2.String())
+	}
+}

pkg/schedule/operator/influence.go

@@ -60,13 +60,13 @@
 }
 
 // GetStoreInfluenceByRole returns the influence of the store according to the role.
-func (s *StoreInfluence) GetStoreInfluenceByRole(r core.Role) int64 {
+func (s *StoreInfluence) GetStoreInfluenceByRole(r core.Rule) int64 {
 	switch r {
-	case core.Leader:
+	case core.LeaderScatter:
 		return s.LeaderCount
-	case core.Follower:
+	case core.PeerScatter:
 		return s.RegionCount
-	case core.Learner:
+	case core.LearnerScatter:
 		return s.RegionCount
 	default:
 		return 0