pdctl: add placement rules commands (#2306)

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

tests/pdctl/config/config_test.go

@@ -16,15 +16,18 @@ package config_test
 import (
 	"context"
 	"encoding/json"
+	"io/ioutil"
 	"reflect"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/coreos/go-semver/semver"
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/pd/v3/server"
 	"github.com/pingcap/pd/v3/server/config"
+	"github.com/pingcap/pd/v3/server/schedule/placement"
 	"github.com/pingcap/pd/v3/tests"
 	"github.com/pingcap/pd/v3/tests/pdctl"
 )
@@ -203,3 +206,93 @@ func (s *configTestSuite) TestConfig(c *C) {
 	_, _, err = pdctl.ExecuteCommandC(cmd, args1...)
 	c.Assert(err, IsNil)
 }
+
+func (s *configTestSuite) TestPlacementRules(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	cluster, err := tests.NewTestCluster(ctx, 1)
+	c.Assert(err, IsNil)
+	err = cluster.RunInitialServers()
+	c.Assert(err, IsNil)
+	cluster.WaitLeader()
+	pdAddr := cluster.GetConfig().GetClientURLs()
+	cmd := pdctl.InitCommand()
+
+	store := metapb.Store{
+		Id:    1,
+		State: metapb.StoreState_Up,
+	}
+	leaderServer := cluster.GetServer(cluster.GetLeader())
+	c.Assert(leaderServer.BootstrapCluster(), IsNil)
+	svr := leaderServer.GetServer()
+	pdctl.MustPutStore(c, svr, store.Id, store.State, store.Labels)
+	defer cluster.Destroy()
+
+	_, output, err := pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "enable")
+	c.Assert(err, IsNil)
+	c.Assert(strings.Contains(string(output), "Success!"), IsTrue)
+
+	// wait config manager reload
+	time.Sleep(time.Second)
+
+	// test show
+	var rules []placement.Rule
+	_, output, err = pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "show")
+	c.Assert(err, IsNil)
+	err = json.Unmarshal(output, &rules)
+	c.Assert(err, IsNil)
+	c.Assert(rules, HasLen, 1)
+	c.Assert(rules[0].Key(), Equals, [2]string{"pd", "default"})
+
+	f, _ := ioutil.TempFile("/tmp", "pd_tests")
+	fname := f.Name()
+	f.Close()
+
+	// test load
+	_, _, err = pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "load", "--out="+fname)
+	c.Assert(err, IsNil)
+	b, _ := ioutil.ReadFile(fname)
+	c.Assert(json.Unmarshal(b, &rules), IsNil)
+	c.Assert(rules, HasLen, 1)
+	c.Assert(rules[0].Key(), Equals, [2]string{"pd", "default"})
+
+	// test save
+	rules = append(rules, placement.Rule{
+		GroupID: "pd",
+		ID:      "test1",
+		Role:    "voter",
+		Count:   1,
+	}, placement.Rule{
+		GroupID: "test-group",
+		ID:      "test2",
+		Role:    "voter",
+		Count:   2,
+	})
+	b, _ = json.Marshal(rules)
+	ioutil.WriteFile(fname, b, 0644)
+	_, _, err = pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "save", "--in="+fname)
+	c.Assert(err, IsNil)
+
+	// test show group
+	var rules2 []placement.Rule
+	_, output, err = pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "show", "--group=pd")
+	c.Assert(err, IsNil)
+	err = json.Unmarshal(output, &rules2)
+	c.Assert(err, IsNil)
+	c.Assert(rules2, HasLen, 2)
+	c.Assert(rules2[0].Key(), Equals, [2]string{"pd", "default"})
+	c.Assert(rules2[1].Key(), Equals, [2]string{"pd", "test1"})
+
+	// test delete
+	rules[0].Count = 0
+	b, _ = json.Marshal(rules)
+	ioutil.WriteFile(fname, b, 0644)
+	_, _, err = pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "save", "--in="+fname)
+	c.Assert(err, IsNil)
+	_, output, err = pdctl.ExecuteCommandC(cmd, "-u", pdAddr, "config", "placement-rules", "show", "--group=pd")
+	c.Assert(err, IsNil)
+	err = json.Unmarshal(output, &rules)
+	c.Assert(err, IsNil)
+	c.Assert(rules, HasLen, 1)
+	c.Assert(rules[0].Key(), Equals, [2]string{"pd", "test1"})
+}

tools/pd-ctl/pdctl/command/config_command.go

@@ -16,10 +16,13 @@ package command
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
+	"io/ioutil"
 	"net/http"
 	"path"
 	"strconv"
 
+	"github.com/pingcap/pd/v3/server/schedule/placement"
 	"github.com/spf13/cobra"
 )
 
@@ -29,6 +32,8 @@ var (
 	replicationPrefix    = "pd/api/v1/config/replicate"
 	labelPropertyPrefix  = "pd/api/v1/config/label-property"
 	clusterVersionPrefix = "pd/api/v1/config/cluster-version"
+	rulesPrefix          = "pd/api/v1/config/rules"
+	rulePrefix           = "pd/api/v1/config/rule"
 )
 
 // NewConfigCommand return a config subcommand of rootCmd
@@ -40,6 +45,7 @@ func NewConfigCommand() *cobra.Command {
 	conf.AddCommand(NewShowConfigCommand())
 	conf.AddCommand(NewSetConfigCommand())
 	conf.AddCommand(NewDeleteConfigCommand())
+	conf.AddCommand(NewPlacementRulesCommand())
 	return conf
 }
 
@@ -297,3 +303,147 @@ func setClusterVersionCommandFunc(cmd *cobra.Command, args []string) {
 	}
 	postJSON(cmd, clusterVersionPrefix, input)
 }
+
+// NewPlacementRulesCommand placement rules subcommand
+func NewPlacementRulesCommand() *cobra.Command {
+	c := &cobra.Command{
+		Use:   "placement-rules",
+		Short: "placement rules configuration",
+	}
+	enable := &cobra.Command{
+		Use:   "enable",
+		Short: "enable placement rules",
+		Run:   enablePlacementRulesFunc,
+	}
+	disable := &cobra.Command{
+		Use:   "disable",
+		Short: "disable placement rules",
+		Run:   disablePlacementRulesFunc,
+	}
+	show := &cobra.Command{
+		Use:   "show",
+		Short: "show placement rules",
+		Run:   getPlacementRulesFunc,
+	}
+	show.Flags().String("group", "", "group id")
+	show.Flags().String("id", "", "rule id")
+	show.Flags().String("region", "", "region id")
+	load := &cobra.Command{
+		Use:   "load",
+		Short: "load placement rules to a file",
+		Run:   getPlacementRulesFunc,
+	}
+	load.Flags().String("group", "", "group id")
+	load.Flags().String("id", "", "rule id")
+	load.Flags().String("region", "", "region id")
+	load.Flags().String("out", "rules.json", "the filename contains rules")
+	save := &cobra.Command{
+		Use:   "save",
+		Short: "save rules from file",
+		Run:   putPlacementRulesFunc,
+	}
+	save.Flags().String("in", "rules.json", "the filename contains rules")
+	c.AddCommand(enable, disable, show, load, save)
+	return c
+}
+
+func enablePlacementRulesFunc(cmd *cobra.Command, args []string) {
+	err := postConfigDataWithPath(cmd, "enable-placement-rules", "true", configPrefix)
+	if err != nil {
+		cmd.Printf("Failed to set config: %s\n", err)
+		return
+	}
+	cmd.Println("Success!")
+}
+
+func disablePlacementRulesFunc(cmd *cobra.Command, args []string) {
+	err := postConfigDataWithPath(cmd, "enable-placement-rules", "false", configPrefix)
+	if err != nil {
+		cmd.Printf("Failed to set config: %s\n", err)
+		return
+	}
+	cmd.Println("Success!")
+}
+
+func getPlacementRulesFunc(cmd *cobra.Command, args []string) {
+	getFlag := func(key string) string {
+		if f := cmd.Flag(key); f != nil {
+			return f.Value.String()
+		}
+		return ""
+	}
+
+	group, id, region, file := getFlag("group"), getFlag("id"), getFlag("region"), getFlag("out")
+	var reqPath string
+	respIsList := true
+	switch {
+	case region == "" && group == "" && id == "": // all rules
+		reqPath = rulesPrefix
+	case region == "" && group == "" && id != "":
+		cmd.Println(`"id" should be specified along with "group"`)
+		return
+	case region == "" && group != "" && id == "": // all rules in a group
+		reqPath = path.Join(rulesPrefix, "group", group)
+	case region == "" && group != "" && id != "": // single rule
+		reqPath, respIsList = path.Join(rulePrefix, group, id), false
+	case region != "" && group == "" && id == "": // rules matches a region
+		reqPath = path.Join(rulesPrefix, "region", region)
+	default:
+		cmd.Println(`"region" should not be specified with "group" or "id" at the same time`)
+		return
+	}
+	res, err := doRequest(cmd, reqPath, http.MethodGet)
+	if err != nil {
+		cmd.Println(err)
+		return
+	}
+	if file == "" {
+		cmd.Println(res)
+		return
+	}
+	if !respIsList {
+		res = "[\n" + res + "]\n"
+	}
+	err = ioutil.WriteFile(file, []byte(res), 0644)
+	if err != nil {
+		cmd.Println(err)
+		return
+	}
+	cmd.Println("rules saved to file " + file)
+}
+
+func putPlacementRulesFunc(cmd *cobra.Command, args []string) {
+	var file string
+	if f := cmd.Flag("in"); f != nil {
+		file = f.Value.String()
+	}
+	content, err := ioutil.ReadFile(file)
+	if err != nil {
+		cmd.Println(err)
+		return
+	}
+	var rules []*placement.Rule
+	if err = json.Unmarshal(content, &rules); err != nil {
+		cmd.Println(err)
+		return
+	}
+	for _, r := range rules {
+		if r.Count > 0 {
+			b, _ := json.Marshal(r)
+			_, err = doRequest(cmd, rulePrefix, http.MethodPost, WithBody("application/json", bytes.NewBuffer(b)))
+			if err != nil {
+				fmt.Printf("failed to save rule %s/%s: %v\n", r.GroupID, r.ID, err)
+				return
+			}
+			fmt.Printf("saved rule %s/%s\n", r.GroupID, r.ID)
+		} else {
+			_, err = doRequest(cmd, path.Join(rulePrefix, r.GroupID, r.ID), http.MethodDelete)
+			if err != nil {
+				fmt.Printf("failed to delete rule %s/%s: %v\n", r.GroupID, r.ID, err)
+				return
+			}
+			fmt.Printf("deleted rule %s/%s\n", r.GroupID, r.ID)
+		}
+	}
+	cmd.Println("Success!")
+}