diff --git a/pkg/core/region.go b/pkg/core/region.go index 713e82cc36d..8e3d4e5dec8 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1541,6 +1541,60 @@ func (r *RegionsInfo) GetStoreRegions(storeID uint64) []*RegionInfo { 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() diff --git a/pkg/response/store.go b/pkg/response/store.go index 1efe11bfb39..8bff1e75e42 100644 --- a/pkg/response/store.go +++ b/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(), }, } diff --git a/server/api/region.go b/server/api/region.go index dac92f247ca..974b5e4fa12 100644 --- a/server/api/region.go +++ b/server/api/region.go @@ -218,7 +218,17 @@ 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 + 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()) diff --git a/server/cluster/cluster.go b/server/cluster/cluster.go index dbc6a6cadf3..75301664b50 100644 --- a/server/cluster/cluster.go +++ b/server/cluster/cluster.go @@ -1198,6 +1198,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) diff --git a/tools/pd-ctl/pdctl/command/region_command.go b/tools/pd-ctl/pdctl/command/region_command.go index d7e19967c7a..3536b01a606 100644 --- a/tools/pd-ctl/pdctl/command/region_command.go +++ b/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) diff --git a/tools/pd-ctl/tests/region/region_test.go b/tools/pd-ctl/tests/region/region_test.go index b328fd88286..2952e137f3b 100644 --- a/tools/pd-ctl/tests/region/region_test.go +++ b/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 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