address

Signed-off-by: nolouch <[email protected]>
nolouch · Apr 29, 2024 · 95099af · 95099af
1 parent 949d15b
commit 95099af
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 16 deletions.
diff --git 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
@@ -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
@@ -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 <store_id> 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}},