server/api: add the api to show the regions in subtree by type

pkg/core/region.go

@@ -1541,6 +1541,60 @@
 	return regions
 }
 
+// SubTreeRegionType is the type of sub tree region.
+type SubTreeRegionType string
+
+const (
+	// AllInSubTree is all sub trees.
+	AllInSubTree SubTreeRegionType = "all"
+	// LeaderInSubTree is the leader sub tree.
+	LeaderInSubTree SubTreeRegionType = "leader"
+	// FollowerInSubTree is the follower sub tree.
+	FollowerInSubTree SubTreeRegionType = "follower"
+	// LearnerInSubTree is the learner sub tree.
+	LearnerInSubTree SubTreeRegionType = "learner"
+	// WitnessInSubTree is the witness sub tree.
+	WitnessInSubTree SubTreeRegionType = "witness"
+	// PendingPeerInSubTree is the pending peer sub tree.
+	PendingPeerInSubTree SubTreeRegionType = "pending"
+)
+
+// GetStoreRegions gets all RegionInfo with a given storeID
+func (r *RegionsInfo) GetStoreRegionsByTypeInSubTree(storeID uint64, typ SubTreeRegionType) ([]*RegionInfo, error) {
+	r.st.RLock()
+	var regions []*RegionInfo
+	switch typ {
+	case LeaderInSubTree:
+		if leaders, ok := r.leaders[storeID]; ok {
+			regions = leaders.scanRanges()
+		}
+	case FollowerInSubTree:
+		if followers, ok := r.followers[storeID]; ok {
+			regions = followers.scanRanges()
+		}
+	case LearnerInSubTree:
+		if learners, ok := r.learners[storeID]; ok {
+			regions = learners.scanRanges()
+		}
+	case WitnessInSubTree:
+		if witnesses, ok := r.witnesses[storeID]; ok {
+			regions = witnesses.scanRanges()
+		}
+	case PendingPeerInSubTree:
+		if pendingPeers, ok := r.pendingPeers[storeID]; ok {
+			regions = pendingPeers.scanRanges()
+		}
+	case AllInSubTree:
+		r.st.RUnlock()
+		return r.GetStoreRegions(storeID), nil
+	default:
+		return nil, errors.Errorf("unknown sub tree region type %v", typ)
+	}
+
+	r.st.RUnlock()
+	return regions, nil
+}
+
 // GetStoreLeaderRegionSize get total size of store's leader regions
 func (r *RegionsInfo) GetStoreLeaderRegionSize(storeID uint64) int64 {
 	r.st.RLock()

pkg/response/store.go

@@ -64,6 +64,7 @@ type StoreStatus struct {
 	RegionSize         int64              `json:"region_size"`
 	LearnerCount       int                `json:"learner_count,omitempty"`
 	WitnessCount       int                `json:"witness_count,omitempty"`
+	PendingPeerCount   int                `json:"pending_peer_count,omitempty"`
 	SlowScore          uint64             `json:"slow_score,omitempty"`
 	SlowTrend          *SlowTrend         `json:"slow_trend,omitempty"`
 	SendingSnapCount   uint32             `json:"sending_snap_count,omitempty"`
@@ -117,6 +118,7 @@ func BuildStoreInfo(opt *sc.ScheduleConfig, store *core.StoreInfo) *StoreInfo {
 			SlowTrend:          slowTrend,
 			SendingSnapCount:   store.GetSendingSnapCount(),
 			ReceivingSnapCount: store.GetReceivingSnapCount(),
+			PendingPeerCount:   store.GetPendingPeerCount(),
 			IsBusy:             store.IsBusy(),
 		},
 	}

server/api/region.go

@@ -218,7 +218,17 @@
 		h.rd.JSON(w, http.StatusBadRequest, err.Error())
 		return
 	}
-	regions := rc.GetStoreRegions(uint64(id))
+	// get type from query
+	typ := r.URL.Query().Get("type")
+	if len(typ) == 0 {
+		typ = string(core.AllInSubTree)
+	}
+
+	regions, err := rc.GetStoreRegionsByTypeInSubTree(uint64(id), core.SubTreeRegionType(typ))
+	if err != nil {
+		h.rd.JSON(w, http.StatusBadRequest, err.Error())
+		return
+	}
 	b, err := response.MarshalRegionsInfoJSON(r.Context(), regions)
 	if err != nil {
 		h.rd.JSON(w, http.StatusInternalServerError, err.Error())

server/cluster/cluster.go

@@ -1198,6 +1198,11 @@
 	return c.core.GetStoreRegions(storeID)
 }
 
+// GetStoreRegions returns all regions' information with a given storeID.
+func (c *RaftCluster) GetStoreRegionsByType(storeID uint64) []*core.RegionInfo {
+	return c.core.GetStoreRegions(storeID)
+}
+
 // RandLeaderRegions returns some random regions that has leader on the store.
 func (c *RaftCluster) RandLeaderRegions(storeID uint64, ranges []core.KeyRange) []*core.RegionInfo {
 	return c.core.RandLeaderRegions(storeID, ranges)

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

@@ -486,6 +486,7 @@ func NewRegionWithStoreCommand() *cobra.Command {
 		Short: "show the regions of a specific store",
 		Run:   showRegionWithStoreCommandFunc,
 	}
+	r.Flags().String("type", "all", "the type of the regions, could be 'all', 'leader', 'learner' or 'pending'")
 	return r
 }
 
@@ -496,6 +497,8 @@ func showRegionWithStoreCommandFunc(cmd *cobra.Command, args []string) {
 	}
 	storeID := args[0]
 	prefix := regionsStorePrefix + "/" + storeID
+	flagType := cmd.Flag("type")
+	prefix += "?type=" + flagType.Value.String()
 	r, err := doRequest(cmd, prefix, http.MethodGet, http.Header{})
 	if err != nil {
 		cmd.Printf("Failed to get regions with the given storeID: %s\n", err)

tools/pd-ctl/tests/region/region_test.go

@@ -108,6 +108,11 @@ func TestRegion(t *testing.T) {
 	)
 	defer cluster.Destroy()
 
+	getRegionsByType := func(storeID uint64, regionType core.SubTreeRegionType) []*core.RegionInfo {
+		regions, _ := leaderServer.GetRaftCluster().GetStoreRegionsByTypeInSubTree(storeID, regionType)
+		return regions
+	}
+
 	var testRegionsCases = []struct {
 		args   []string
 		expect []*core.RegionInfo
@@ -118,7 +123,12 @@ func TestRegion(t *testing.T) {
 		{[]string{"region", "sibling", "2"}, leaderServer.GetAdjacentRegions(leaderServer.GetRegionInfoByID(2))},
 		// region store <store_id> command
 		{[]string{"region", "store", "1"}, leaderServer.GetStoreRegions(1)},
-		{[]string{"region", "store", "1"}, []*core.RegionInfo{r1, r2, r3, r4}},
+		{[]string{"region", "store", "1", "--type=leader"}, getRegionsByType(1, core.LeaderInSubTree)},
+		{[]string{"region", "store", "1", "--type=follower"}, getRegionsByType(1, core.FollowerInSubTree)},
+		{[]string{"region", "store", "1", "--type=learner"}, getRegionsByType(1, core.LearnerInSubTree)},
+		{[]string{"region", "store", "1", "--type=witness"}, getRegionsByType(1, core.WitnessInSubTree)},
+		{[]string{"region", "store", "1", "--type=pending"}, getRegionsByType(1, core.PendingPeerInSubTree)},
+		{[]string{"region", "store", "1", "--type=all"}, []*core.RegionInfo{r1, r2, r3, r4}},
 		// region check extra-peer command
 		{[]string{"region", "check", "extra-peer"}, []*core.RegionInfo{r1}},
 		// region check miss-peer command