api: support query regions by read/write flow.

pdctl/command/region_command.go

@@ -25,9 +25,11 @@ import (
 )
 
 var (
-	regionsPrefix   = "pd/api/v1/regions"
-	regionIDPrefix  = "pd/api/v1/region/id"
-	regionKeyPrefix = "pd/api/v1/region/key"
+	regionsPrefix          = "pd/api/v1/regions"
+	regionsWriteflowPrefix = "pd/api/v1/regions/writeflow"
+	regionsReadflowPrefix  = "pd/api/v1/regions/readflow"
+	regionIDPrefix         = "pd/api/v1/region/id"
+	regionKeyPrefix        = "pd/api/v1/region/key"
 )
 
 type regionInfo struct {
@@ -43,6 +45,21 @@ func NewRegionCommand() *cobra.Command {
 		Run:   showRegionCommandFunc,
 	}
 	r.AddCommand(NewRegionWithKeyCommand())
+
+	topRead := &cobra.Command{
+		Use:   "topread <limit>",
+		Short: "show regions with top read flow",
+		Run:   showRegionTopReadCommandFunc,
+	}
+	r.AddCommand(topRead)
+
+	topWrite := &cobra.Command{
+		Use:   "topwrite <limit>",
+		Short: "show regions with top write flow",
+		Run:   showRegionTopWriteCommandFunc,
+	}
+	r.AddCommand(topWrite)
+
 	return r
 }
 
@@ -64,6 +81,40 @@ func showRegionCommandFunc(cmd *cobra.Command, args []string) {
 	fmt.Println(r)
 }
 
+func showRegionTopWriteCommandFunc(cmd *cobra.Command, args []string) {
+	prefix := regionsWriteflowPrefix
+	if len(args) == 1 {
+		if _, err := strconv.Atoi(args[0]); err != nil {
+			fmt.Println("limit should be a number")
+			return
+		}
+		prefix += "?limit=" + args[0]
+	}
+	r, err := doRequest(cmd, prefix, http.MethodGet)
+	if err != nil {
+		fmt.Printf("Failed to get regions: %s\n", err)
+		return
+	}
+	fmt.Println(r)
+}
+
+func showRegionTopReadCommandFunc(cmd *cobra.Command, args []string) {
+	prefix := regionsReadflowPrefix
+	if len(args) == 1 {
+		if _, err := strconv.Atoi(args[0]); err != nil {
+			fmt.Println("limit should be a number")
+			return
+		}
+		prefix += "?limit=" + args[0]
+	}
+	r, err := doRequest(cmd, prefix, http.MethodGet)
+	if err != nil {
+		fmt.Printf("Failed to get regions: %s\n", err)
+		return
+	}
+	fmt.Println(r)
+}
+
 // NewRegionWithKeyCommand return a region with key subcommand of regionCmd
 func NewRegionWithKeyCommand() *cobra.Command {
 	r := &cobra.Command{

server/api/region.go

@@ -14,6 +14,7 @@
 package api
 
 import (
+	"container/heap"
 	"fmt"
 	"net/http"
 	"strconv"
@@ -122,14 +123,14 @@ func newRegionsHandler(svr *server.Server, rd *render.Render) *regionsHandler {
 	}
 }
 
-func (h *regionsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+func (h *regionsHandler) GetAll(w http.ResponseWriter, r *http.Request) {
 	cluster := h.svr.GetRaftCluster()
 	if cluster == nil {
 		h.rd.JSON(w, http.StatusInternalServerError, server.ErrNotBootstrapped.Error())
 		return
 	}
 
-	regions := cluster.GetRegions()
+	regions := cluster.GetMetaRegions()
 	regionInfos := make([]*regionInfo, len(regions))
 	for i, r := range regions {
 		regionInfos[i] = newRegionInfo(&core.RegionInfo{Region: r})
@@ -140,3 +141,105 @@ func (h *regionsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 	h.rd.JSON(w, http.StatusOK, regionsInfo)
 }
+
+const (
+	defaultRegionLimit = 16
+	maxRegionLimit     = 10240
+)
+
+func (h *regionsHandler) GetTopWriteFlow(w http.ResponseWriter, r *http.Request) {
+	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool { return a.WrittenBytes < b.WrittenBytes })
+}
+
+func (h *regionsHandler) GetTopReadFlow(w http.ResponseWriter, r *http.Request) {
+	h.GetTopNRegions(w, r, func(a, b *core.RegionInfo) bool { return a.ReadBytes < b.ReadBytes })
+}
+
+func (h *regionsHandler) GetTopNRegions(w http.ResponseWriter, r *http.Request, less func(a, b *core.RegionInfo) bool) {
+	cluster := h.svr.GetRaftCluster()
+	if cluster == nil {
+		h.rd.JSON(w, http.StatusInternalServerError, server.ErrNotBootstrapped.Error())
+		return
+	}
+	limit := defaultRegionLimit
+	if limitStr := r.URL.Query().Get("limit"); limitStr != "" {
+		var err error
+		limit, err = strconv.Atoi(limitStr)
+		if err != nil {
+			h.rd.JSON(w, http.StatusBadRequest, err.Error())
+			return
+		}
+	}
+	if limit > maxRegionLimit {
+		limit = maxRegionLimit
+	}
+	regions := topNRegions(cluster.GetRegions(), less, limit)
+	regionInfos := make([]*regionInfo, len(regions))
+	for i, r := range regions {
+		regionInfos[i] = newRegionInfo(r)
+	}
+	res := &regionsInfo{
+		Count:   len(regions),
+		Regions: regionInfos,
+	}
+	h.rd.JSON(w, http.StatusOK, res)
+}
+
+// RegionHeap implements heap.Interface, used for selecting top n regions.
+type RegionHeap struct {
+	regions []*core.RegionInfo
+	less    func(a, b *core.RegionInfo) bool
+}
+
+func (h *RegionHeap) Len() int           { return len(h.regions) }
+func (h *RegionHeap) Less(i, j int) bool { return h.less(h.regions[i], h.regions[j]) }
+func (h *RegionHeap) Swap(i, j int)      { h.regions[i], h.regions[j] = h.regions[j], h.regions[i] }
+
+// Push pushes an element x onto the heap.
+func (h *RegionHeap) Push(x interface{}) {
+	h.regions = append(h.regions, x.(*core.RegionInfo))
+}
+
+// Pop removes the minimum element (according to Less) from the heap and returns
+// it.
+func (h *RegionHeap) Pop() interface{} {
+	pos := len(h.regions) - 1
+	x := h.regions[pos]
+	h.regions = h.regions[:pos]
+	return x
+}
+
+// Min returns the minimum region from the heap.
+func (h *RegionHeap) Min() *core.RegionInfo {
+	if h.Len() == 0 {
+		return nil
+	}
+	return h.regions[0]
+}
+
+func topNRegions(regions []*core.RegionInfo, less func(a, b *core.RegionInfo) bool, n int) []*core.RegionInfo {
+	if n <= 0 {
+		return nil
+	}
+
+	hp := &RegionHeap{
+		regions: make([]*core.RegionInfo, 0, n),
+		less:    less,
+	}
+	for _, r := range regions {
+		if hp.Len() < n {
+			heap.Push(hp, r)
+			continue
+		}
+		if less(hp.Min(), r) {
+			heap.Pop(hp)
+			heap.Push(hp, r)
+		}
+	}
+
+	res := make([]*core.RegionInfo, hp.Len())
+	for i := hp.Len() - 1; i >= 0; i-- {
+		res[i] = heap.Pop(hp).(*core.RegionInfo)
+	}
+	return res
+}

server/api/region_test.go

@@ -15,6 +15,7 @@ package api
 
 import (
 	"fmt"
+	"math/rand"
 
 	. "github.com/pingcap/check"
 	"github.com/pingcap/kvproto/pkg/metapb"
@@ -76,3 +77,50 @@ func (s *testRegionSuite) TestRegion(c *C) {
 	c.Assert(err, IsNil)
 	c.Assert(r2, DeepEquals, newRegionInfo(r))
 }
+
+func (s *testRegionSuite) TestTopFlow(c *C) {
+	r1 := newTestRegionInfo(1, 1, []byte("a"), []byte("b"))
+	r1.WrittenBytes, r1.ReadBytes = 1000, 1000
+	mustRegionHeartbeat(c, s.svr, r1)
+	r2 := newTestRegionInfo(2, 1, []byte("b"), []byte("c"))
+	r2.WrittenBytes, r2.ReadBytes = 2000, 0
+	mustRegionHeartbeat(c, s.svr, r2)
+	r3 := newTestRegionInfo(3, 1, []byte("c"), []byte("d"))
+	r3.WrittenBytes, r3.ReadBytes = 500, 800
+	mustRegionHeartbeat(c, s.svr, r3)
+	s.checkTopFlow(c, fmt.Sprintf("%s/regions/writeflow", s.urlPrefix), []uint64{2, 1, 3})
+	s.checkTopFlow(c, fmt.Sprintf("%s/regions/readflow", s.urlPrefix), []uint64{1, 3, 2})
+	s.checkTopFlow(c, fmt.Sprintf("%s/regions/writeflow?limit=2", s.urlPrefix), []uint64{2, 1})
+}
+
+func (s *testRegionSuite) checkTopFlow(c *C, url string, regionIDs []uint64) {
+	regions := &regionsInfo{}
+	err := readJSONWithURL(url, regions)
+	c.Assert(err, IsNil)
+	c.Assert(regions.Count, Equals, len(regionIDs))
+	for i, r := range regions.Regions {
+		c.Assert(r.ID, Equals, regionIDs[i])
+	}
+}
+
+func (s *testRegionSuite) TestTopN(c *C) {
+	writtenBytes := []uint64{10, 10, 9, 5, 3, 2, 2, 1, 0, 0}
+	for n := 0; n <= len(writtenBytes)+1; n++ {
+		regions := make([]*core.RegionInfo, 0, len(writtenBytes))
+		for _, i := range rand.Perm(len(writtenBytes)) {
+			id := uint64(i + 1)
+			region := newTestRegionInfo(id, id, nil, nil)
+			region.WrittenBytes = uint64(writtenBytes[i])
+			regions = append(regions, region)
+		}
+		topN := topNRegions(regions, func(a, b *core.RegionInfo) bool { return a.WrittenBytes < b.WrittenBytes }, n)
+		if n > len(writtenBytes) {
+			c.Assert(len(topN), Equals, len(writtenBytes))
+		} else {
+			c.Assert(len(topN), Equals, n)
+		}
+		for i := range topN {
+			c.Assert(topN[i].WrittenBytes, Equals, writtenBytes[i])
+		}
+	}
+}

server/api/router.go

@@ -76,7 +76,11 @@ func createRouter(prefix string, svr *server.Server) *mux.Router {
 	router.HandleFunc("/api/v1/region/id/{id}", regionHandler.GetRegionByID).Methods("GET")
 	router.HandleFunc("/api/v1/region/key/{key}", regionHandler.GetRegionByKey).Methods("GET")
 
-	router.Handle("/api/v1/regions", newRegionsHandler(svr, rd)).Methods("GET")
+	regionsHandler := newRegionsHandler(svr, rd)
+	router.HandleFunc("/api/v1/regions", regionsHandler.GetAll).Methods("GET")
+	router.HandleFunc("/api/v1/regions/writeflow", regionsHandler.GetTopWriteFlow).Methods("GET")
+	router.HandleFunc("/api/v1/regions/readflow", regionsHandler.GetTopReadFlow).Methods("GET")
+
 	router.Handle("/api/v1/version", newVersionHandler(rd)).Methods("GET")
 	router.Handle("/api/v1/status", newStatusHandler(rd)).Methods("GET")
 

server/cluster.go

@@ -227,11 +227,16 @@ func (c *RaftCluster) GetRegionInfoByID(regionID uint64) *core.RegionInfo {
 	return c.cachedCluster.GetRegion(regionID)
 }
 
-// GetRegions gets regions from cluster.
-func (c *RaftCluster) GetRegions() []*metapb.Region {
+// GetMetaRegions gets regions from cluster.
+func (c *RaftCluster) GetMetaRegions() []*metapb.Region {
 	return c.cachedCluster.getMetaRegions()
 }
 
+// GetRegions returns all regions info in detail.
+func (c *RaftCluster) GetRegions() []*core.RegionInfo {
+	return c.cachedCluster.getRegions()
+}
+
 // GetRegionStats returns region statistics from cluster.
 func (c *RaftCluster) GetRegionStats(startKey, endKey []byte) *core.RegionStats {
 	return c.cachedCluster.getRegionStats(startKey, endKey)