diff --git a/pkg/core/region.go b/pkg/core/region.go index 0aff9ac5e2f..d6aa9d8df8a 100644 --- a/pkg/core/region.go +++ b/pkg/core/region.go @@ -1371,23 +1371,23 @@ func (r *RegionsInfo) GetStoreRegions(storeID uint64) []*RegionInfo { // SubTreeRegionType is the type of sub tree region. type SubTreeRegionType int -var ( +const ( // AllInSubTree is all sub trees. - AllInSubTree SubTreeRegionType = 0 + AllInSubTree SubTreeRegionType = iota // LeaderInSubTree is the leader sub tree. - LeaderInSubTree SubTreeRegionType = 1 + LeaderInSubTree // FollowerInSubTree is the follower sub tree. - FollowerInSubTree SubTreeRegionType = 2 + FollowerInSubTree // LearnerInSubTree is the learner sub tree. - LearnerInSubTree SubTreeRegionType = 3 + LearnerInSubTree // WitnessInSubTree is the witness sub tree. - WitnessInSubTree SubTreeRegionType = 4 + WitnessInSubTree // PendingPeerInSubTree is the pending peer sub tree. - PendingPeerInSubTree SubTreeRegionType = 5 + PendingPeerInSubTree ) // GetStoreRegions gets all RegionInfo with a given storeID -func (r *RegionsInfo) GetStoreRegionsByTypeInSubTree(storeID uint64, typ SubTreeRegionType) []*RegionInfo { +func (r *RegionsInfo) GetStoreRegionsByTypeInSubTree(storeID uint64, typ SubTreeRegionType) ([]*RegionInfo, error) { r.st.RLock() var regions []*RegionInfo switch typ { @@ -1413,10 +1413,13 @@ func (r *RegionsInfo) GetStoreRegionsByTypeInSubTree(storeID uint64, typ SubTree } case AllInSubTree: r.st.RUnlock() - return r.GetStoreRegions(storeID) + return r.GetStoreRegions(storeID), nil + default: + return nil, errors.Errorf("unknown sub tree region type %v", typ) } + r.st.RUnlock() - return regions + return regions, nil } // GetStoreLeaderRegionSize get total size of store's leader regions diff --git a/server/api/region.go b/server/api/region.go index dc466bd33c7..2f1d07988c4 100644 --- a/server/api/region.go +++ b/server/api/region.go @@ -223,7 +223,11 @@ func (h *regionsHandler) GetStoreRegions(w http.ResponseWriter, r *http.Request) if err != nil { typID = int(core.AllInSubTree) } - regions := rc.GetStoreRegionsByTypeInSubTree(uint64(id), core.SubTreeRegionType(typID)) + regions, err := rc.GetStoreRegionsByTypeInSubTree(uint64(id), core.SubTreeRegionType(typID)) + 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/tools/pd-ctl/tests/region/region_test.go b/tools/pd-ctl/tests/region/region_test.go index 109432ca9fa..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,11 +123,11 @@ 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", "--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=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}},