ctl: pd-ctl should support to check regions that's isolation label eq…

…uals the given labels (#4697) close #4696 Signed-off-by: bufferflies <1045931706@qq.com> Co-authored-by: Ti Chi Robot <ti-community-prow-bot@tidb.io>
tikv · Apr 11, 2022 · 13c68cf · 13c68cf
1 parent e3fd147
commit 13c68cf
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 6 deletions.
diff --git a/server/statistics/region_collection.go b/server/statistics/region_collection.go
@@ -265,7 +265,7 @@ func NewLabelStatistics() *LabelStatistics {
 // Observe records the current label status.
 func (l *LabelStatistics) Observe(region *core.RegionInfo, stores []*core.StoreInfo, labels []string) {
 	regionID := region.GetID()
-	regionIsolation := getRegionLabelIsolation(stores, labels)
+	regionIsolation := GetRegionLabelIsolation(stores, labels)
 	if label, ok := l.regionLabelStats[regionID]; ok {
 		if label == regionIsolation {
 			return
@@ -301,7 +301,8 @@ func (l *LabelStatistics) GetLabelCounter() map[string]int {
 	return l.labelCounter
 }
 
-func getRegionLabelIsolation(stores []*core.StoreInfo, labels []string) string {
+// GetRegionLabelIsolation returns the isolation level of the region.
+func GetRegionLabelIsolation(stores []*core.StoreInfo, labels []string) string {
 	if len(stores) == 0 || len(labels) == 0 {
 		return nonIsolation
 	}

diff --git a/server/statistics/region_collection_test.go b/server/statistics/region_collection_test.go
@@ -252,7 +252,7 @@ func (t *testRegionStatisticsSuite) TestRegionLabelIsolationLevel(c *C) {
 			stores = append(stores, s)
 		}
 		region := core.NewRegionInfo(&metapb.Region{Id: uint64(regionID)}, nil)
-		label := getRegionLabelIsolation(stores, locationLabels)
+		label := GetRegionLabelIsolation(stores, locationLabels)
 		labelLevelStats.Observe(region, stores, locationLabels)
 		c.Assert(label, Equals, res)
 		regionID++
@@ -265,12 +265,12 @@ func (t *testRegionStatisticsSuite) TestRegionLabelIsolationLevel(c *C) {
 		c.Assert(labelLevelStats.labelCounter[i], Equals, res)
 	}
 
-	label := getRegionLabelIsolation(nil, locationLabels)
+	label := GetRegionLabelIsolation(nil, locationLabels)
 	c.Assert(label, Equals, nonIsolation)
-	label = getRegionLabelIsolation(nil, nil)
+	label = GetRegionLabelIsolation(nil, nil)
 	c.Assert(label, Equals, nonIsolation)
 	store := core.NewStoreInfo(&metapb.Store{Id: 1, Address: "mock://tikv-1"}, core.SetStoreLabels([]*metapb.StoreLabel{{Key: "foo", Value: "bar"}}))
-	label = getRegionLabelIsolation([]*core.StoreInfo{store}, locationLabels)
+	label = GetRegionLabelIsolation([]*core.StoreInfo{store}, locationLabels)
 	c.Assert(label, Equals, "zone")
 
 	regionID = 1

diff --git a/tests/pdctl/label/label_test.go b/tests/pdctl/label/label_test.go
@@ -144,4 +144,11 @@ func (s *labelTestSuite) TestLabel(c *C) {
 	c.Assert(json.Unmarshal(output, &storesInfo), IsNil)
 	sss := []*api.StoreInfo{stores[0], stores[2]}
 	pdctl.CheckStoresInfo(c, storesInfo.Stores, sss)
+
+	// label isolation [label]
+	args = []string{"-u", pdAddr, "label", "isolation"}
+	output, err = pdctl.ExecuteCommand(cmd, args...)
+	c.Assert(err, IsNil)
+	c.Assert(strings.Contains(string(output), "none"), IsTrue)
+	c.Assert(strings.Contains(string(output), "2"), IsTrue)
 }
diff --git a/tools/pd-ctl/pdctl/command/label_command.go b/tools/pd-ctl/pdctl/command/label_command.go
@@ -15,10 +15,15 @@
 package command
 
 import (
+	"encoding/json"
 	"fmt"
 	"net/http"
 
 	"github.com/spf13/cobra"
+	"github.com/tikv/pd/server/api"
+	"github.com/tikv/pd/server/config"
+	"github.com/tikv/pd/server/core"
+	"github.com/tikv/pd/server/statistics"
 )
 
 var (
@@ -34,6 +39,7 @@ func NewLabelCommand() *cobra.Command {
 		Run:   showLabelsCommandFunc,
 	}
 	l.AddCommand(NewLabelListStoresCommand())
+	l.AddCommand(NewCheckLabels())
 	return l
 }
 
@@ -78,3 +84,105 @@ func showLabelListStoresCommandFunc(cmd *cobra.Command, args []string) {
 	}
 	cmd.Println(r)
 }
+
+// NewCheckLabels returns a isolation label check.
+func NewCheckLabels() *cobra.Command {
+	return &cobra.Command{
+		Use:     "isolation [label]",
+		Short:   "isolation labels",
+		Example: "label-count: map[isolation-label:count], region-map: [region-id:isolation-label]",
+		Run:     checkIsolationLabel,
+	}
+}
+
+func getReplicationConfig(cmd *cobra.Command, _ []string) (*config.ReplicationConfig, error) {
+	prefix := configPrefix + "/replicate"
+	body, err := doRequest(cmd, prefix, http.MethodGet, http.Header{})
+	if err != nil {
+		return nil, err
+	}
+	var config config.ReplicationConfig
+	if err := json.Unmarshal([]byte(body), &config); err != nil {
+		return nil, err
+	}
+	return &config, nil
+}
+
+func getStores(cmd *cobra.Command, _ []string) ([]*core.StoreInfo, error) {
+	prefix := storesPrefix
+	body, err := doRequest(cmd, prefix, http.MethodGet, http.Header{})
+	if err != nil {
+		return nil, err
+	}
+	var storesInfo api.StoresInfo
+	if err := json.Unmarshal([]byte(body), &storesInfo); err != nil {
+		return nil, err
+	}
+	stores := make([]*core.StoreInfo, 0)
+	for _, storeInfo := range storesInfo.Stores {
+		stores = append(stores, core.NewStoreInfo(storeInfo.Store.Store))
+	}
+	return stores, nil
+}
+
+func getRegions(cmd *cobra.Command, _ []string) ([]api.RegionInfo, error) {
+	prefix := regionsPrefix
+	body, err := doRequest(cmd, prefix, http.MethodGet, http.Header{})
+	if err != nil {
+		return nil, err
+	}
+	var RegionsInfo api.RegionsInfo
+	if err := json.Unmarshal([]byte(body), &RegionsInfo); err != nil {
+		return nil, err
+	}
+	return RegionsInfo.Regions, nil
+}
+
+func checkIsolationLabel(cmd *cobra.Command, args []string) {
+	var checkLabel string
+	if len(args) == 1 {
+		checkLabel = args[0]
+	}
+
+	storesInfo, err := getStores(cmd, args)
+	if err != nil {
+		cmd.Printf("Failed to get stores info: %s\n", err)
+		return
+	}
+
+	regionsInfo, err := getRegions(cmd, args)
+	if err != nil {
+		cmd.Printf("Failed to get regions info: %s\n", err)
+		return
+	}
+
+	config, err := getReplicationConfig(cmd, args)
+	if err != nil {
+		cmd.Printf("Failed to get labels: %s\n", err)
+		return
+	}
+	locationLabels := config.LocationLabels
+
+	storesMap := make(map[uint64]*core.StoreInfo, len(storesInfo))
+	for _, store := range storesInfo {
+		storesMap[store.GetID()] = store
+	}
+
+	regionMap := make(map[uint64]string, len(regionsInfo))
+	labelCount := make(map[string]int, len(locationLabels))
+
+	for _, region := range regionsInfo {
+		stores := make([]*core.StoreInfo, 0)
+		for _, peer := range region.Peers {
+			if s, ok := storesMap[peer.StoreId]; ok {
+				stores = append(stores, s)
+			}
+		}
+		isolationLabel := statistics.GetRegionLabelIsolation(stores, locationLabels)
+		if len(checkLabel) == 0 || isolationLabel == checkLabel {
+			regionMap[region.ID] = isolationLabel
+			labelCount[isolationLabel]++
+		}
+	}
+	cmd.Printf("label-count:%+v,region-map:%+v\n", labelCount, regionMap)
+}