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

Signed-off-by: nolouch <[email protected]>

pkg/core/region.go

@@ -1368,6 +1368,56 @@ func (r *RegionsInfo) GetStoreRegions(storeID uint64) []*RegionInfo {
 	return regions
 }
 
+// SubTreeRegionType is the type of sub tree region.
+type SubTreeRegionType int
+
+var (
+	// AllInSubTree is all sub trees.
+	AllInSubTree SubTreeRegionType = 0
+	// LeaderInSubTree is the leader sub tree.
+	LeaderInSubTree SubTreeRegionType = 1
+	// FollowerInSubTree is the follower sub tree.
+	FollowerInSubTree SubTreeRegionType = 2
+	// LearnerInSubTree is the learner sub tree.
+	LearnerInSubTree SubTreeRegionType = 3
+	// WitnessInSubTree is the witness sub tree.
+	WitnessInSubTree SubTreeRegionType = 4
+	// PendingPeerInSubTree is the pending peer sub tree.
+	PendingPeerInSubTree SubTreeRegionType = 5
+)
+
+// GetStoreRegions gets all RegionInfo with a given storeID
+func (r *RegionsInfo) GetStoreRegionsByTypeInSubTree(storeID uint64, typ SubTreeRegionType) []*RegionInfo {
+	r.st.RLock()
+	defer r.st.RUnlock()
+	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:
+		return r.GetStoreRegions(storeID)
+	}
+	return regions
+}
+
 // 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,12 @@ func (h *regionsHandler) GetStoreRegions(w http.ResponseWriter, r *http.Request)
 		h.rd.JSON(w, http.StatusBadRequest, err.Error())
 		return
 	}
-	regions := rc.GetStoreRegions(uint64(id))
+	// get type from query
+	typID, err := strconv.Atoi(r.URL.Query().Get("type"))
+	if err != nil {
+		typID = int(core.AllInSubTree)
+	}
+	regions := rc.GetStoreRegionsByTypeInSubTree(uint64(id), core.SubTreeRegionType(typID))
 	b, err := response.MarshalRegionsInfoJSON(r.Context(), regions)
 	if err != nil {
 		h.rd.JSON(w, http.StatusInternalServerError, err.Error())

server/cluster/cluster.go

@@ -1204,6 +1204,11 @@ func (c *RaftCluster) GetStoreRegions(storeID uint64) []*core.RegionInfo {
 	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,26 @@ func showRegionWithStoreCommandFunc(cmd *cobra.Command, args []string) {
 	}
 	storeID := args[0]
 	prefix := regionsStorePrefix + "/" + storeID
+	flag := cmd.Flag("type")
+	typ := 0
+	switch flag.Value.String() {
+	case "all":
+		typ = 0
+	case "leader":
+		typ = 1
+	case "follower":
+		typ = 2
+	case "learner":
+		typ = 3
+	case "witness":
+		typ = 4
+	case "pending":
+		typ = 5
+	default:
+		cmd.Println("unknown type")
+		return
+	}
+	prefix += "?type=" + strconv.Itoa(typ)
 	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

@@ -118,7 +118,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"}, leaderServer.GetRaftCluster().GetStoreRegionsByTypeInSubTree(1, core.LeaderInSubTree)},
+		{[]string{"region", "store", "1", "--type=follower"}, leaderServer.GetRaftCluster().GetStoreRegionsByTypeInSubTree(1, core.FollowerInSubTree)},
+		{[]string{"region", "store", "1", "--type=learner"}, leaderServer.GetRaftCluster().GetStoreRegionsByTypeInSubTree(1, core.LearnerInSubTree)},
+		{[]string{"region", "store", "1", "--type=witness"}, leaderServer.GetRaftCluster().GetStoreRegionsByTypeInSubTree(1, core.WitnessInSubTree)},
+		{[]string{"region", "store", "1", "--type=pending"}, leaderServer.GetRaftCluster().GetStoreRegionsByTypeInSubTree(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