api: use easyjson and context in regions interface (#6838)

go.mod

@@ -51,6 +51,7 @@ require (
 )
 
 require (
+	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/samber/lo v1.37.0 // indirect
 	gorm.io/datatypes v1.1.0 // indirect
 )
@@ -115,7 +116,7 @@ require (
 	github.com/json-iterator/go v1.1.12 // indirect
 	github.com/konsorten/go-windows-terminal-sequences v1.0.3 // indirect
 	github.com/leodido/go-urn v1.2.0 // indirect
-	github.com/mailru/easyjson v0.7.6 // indirect
+	github.com/mailru/easyjson v0.7.6
 	github.com/mattn/go-colorable v0.1.8 // indirect
 	github.com/mattn/go-isatty v0.0.12 // indirect
 	github.com/mattn/go-runewidth v0.0.8 // indirect

go.sum

@@ -216,8 +216,9 @@ github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMyw
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
-github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
+github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
+github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/pprof v0.0.0-20211122183932-1daafda22083 h1:c8EUapQFi+kjzedr4c6WqbwMdmB95+oDBWZ5XFHFYxY=
 github.com/google/pprof v0.0.0-20211122183932-1daafda22083/go.mod h1:KgnwoLYCZ8IQu3XUZ8Nc/bM9CCZFOyjUNOSygVozoDg=
@@ -715,7 +716,6 @@ golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
-golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 h1:go1bK/D/BFZV2I8cIQd1NKEZ+0owSTG1fDTci4IqFcE=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM=
 google.golang.org/appengine v1.4.0 h1:/wp5JvzpHIxhs/dumFmF7BXTf3Z+dd4uXta4kVyO508=

server/api/region.go

@@ -16,6 +16,7 @@
 
 import (
 	"container/heap"
+	"context"
 	"encoding/hex"
 	"fmt"
 	"net/http"
@@ -25,6 +26,7 @@
 	"strings"
 
 	"github.com/gorilla/mux"
+	jwriter "github.com/mailru/easyjson/jwriter"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
@@ -53,13 +55,34 @@
 	IsLearner bool `json:"is_learner,omitempty"`
 }
 
+func (m *MetaPeer) setDefaultIfNil() {
+	if m.Peer == nil {
+		m.Peer = &metapb.Peer{
+			Id:        m.GetId(),
+			StoreId:   m.GetStoreId(),
+			Role:      m.GetRole(),
+			IsWitness: m.GetIsWitness(),
+		}
+	}
+}
+
 // PDPeerStats is api compatible with *pdpb.PeerStats.
 // NOTE: This type is exported by HTTP API. Please pay more attention when modifying it.
 type PDPeerStats struct {
 	*pdpb.PeerStats
 	Peer MetaPeer `json:"peer"`
 }
 
+func (s *PDPeerStats) setDefaultIfNil() {
+	if s.PeerStats == nil {
+		s.PeerStats = &pdpb.PeerStats{
+			Peer:        s.GetPeer(),
+			DownSeconds: s.GetDownSeconds(),
+		}
+	}
+	s.Peer.setDefaultIfNil()
+}
+
 func fromPeer(peer *metapb.Peer) MetaPeer {
 	if peer == nil {
 		return MetaPeer{}
@@ -102,6 +125,7 @@
 
 // RegionInfo records detail region info for api usage.
 // NOTE: This type is exported by HTTP API. Please pay more attention when modifying it.
+// easyjson:json
 type RegionInfo struct {
 	ID          uint64              `json:"id"`
 	StartKey    string              `json:"start_key"`
@@ -168,9 +192,9 @@
 	s.ApproximateSize = r.GetApproximateSize()
 	s.ApproximateKeys = r.GetApproximateKeys()
 	s.ReplicationStatus = fromPBReplicationStatus(r.GetReplicationStatus())
+	s.Buckets = nil
 
 	keys := r.GetBuckets().GetKeys()
-
 	if len(keys) > 0 {
 		s.Buckets = make([]string, len(keys))
 		for i, key := range keys {
@@ -312,15 +336,48 @@
 	}
 }
 
-func convertToAPIRegions(regions []*core.RegionInfo) *RegionsInfo {
-	regionInfos := make([]RegionInfo, len(regions))
+// marshalRegionsInfoJSON marshals regions to bytes in `RegionsInfo`'s JSON format.
+// It is used to reduce the cost of JSON serialization.
+func marshalRegionsInfoJSON(ctx context.Context, regions []*core.RegionInfo) ([]byte, error) {
+	out := &jwriter.Writer{}
+	out.RawByte('{')
+
+	out.RawString("\"count\":")
+	out.Int(len(regions))
+
+	out.RawString(",\"regions\":")
+	out.RawByte('[')
+	region := &RegionInfo{}
 	for i, r := range regions {
-		InitRegion(r, &regionInfos[i])
-	}
-	return &RegionsInfo{
-		Count:   len(regions),
-		Regions: regionInfos,
+		select {
+		case <-ctx.Done():
+			// Return early, avoid the unnecessary computation.
+			// See more details in https://github.com/tikv/pd/issues/6835
+			return nil, ctx.Err()
+		default:
+		}
+		if i > 0 {
+			out.RawByte(',')
+		}
+		InitRegion(r, region)
+		// EasyJSON will not check anonymous struct pointer field and will panic if the field is nil.
+		// So we need to set the field to default value explicitly when the anonymous struct pointer is nil.
+		region.Leader.setDefaultIfNil()
+		for i := range region.Peers {
+			region.Peers[i].setDefaultIfNil()
+		}
+		for i := range region.PendingPeers {
+			region.PendingPeers[i].setDefaultIfNil()
+		}
+		for i := range region.DownPeers {
+			region.DownPeers[i].setDefaultIfNil()
+		}
+		region.MarshalEasyJSON(out)
 	}
+	out.RawByte(']')
+
+	out.RawByte('}')
+	return out.Buffer.BuildBytes(), out.Error
 }
 
 // @Tags     region
@@ -331,8 +388,7 @@
 func (h *regionsHandler) GetRegions(w http.ResponseWriter, r *http.Request) {
 	rc := getCluster(r)
 	regions := rc.GetRegions()
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.returnWithRegions(w, r, regions)
 }
 
 // @Tags     region
@@ -362,8 +418,7 @@
 		limit = maxRegionLimit
 	}
 	regions := rc.ScanRegions([]byte(startKey), []byte(endKey), limit)
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.returnWithRegions(w, r, regions)
 }
 
 // @Tags     region
@@ -394,8 +449,7 @@
 		return
 	}
 	regions := rc.GetStoreRegions(uint64(id))
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.returnWithRegions(w, r, regions)
 }
 
 // @Tags     region
@@ -405,14 +459,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/miss-peer [get]
 func (h *regionsHandler) GetMissPeerRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.MissPeer)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.MissPeer, r)
 }
 
 // @Tags     region
@@ -422,14 +469,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/extra-peer [get]
 func (h *regionsHandler) GetExtraPeerRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.ExtraPeer)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.ExtraPeer, r)
 }
 
 // @Tags     region
@@ -439,14 +479,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/pending-peer [get]
 func (h *regionsHandler) GetPendingPeerRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.PendingPeer)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.PendingPeer, r)
 }
 
 // @Tags     region
@@ -456,14 +489,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/down-peer [get]
 func (h *regionsHandler) GetDownPeerRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.DownPeer)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.DownPeer, r)
 }
 
 // @Tags     region
@@ -473,14 +499,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/learner-peer [get]
 func (h *regionsHandler) GetLearnerPeerRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.LearnerPeer)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.LearnerPeer, r)
 }
 
 // @Tags     region
@@ -490,14 +509,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/offline-peer [get]
 func (h *regionsHandler) GetOfflinePeerRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetOfflinePeer(statistics.OfflinePeer)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.OfflinePeer, r)
 }
 
 // @Tags     region
@@ -507,14 +519,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/oversized-region [get]
 func (h *regionsHandler) GetOverSizedRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.OversizedRegion)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.OversizedRegion, r)
 }
 
 // @Tags     region
@@ -524,14 +529,7 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/undersized-region [get]
 func (h *regionsHandler) GetUndersizedRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.UndersizedRegion)
-	if err != nil {
-		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
-		return
-	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.getRegionsByType(w, statistics.UndersizedRegion, r)
 }
 
 // @Tags     region
@@ -541,14 +539,24 @@
 // @Failure  500  {string}  string  "PD server failed to proceed the request."
 // @Router   /regions/check/empty-region [get]
 func (h *regionsHandler) GetEmptyRegions(w http.ResponseWriter, r *http.Request) {
-	handler := h.svr.GetHandler()
-	regions, err := handler.GetRegionsByType(statistics.EmptyRegion)
+	h.getRegionsByType(w, statistics.EmptyRegion, r)
+}
+
+func (h *regionsHandler) getRegionsByType(w http.ResponseWriter, t statistics.RegionStatisticType, r *http.Request) {
+	regions, err := h.svr.GetHandler().GetRegionsByType(t)
 	if err != nil {
 		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
 		return
 	}
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	h.returnWithRegions(w, r, regions)
+}
+
+func (h *regionsHandler) returnWithRegions(w http.ResponseWriter, r *http.Request, regions []*core.RegionInfo) {
+	b, err := marshalRegionsInfoJSON(r.Context(), regions)
+	if err != nil {
+		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
+	}
+	h.rd.Data(w, http.StatusOK, b)
 }
 
 type histItem struct {
@@ -688,8 +696,12 @@
 	}
 
 	left, right := rc.GetAdjacentRegions(region)
-	regionsInfo := convertToAPIRegions([]*core.RegionInfo{left, right})
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	b, err := marshalRegionsInfoJSON(r.Context(), []*core.RegionInfo{left, right})
+	if err != nil {
+		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
+		return
+	}
+	h.rd.Data(w, http.StatusOK, b)
 }
 
 const (
@@ -910,8 +922,12 @@
 		limit = maxRegionLimit
 	}
 	regions := TopNRegions(rc.GetRegions(), less, limit)
-	regionsInfo := convertToAPIRegions(regions)
-	h.rd.JSON(w, http.StatusOK, regionsInfo)
+	b, err := marshalRegionsInfoJSON(r.Context(), regions)
+	if err != nil {
+		h.rd.JSON(w, http.StatusInternalServerError, err.Error())
+		return
+	}
+	h.rd.Data(w, http.StatusOK, b)
 }
 
 // @Tags     region