pd-ctl: add policy selection for split region (#1149)

* add policy for split region

* update kvproto

Gopkg.toml

@@ -27,3 +27,4 @@
 [[constraint]]
   name = "github.com/pingcap/kvproto"
   branch = "master"
+

pdctl/command/operator.go

@@ -262,10 +262,11 @@ func removePeerCommandFunc(cmd *cobra.Command, args []string) {
 // NewSplitRegionCommand returns a command to split a region.
 func NewSplitRegionCommand() *cobra.Command {
 	c := &cobra.Command{
-		Use:   "split-region <region_id>",
+		Use:   "split-region <region_id> [--policy=scan|approximate]",
 		Short: "split a region",
 		Run:   splitRegionCommandFunc,
 	}
+	c.Flags().String("policy", "scan", "the policy to get region split key")
 	return c
 }
 
@@ -281,9 +282,19 @@ func splitRegionCommandFunc(cmd *cobra.Command, args []string) {
 		return
 	}
 
+	policy := cmd.Flags().Lookup("policy").Value.String()
+	switch policy {
+	case "scan", "approximate":
+		break
+	default:
+		fmt.Println("Error: unknown policy")
+		return
+	}
+
 	input := make(map[string]interface{})
 	input["name"] = cmd.Name()
 	input["region_id"] = ids[0]
+	input["policy"] = policy
 	postJSON(cmd, operatorsPrefix, input)
 }
 

server/api/operator.go

@@ -206,7 +206,12 @@ func (h *operatorHandler) Post(w http.ResponseWriter, r *http.Request) {
 			h.r.JSON(w, http.StatusBadRequest, "missing region id")
 			return
 		}
-		if err := h.AddSplitRegionOperator(uint64(regionID)); err != nil {
+		policy, ok := input["policy"].(string)
+		if !ok {
+			h.r.JSON(w, http.StatusBadRequest, "missing split policy")
+			return
+		}
+		if err := h.AddSplitRegionOperator(uint64(regionID), policy); err != nil {
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
 			return
 		}

server/coordinator.go

@@ -639,7 +639,9 @@ func (c *coordinator) sendScheduleCommand(region *core.RegionInfo, step schedule
 		c.hbStreams.sendMsg(region, cmd)
 	case schedule.SplitRegion:
 		cmd := &pdpb.RegionHeartbeatResponse{
-			SplitRegion: &pdpb.SplitRegion{},
+			SplitRegion: &pdpb.SplitRegion{
+				Policy: s.Policy,
+			},
 		}
 		c.hbStreams.sendMsg(region, cmd)
 	default:

server/handler.go

@@ -16,10 +16,12 @@ package server
 import (
 	"bytes"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/juju/errors"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/pd/server/core"
 	"github.com/pingcap/pd/server/schedule"
 	log "github.com/sirupsen/logrus"
@@ -516,7 +518,7 @@ func (h *Handler) AddMergeRegionOperator(regionID uint64, targetID uint64) error
 }
 
 // AddSplitRegionOperator adds an operator to split a region.
-func (h *Handler) AddSplitRegionOperator(regionID uint64) error {
+func (h *Handler) AddSplitRegionOperator(regionID uint64, policy string) error {
 	c, err := h.getCoordinator()
 	if err != nil {
 		return errors.Trace(err)
@@ -527,7 +529,11 @@ func (h *Handler) AddSplitRegionOperator(regionID uint64) error {
 		return ErrRegionNotFound(regionID)
 	}
 
-	step := schedule.SplitRegion{StartKey: region.StartKey, EndKey: region.EndKey}
+	step := schedule.SplitRegion{
+		StartKey: region.StartKey,
+		EndKey:   region.EndKey,
+		Policy:   pdpb.CheckPolicy(pdpb.CheckPolicy_value[strings.ToUpper(policy)]),
+	}
 	op := schedule.NewOperator("adminSplitRegion", regionID, region.GetRegionEpoch(), schedule.OpAdmin, step)
 	if ok := c.addOperator(op); !ok {
 		return errors.Trace(errAddOperator)

server/schedule/operator.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/juju/errors"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/pingcap/kvproto/pkg/pdpb"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/pingcap/pd/server/core"
@@ -212,10 +213,11 @@ func (mr MergeRegion) Influence(opInfluence OpInfluence, region *core.RegionInfo
 // SplitRegion is an OperatorStep that splits a region.
 type SplitRegion struct {
 	StartKey, EndKey []byte
+	Policy           pdpb.CheckPolicy
 }
 
 func (sr SplitRegion) String() string {
-	return "split region"
+	return fmt.Sprintf("split region with policy %s", sr.Policy.String())
 }
 
 // IsFinish checks if current step is finished.