update kvproto

Signed-off-by: Ryan Leung <[email protected]>

client/go.mod

@@ -40,5 +40,3 @@ require (
 	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
-
-replace github.com/pingcap/kvproto => github.com/rleungx/kvproto v0.0.0-20230824090226-d21284147287

client/go.sum

@@ -91,6 +91,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c h1:xpW9bvK+HuuTm
 github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 h1:C3N3itkduZXDZFh4N3vQ5HEtld3S+Y+StULhWVvumU0=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
+github.com/pingcap/kvproto v0.0.0-20230727073445-53e1f8730c30 h1:EvqKcDT7ceGLW0mXqM8Cp5Z8DfgQRnwj2YTnlCLj2QI=
+github.com/pingcap/kvproto v0.0.0-20230727073445-53e1f8730c30/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -117,8 +119,6 @@ github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsT
 github.com/prometheus/procfs v0.1.3/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4OA4YeYWdaU=
 github.com/prometheus/procfs v0.6.0 h1:mxy4L2jP6qMonqmq+aTtOx1ifVWUgG/TAmntgbh3xv4=
 github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
-github.com/rleungx/kvproto v0.0.0-20230824090226-d21284147287 h1:k5LiNAjzBL5TRaA/ayyUNfPbSkNCUMuYnaLAPb02euE=
-github.com/rleungx/kvproto v0.0.0-20230824090226-d21284147287/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
 github.com/sergi/go-diff v1.1.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
 github.com/sirupsen/logrus v1.2.0/go.mod h1:LxeOpSwHxABJmUn/MG1IvRgCAasNZTLOkJPxbbu5VWo=
 github.com/sirupsen/logrus v1.4.2/go.mod h1:tLMulIdttU9McNUspp0xgXVQah82FyeX6MwdIuYE2rE=

go.mod

@@ -33,7 +33,7 @@ require (
 	github.com/pingcap/errcode v0.3.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
-	github.com/pingcap/kvproto v0.0.0-20230911090708-d603cce32b96
+	github.com/pingcap/kvproto v0.0.0-20230914063002-3215a502e9a1
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21
 	github.com/pingcap/tidb-dashboard v0.0.0-20230911054332-22add1e00511
@@ -147,7 +147,7 @@ require (
 	github.com/pelletier/go-toml/v2 v2.0.1 // indirect
 	github.com/petermattis/goid v0.0.0-20211229010228-4d14c490ee36 // indirect
 	github.com/pingcap/tipb v0.0.0-20220718022156-3e2483c20a9e // indirect
-	github.com/pkg/errors v0.9.1 // indirect
+	github.com/pkg/errors v0.9.1
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/power-devops/perfstat v0.0.0-20221212215047-62379fc7944b // indirect
 	github.com/prometheus/client_model v0.2.0 // indirect
@@ -206,4 +206,4 @@ replace google.golang.org/grpc v1.54.0 => google.golang.org/grpc v1.26.0
 // After the PR to kvproto is merged, remember to comment this out and run `go mod tidy`.
 // replace github.com/pingcap/kvproto => github.com/$YourPrivateRepo $YourPrivateBranch
 
-replace github.com/pingcap/kvproto => github.com/rleungx/kvproto v0.0.0-20230824090226-d21284147287
+replace github.com/pingcap/kvproto => github.com/rleungx/kvproto v0.0.0-20230913074131-10a8f7b8a96c

go.sum

@@ -446,9 +446,6 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c h1:xpW9bvK+HuuTm
 github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ueLEUZgtx2cIogM0v4Zj5uvvzhuuiu7Pn8HzMPg=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00 h1:C3N3itkduZXDZFh4N3vQ5HEtld3S+Y+StULhWVvumU0=
 github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
-github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20230911090708-d603cce32b96 h1:Upb52Po0Ev1lPKQdUT4suRwQ5Z49A7gEmJ0trADKftM=
-github.com/pingcap/kvproto v0.0.0-20230911090708-d603cce32b96/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
@@ -492,8 +489,8 @@ github.com/prometheus/procfs v0.1.3/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4O
 github.com/prometheus/procfs v0.6.0 h1:mxy4L2jP6qMonqmq+aTtOx1ifVWUgG/TAmntgbh3xv4=
 github.com/prometheus/procfs v0.6.0/go.mod h1:cz+aTbrPOrUb4q7XlbU9ygM+/jj0fzG6c1xBZuNvfVA=
 github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40TwIPHuXU=
-github.com/rleungx/kvproto v0.0.0-20230824090226-d21284147287 h1:k5LiNAjzBL5TRaA/ayyUNfPbSkNCUMuYnaLAPb02euE=
-github.com/rleungx/kvproto v0.0.0-20230824090226-d21284147287/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
+github.com/rleungx/kvproto v0.0.0-20230913074131-10a8f7b8a96c h1:Ojpi9zrQ1gSysUsmb3vWjvqo91Mb/iggevAm1NBX530=
+github.com/rleungx/kvproto v0.0.0-20230913074131-10a8f7b8a96c/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
 github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
 github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4=

pkg/core/region.go

@@ -31,6 +31,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/kvproto/pkg/replication_modepb"
+	"github.com/pingcap/kvproto/pkg/schedulingpb"
 	"github.com/pingcap/log"
 	"github.com/tikv/pd/pkg/errs"
 	"github.com/tikv/pd/pkg/utils/logutil"
@@ -194,6 +195,54 @@ func RegionFromHeartbeat(heartbeat *pdpb.RegionHeartbeatRequest, opts ...RegionC
 	return region
 }
 
+// RegionFromForward constructs a Region from forwarded region heartbeat.
+func RegionFromForward(heartbeat *schedulingpb.RegionHeartbeatRequest, opts ...RegionCreateOption) *RegionInfo {
+	// Convert unit to MB.
+	// If region isn't empty and less than 1MB, use 1MB instead.
+	regionSize := heartbeat.GetApproximateSize() / units.MiB
+	// Due to https://github.com/tikv/tikv/pull/11170, if region size is not initialized,
+	// approximate size will be zero, and region size is zero not EmptyRegionApproximateSize
+	if heartbeat.GetApproximateSize() > 0 && regionSize < EmptyRegionApproximateSize {
+		regionSize = EmptyRegionApproximateSize
+	}
+
+	region := &RegionInfo{
+		term:            heartbeat.GetTerm(),
+		meta:            heartbeat.GetRegion(),
+		leader:          heartbeat.GetLeader(),
+		downPeers:       heartbeat.GetDownPeers(),
+		pendingPeers:    heartbeat.GetPendingPeers(),
+		cpuUsage:        heartbeat.GetCpuUsage(),
+		writtenBytes:    heartbeat.GetBytesWritten(),
+		writtenKeys:     heartbeat.GetKeysWritten(),
+		readBytes:       heartbeat.GetBytesRead(),
+		readKeys:        heartbeat.GetKeysRead(),
+		approximateSize: int64(regionSize),
+		approximateKeys: int64(heartbeat.GetApproximateKeys()),
+		interval:        heartbeat.GetInterval(),
+		queryStats:      heartbeat.GetQueryStats(),
+	}
+
+	for _, opt := range opts {
+		opt(region)
+	}
+
+	if region.writtenKeys >= ImpossibleFlowSize || region.writtenBytes >= ImpossibleFlowSize {
+		region.writtenKeys = 0
+		region.writtenBytes = 0
+	}
+	if region.readKeys >= ImpossibleFlowSize || region.readBytes >= ImpossibleFlowSize {
+		region.readKeys = 0
+		region.readBytes = 0
+	}
+
+	sort.Sort(peerStatsSlice(region.downPeers))
+	sort.Sort(peerSlice(region.pendingPeers))
+
+	classifyVoterAndLearner(region)
+	return region
+}
+
 // InheritBuckets inherits the buckets from the parent region if bucket enabled.
 func (r *RegionInfo) InheritBuckets(origin *RegionInfo) {
 	if origin != nil && r.buckets == nil {

pkg/mcs/scheduling/server/cluster.go

@@ -59,7 +59,6 @@ func NewCluster(ctx context.Context, persistConfig *config.PersistConfig, storag
 		hotStat:           statistics.NewHotStat(ctx),
 		labelLevelStats:   statistics.NewLabelStatistics(),
 		regionStats:       statistics.NewRegionStatistics(basicCluster, persistConfig, ruleManager),
-		labelLevelStats:   statistics.NewLabelStatistics(),
 		storage:           storage,
 		clusterID:         clusterID,
 		checkMembershipCh: checkMembershipCh,
@@ -269,7 +268,9 @@ func (c *Cluster) processRegionHeartbeat(region *core.RegionInfo) error {
 	if err != nil {
 		return err
 	}
-	region.Inherit(origin, c.GetStoreConfig().IsEnableRegionBucket())
+	if c.GetStoreConfig().IsEnableRegionBucket() {
+		region.InheritBuckets(origin)
+	}
 
 	c.hotStat.CheckWriteAsync(statistics.NewCheckExpiredItemTask(region))
 	c.hotStat.CheckReadAsync(statistics.NewCheckExpiredItemTask(region))
@@ -329,8 +330,3 @@ func (c *Cluster) processRegionHeartbeat(region *core.RegionInfo) error {
 func (c *Cluster) IsPrepared() bool {
 	return c.coordinator.GetPrepareChecker().IsPrepared()
 }
-
-// TODO: implement the following methods
-
-// AllocID allocates a new ID.
-func (c *Cluster) AllocID() (uint64, error) { return 0, nil }

pkg/mcs/scheduling/server/config/config.go

@@ -110,7 +110,6 @@ func (c *Config) Parse(flagSet *pflag.FlagSet) error {
 	configutil.AdjustCommandLineString(flagSet, &c.BackendEndpoints, "backend-endpoints")
 	configutil.AdjustCommandLineString(flagSet, &c.ListenAddr, "listen-addr")
 	configutil.AdjustCommandLineString(flagSet, &c.AdvertiseListenAddr, "advertise-listen-addr")
-
 	return c.adjust(meta)
 }
 

pkg/mcs/scheduling/server/grpc_service.go

@@ -21,13 +21,13 @@ import (
 	"sync/atomic"
 	"time"
 
-	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/pingcap/kvproto/pkg/schedulingpb"
 	"github.com/pingcap/log"
 	"github.com/pkg/errors"
 	bs "github.com/tikv/pd/pkg/basicserver"
 	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/mcs/registry"
+	"github.com/tikv/pd/pkg/schedule/hbstream"
 	"github.com/tikv/pd/pkg/utils/apiutil"
 	"github.com/tikv/pd/pkg/utils/logutil"
 	"go.uber.org/zap"
@@ -74,21 +74,21 @@ func NewService[T ConfigProvider](svr bs.Server) registry.RegistrableService {
 	}
 }
 
-// heartbeatServer wraps PD_RegionHeartbeatServer to ensure when any error
+// heartbeatServer wraps Scheduling_RegionHeartbeatServer to ensure when any error
 // occurs on Send() or Recv(), both endpoints will be closed.
 type heartbeatServer struct {
-	stream pdpb.PD_RegionHeartbeatServer
+	stream schedulingpb.Scheduling_RegionHeartbeatServer
 	closed int32
 }
 
-func (s *heartbeatServer) Send(m *pdpb.RegionHeartbeatResponse) error {
+func (s *heartbeatServer) Send(m hbstream.RegionHeartbeatResponse) error {
 	if atomic.LoadInt32(&s.closed) == 1 {
 		return io.EOF
 	}
 	done := make(chan error, 1)
 	go func() {
 		defer logutil.LogPanic()
-		done <- s.stream.Send(m)
+		done <- s.stream.Send(m.(*schedulingpb.RegionHeartbeatResponse))
 	}()
 	timer := time.NewTimer(5 * time.Second)
 	defer timer.Stop()
@@ -104,7 +104,7 @@ func (s *heartbeatServer) Send(m *pdpb.RegionHeartbeatResponse) error {
 	}
 }
 
-func (s *heartbeatServer) Recv() (*pdpb.RegionHeartbeatRequest, error) {
+func (s *heartbeatServer) Recv() (*schedulingpb.RegionHeartbeatRequest, error) {
 	if atomic.LoadInt32(&s.closed) == 1 {
 		return nil, io.EOF
 	}
@@ -141,10 +141,10 @@ func (s *Service) RegionHeartbeat(stream schedulingpb.Scheduling_RegionHeartbeat
 
 		c := s.GetCluster()
 		if c == nil {
-			resp := &pdpb.RegionHeartbeatResponse{Header: &pdpb.ResponseHeader{
+			resp := &schedulingpb.RegionHeartbeatResponse{Header: &schedulingpb.ResponseHeader{
 				ClusterId: s.clusterID,
-				Error: &pdpb.Error{
-					Type:    pdpb.ErrorType_NOT_BOOTSTRAPPED,
+				Error: &schedulingpb.Error{
+					Type:    schedulingpb.ErrorType_NOT_BOOTSTRAPPED,
 					Message: "scheduling server is not initialized yet",
 				},
 			}}
@@ -162,11 +162,11 @@ func (s *Service) RegionHeartbeat(stream schedulingpb.Scheduling_RegionHeartbeat
 			s.hbStreams.BindStream(storeID, server)
 			lastBind = time.Now()
 		}
-		region := core.RegionFromHeartbeat(request, core.SetFromHeartbeat(true))
+		region := core.RegionFromForward(request, core.SetFromHeartbeat(true))
 		err = c.HandleRegionHeartbeat(region)
 		if err != nil {
-			msg := err.Error()
-			s.hbStreams.SendErr(pdpb.ErrorType_UNKNOWN, msg, request.GetLeader())
+			// TODO: if we need to send the error back to API server.
+			log.Error("failed handle region heartbeat", zap.Error(err))
 			continue
 		}
 	}

pkg/mcs/scheduling/server/server.go

@@ -437,7 +437,7 @@ func (s *Server) startCluster(context.Context) error {
 	if err != nil {
 		return err
 	}
-	s.hbStreams = hbstream.NewHeartbeatStreams(s.Context(), s.clusterID, s.basicCluster)
+	s.hbStreams = hbstream.NewHeartbeatStreams(s.Context(), s.clusterID, utils.SchedulingServiceName, s.basicCluster)
 	s.cluster, err = NewCluster(s.Context(), s.persistConfig, s.storage, s.basicCluster, s.hbStreams, s.clusterID, s.checkMembershipCh)
 	if err != nil {
 		return err

pkg/mock/mockhbstream/mockhbstream.go

@@ -25,18 +25,18 @@ import (
 
 // HeartbeatStream is used to mock HeartbeatStream for test use.
 type HeartbeatStream struct {
-	ch chan *pdpb.RegionHeartbeatResponse
+	ch chan hbstream.RegionHeartbeatResponse
 }
 
 // NewHeartbeatStream creates a new HeartbeatStream.
 func NewHeartbeatStream() HeartbeatStream {
 	return HeartbeatStream{
-		ch: make(chan *pdpb.RegionHeartbeatResponse),
+		ch: make(chan hbstream.RegionHeartbeatResponse),
 	}
 }
 
 // Send mocks method.
-func (s HeartbeatStream) Send(m *pdpb.RegionHeartbeatResponse) error {
+func (s HeartbeatStream) Send(m hbstream.RegionHeartbeatResponse) error {
 	select {
 	case <-time.After(time.Second):
 		return errors.New("timeout")
@@ -52,7 +52,7 @@ func (s HeartbeatStream) SendMsg(region *core.RegionInfo, msg *pdpb.RegionHeartb
 func (s HeartbeatStream) BindStream(storeID uint64, stream hbstream.HeartbeatStream) {}
 
 // Recv mocks method.
-func (s HeartbeatStream) Recv() *pdpb.RegionHeartbeatResponse {
+func (s HeartbeatStream) Recv() hbstream.RegionHeartbeatResponse {
 	select {
 	case <-time.After(time.Millisecond * 10):
 		return nil

pkg/mock/mockhbstream/mockhbstream_test.go

@@ -22,12 +22,10 @@ import (
 	"github.com/pingcap/kvproto/pkg/metapb"
 	"github.com/pingcap/kvproto/pkg/pdpb"
 	"github.com/stretchr/testify/require"
-	"github.com/tikv/pd/pkg/core"
 	"github.com/tikv/pd/pkg/mock/mockcluster"
 	"github.com/tikv/pd/pkg/mock/mockconfig"
 	"github.com/tikv/pd/pkg/schedule/hbstream"
 	"github.com/tikv/pd/pkg/utils/testutil"
-	"github.com/tikv/pd/pkg/utils/typeutil"
 )
 
 func TestActivity(t *testing.T) {
@@ -41,37 +39,33 @@ func TestActivity(t *testing.T) {
 	cluster.AddRegionStore(2, 0)
 	cluster.AddLeaderRegion(1, 1)
 	region := cluster.GetRegion(1)
-	msg := &pdpb.RegionHeartbeatResponse{
-		ChangePeer: &pdpb.ChangePeer{Peer: &metapb.Peer{Id: 2, StoreId: 2}, ChangeType: eraftpb.ConfChangeType_AddLearnerNode},
-	}
-
 	hbs := hbstream.NewTestHeartbeatStreams(ctx, cluster.ID, cluster, true)
 	stream1, stream2 := NewHeartbeatStream(), NewHeartbeatStream()
 
 	// Active stream is stream1.
 	hbs.BindStream(1, stream1)
 	testutil.Eventually(re, func() bool {
-		newMsg := typeutil.DeepClone(msg, core.RegionHeartbeatResponseFactory)
-		hbs.SendMsg(region, newMsg)
+		msg := &hbstream.Operation{ChangePeer: &pdpb.ChangePeer{Peer: &metapb.Peer{Id: 2, StoreId: 2}, ChangeType: eraftpb.ConfChangeType_AddLearnerNode}}
+		hbs.SendMsg(region, msg)
 		return stream1.Recv() != nil && stream2.Recv() == nil
 	})
 	// Rebind to stream2.
 	hbs.BindStream(1, stream2)
 	testutil.Eventually(re, func() bool {
-		newMsg := typeutil.DeepClone(msg, core.RegionHeartbeatResponseFactory)
-		hbs.SendMsg(region, newMsg)
+		msg := &hbstream.Operation{ChangePeer: &pdpb.ChangePeer{Peer: &metapb.Peer{Id: 2, StoreId: 2}, ChangeType: eraftpb.ConfChangeType_AddLearnerNode}}
+		hbs.SendMsg(region, msg)
 		return stream1.Recv() == nil && stream2.Recv() != nil
 	})
 	// SendErr to stream2.
 	hbs.SendErr(pdpb.ErrorType_UNKNOWN, "test error", &metapb.Peer{Id: 1, StoreId: 1})
 	res := stream2.Recv()
 	re.NotNil(res)
-	re.NotNil(res.GetHeader().GetError())
+	re.NotNil(res.(*pdpb.RegionHeartbeatResponse).GetHeader().GetError())
 	// Switch back to 1 again.
 	hbs.BindStream(1, stream1)
 	testutil.Eventually(re, func() bool {
-		newMsg := typeutil.DeepClone(msg, core.RegionHeartbeatResponseFactory)
-		hbs.SendMsg(region, newMsg)
+		msg := &hbstream.Operation{ChangePeer: &pdpb.ChangePeer{Peer: &metapb.Peer{Id: 2, StoreId: 2}, ChangeType: eraftpb.ConfChangeType_AddLearnerNode}}
+		hbs.SendMsg(region, msg)
 		return stream1.Recv() != nil && stream2.Recv() == nil
 	})
 }