pd-ctl: list all regions in a specified store (#1231)

pdctl/command/region_command.go

@@ -26,6 +26,7 @@ import (
 
 var (
 	regionsPrefix          = "pd/api/v1/regions"
+	regionsStorePrefix     = "pd/api/v1/regions/store"
 	regionsCheckPrefix     = "pd/api/v1/regions/check"
 	regionsWriteflowPrefix = "pd/api/v1/regions/writeflow"
 	regionsReadflowPrefix  = "pd/api/v1/regions/readflow"
@@ -34,7 +35,7 @@ var (
 	regionKeyPrefix        = "pd/api/v1/region/key"
 )
 
-// NewRegionCommand return a region subcommand of rootCmd
+// NewRegionCommand returns a region subcommand of rootCmd
 func NewRegionCommand() *cobra.Command {
 	r := &cobra.Command{
 		Use:   `region <region_id> [-jq="<query string>"]`,
@@ -44,6 +45,7 @@ func NewRegionCommand() *cobra.Command {
 	r.AddCommand(NewRegionWithKeyCommand())
 	r.AddCommand(NewRegionWithCheckCommand())
 	r.AddCommand(NewRegionWithSiblingCommand())
+	r.AddCommand(NewRegionWithStoreCommand())
 
 	topRead := &cobra.Command{
 		Use:   "topread <limit>",
@@ -119,7 +121,7 @@ func showRegionTopReadCommandFunc(cmd *cobra.Command, args []string) {
 	fmt.Println(r)
 }
 
-// NewRegionWithKeyCommand return a region with key subcommand of regionCmd
+// NewRegionWithKeyCommand returns a region with key subcommand of regionCmd
 func NewRegionWithKeyCommand() *cobra.Command {
 	r := &cobra.Command{
 		Use:   "key [--format=raw|pb|proto|protobuf] <key>",
@@ -188,7 +190,7 @@ func decodeProtobufText(text string) (string, error) {
 	return string(buf), nil
 }
 
-// NewRegionWithCheckCommand return a region with check subcommand of regionCmd
+// NewRegionWithCheckCommand returns a region with check subcommand of regionCmd
 func NewRegionWithCheckCommand() *cobra.Command {
 	r := &cobra.Command{
 		Use:   "check [miss-peer|extra-peer|down-peer|pending-peer|incorrect-ns]",
@@ -213,7 +215,7 @@ func showRegionWithCheckCommandFunc(cmd *cobra.Command, args []string) {
 	fmt.Println(r)
 }
 
-// NewRegionWithSiblingCommand return a region with check subcommand of regionCmd
+// NewRegionWithSiblingCommand returns a region with sibling subcommand of regionCmd
 func NewRegionWithSiblingCommand() *cobra.Command {
 	r := &cobra.Command{
 		Use:   "sibling <region_id>",
@@ -238,6 +240,31 @@ func showRegionWithSiblingCommandFunc(cmd *cobra.Command, args []string) {
 	fmt.Println(r)
 }
 
+// NewRegionWithStoreCommand returns regions with store subcommand of regionCmd
+func NewRegionWithStoreCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:   "store <store_id>",
+		Short: "show the regions of a specific store",
+		Run:   showRegionWithStoreCommandFunc,
+	}
+	return r
+}
+
+func showRegionWithStoreCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) != 1 {
+		fmt.Println(cmd.UsageString())
+		return
+	}
+	storeID := args[0]
+	prefix := regionsStorePrefix + "/" + storeID
+	r, err := doRequest(cmd, prefix, http.MethodGet)
+	if err != nil {
+		fmt.Printf("Failed to get regions with the given storeID: %s\n", err)
+		return
+	}
+	fmt.Println(r)
+}
+
 func printWithJQFilter(data, filter string) {
 	cmd := exec.Command("jq", "-c", filter)
 	stdin, err := cmd.StdinPipe()

server/api/api.raml

@@ -861,6 +861,20 @@ types:
           description: The region does not exist.
         500:
           description: PD server failed to proceed the request.
+  /store/{id}:
+    uriParameters:
+      id: integer
+    get:
+      description: List all regions of a specific store.
+      responses:
+        200:
+          body:
+            application/json:
+              type: Regions
+        400:
+          description: The input is invalid.
+        500:
+          description: PD server failed to proceed the request.
 
 /schedulers:
   description: Running schedulers.

server/api/region.go

@@ -143,6 +143,31 @@ func (h *regionsHandler) GetAll(w http.ResponseWriter, r *http.Request) {
 	h.rd.JSON(w, http.StatusOK, regionsInfo)
 }
 
+func (h *regionsHandler) GetStoreRegions(w http.ResponseWriter, r *http.Request) {
+	cluster := h.svr.GetRaftCluster()
+	if cluster == nil {
+		h.rd.JSON(w, http.StatusInternalServerError, server.ErrNotBootstrapped.Error())
+		return
+	}
+
+	vars := mux.Vars(r)
+	id, err := strconv.Atoi(vars["id"])
+	if err != nil {
+		h.rd.JSON(w, http.StatusBadRequest, err.Error())
+		return
+	}
+	regions := cluster.GetStoreRegions(uint64(id))
+	regionInfos := make([]*regionInfo, len(regions))
+	for i, r := range regions {
+		regionInfos[i] = newRegionInfo(r)
+	}
+	regionsInfo := &regionsInfo{
+		Count:   len(regions),
+		Regions: regionInfos,
+	}
+	h.rd.JSON(w, http.StatusOK, regionsInfo)
+}
+
 func (h *regionsHandler) GetMissPeerRegions(w http.ResponseWriter, r *http.Request) {
 	handler := h.svr.GetHandler()
 	res, err := handler.GetMissPeerRegions()

server/api/region_test.go

@@ -16,6 +16,7 @@ package api
 import (
 	"fmt"
 	"math/rand"
+	"sort"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
@@ -81,6 +82,43 @@ func (s *testRegionSuite) TestRegion(c *C) {
 	c.Assert(r2, DeepEquals, newRegionInfo(r))
 }
 
+func (s *testRegionSuite) TestStoreRegions(c *C) {
+	r1 := newTestRegionInfo(2, 1, []byte("a"), []byte("b"))
+	r2 := newTestRegionInfo(3, 1, []byte("b"), []byte("c"))
+	r3 := newTestRegionInfo(4, 2, []byte("c"), []byte("d"))
+	mustRegionHeartbeat(c, s.svr, r1)
+	mustRegionHeartbeat(c, s.svr, r2)
+	mustRegionHeartbeat(c, s.svr, r3)
+
+	regionIDs := []uint64{2, 3}
+	url := fmt.Sprintf("%s/regions/store/%d", s.urlPrefix, 1)
+	r4 := &regionsInfo{}
+	err := readJSONWithURL(url, r4)
+	c.Assert(err, IsNil)
+	c.Assert(r4.Count, Equals, len(regionIDs))
+	sort.Slice(r4.Regions, func(i, j int) bool { return r4.Regions[i].ID < r4.Regions[j].ID })
+	for i, r := range r4.Regions {
+		c.Assert(r.ID, Equals, regionIDs[i])
+	}
+
+	regionIDs = []uint64{4}
+	url = fmt.Sprintf("%s/regions/store/%d", s.urlPrefix, 2)
+	r5 := &regionsInfo{}
+	err = readJSONWithURL(url, r5)
+	c.Assert(err, IsNil)
+	c.Assert(r5.Count, Equals, len(regionIDs))
+	for i, r := range r5.Regions {
+		c.Assert(r.ID, Equals, regionIDs[i])
+	}
+
+	regionIDs = []uint64{}
+	url = fmt.Sprintf("%s/regions/store/%d", s.urlPrefix, 3)
+	r6 := &regionsInfo{}
+	err = readJSONWithURL(url, r6)
+	c.Assert(err, IsNil)
+	c.Assert(r6.Count, Equals, len(regionIDs))
+}
+
 func (s *testRegionSuite) TestTopFlow(c *C) {
 	r1 := newTestRegionInfo(1, 1, []byte("a"), []byte("b"), core.SetWrittenBytes(1000), core.SetReadBytes(1000))
 	mustRegionHeartbeat(c, s.svr, r1)

server/api/router.go

@@ -84,6 +84,7 @@ func createRouter(prefix string, svr *server.Server) *mux.Router {
 
 	regionsHandler := newRegionsHandler(svr, rd)
 	router.HandleFunc("/api/v1/regions", regionsHandler.GetAll).Methods("GET")
+	router.HandleFunc("/api/v1/regions/store/{id}", regionsHandler.GetStoreRegions).Methods("GET")
 	router.HandleFunc("/api/v1/regions/writeflow", regionsHandler.GetTopWriteFlow).Methods("GET")
 	router.HandleFunc("/api/v1/regions/readflow", regionsHandler.GetTopReadFlow).Methods("GET")
 	router.HandleFunc("/api/v1/regions/check/miss-peer", regionsHandler.GetMissPeerRegions).Methods("GET")

server/cluster.go

@@ -252,6 +252,11 @@ func (c *RaftCluster) GetRegions() []*core.RegionInfo {
 	return c.cachedCluster.getRegions()
 }
 
+// GetStoreRegions returns all regions info with a given storeID.
+func (c *RaftCluster) GetStoreRegions(storeID uint64) []*core.RegionInfo {
+	return c.cachedCluster.getStoreRegions(storeID)
+}
+
 // GetRegionStats returns region statistics from cluster.
 func (c *RaftCluster) GetRegionStats(startKey, endKey []byte) *core.RegionStats {
 	return c.cachedCluster.getRegionStats(startKey, endKey)

server/cluster_info.go

@@ -312,6 +312,12 @@ func (c *clusterInfo) getRegions() []*core.RegionInfo {
 	return c.core.Regions.GetRegions()
 }
 
+func (c *clusterInfo) getStoreRegions(storeID uint64) []*core.RegionInfo {
+	c.RLock()
+	defer c.RUnlock()
+	return c.core.Regions.GetStoreRegions(storeID)
+}
+
 func (c *clusterInfo) getMetaRegions() []*metapb.Region {
 	c.RLock()
 	defer c.RUnlock()

server/core/region.go

@@ -630,6 +630,22 @@ func (r *RegionsInfo) GetRegions() []*RegionInfo {
 	return regions
 }
 
+// GetStoreRegions gets all RegionInfo with a given storeID
+func (r *RegionsInfo) GetStoreRegions(storeID uint64) []*RegionInfo {
+	regions := make([]*RegionInfo, 0, r.GetStoreLeaderCount(storeID)+r.GetStoreFollowerCount(storeID))
+	if leaders, ok := r.leaders[storeID]; ok {
+		for _, region := range leaders.m {
+			regions = append(regions, region.RegionInfo)
+		}
+	}
+	if followers, ok := r.followers[storeID]; ok {
+		for _, region := range followers.m {
+			regions = append(regions, region.RegionInfo)
+		}
+	}
+	return regions
+}
+
 // GetStoreLeaderRegionSize get total size of store's leader regions
 func (r *RegionsInfo) GetStoreLeaderRegionSize(storeID uint64) int64 {
 	return r.leaders[storeID].TotalSize()