mcs: support forwarding split and scatter request

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-20230920042517-db656f45023b
+	github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a
 	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

go.sum

@@ -439,8 +439,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 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-20230920042517-db656f45023b h1:XwwIxepR+uuSYWhdQtstEdr67XUE7X6lpSIHVh5iWjs=
-github.com/pingcap/kvproto v0.0.0-20230920042517-db656f45023b/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
+github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a h1:xbak0nFiTkinvdtr/EUK0dvdkX32twKRpIgsMRiiVQA=
+github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a/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=

pkg/mcs/scheduling/server/cluster.go

@@ -20,7 +20,9 @@
 	"github.com/tikv/pd/pkg/schedule/labeler"
 	"github.com/tikv/pd/pkg/schedule/operator"
 	"github.com/tikv/pd/pkg/schedule/placement"
+	"github.com/tikv/pd/pkg/schedule/scatter"
 	"github.com/tikv/pd/pkg/schedule/schedulers"
+	"github.com/tikv/pd/pkg/schedule/splitter"
 	"github.com/tikv/pd/pkg/slice"
 	"github.com/tikv/pd/pkg/statistics"
 	"github.com/tikv/pd/pkg/statistics/buckets"
@@ -130,6 +132,16 @@
 	return c.labelerManager
 }
 
+// GetRegionSplitter returns the region splitter.
+func (c *Cluster) GetRegionSplitter() *splitter.RegionSplitter {
+	return c.coordinator.GetRegionSplitter()
+}
+
+// GetRegionScatterer returns the region scatter.
+func (c *Cluster) GetRegionScatterer() *scatter.RegionScatterer {
+	return c.coordinator.GetRegionScatterer()
+}
+
 // GetStoresLoads returns load stats of all stores.
 func (c *Cluster) GetStoresLoads() map[uint64][]float64 {
 	return c.hotStat.GetStoresLoads()

pkg/mcs/scheduling/server/grpc_service.go

@@ -140,13 +140,7 @@
 
 		c := s.GetCluster()
 		if c == nil {
-			resp := &schedulingpb.RegionHeartbeatResponse{Header: &schedulingpb.ResponseHeader{
-				ClusterId: s.clusterID,
-				Error: &schedulingpb.Error{
-					Type:    schedulingpb.ErrorType_NOT_BOOTSTRAPPED,
-					Message: "scheduling server is not initialized yet",
-				},
-			}}
+			resp := &schedulingpb.RegionHeartbeatResponse{Header: s.notBootstrappedHeader()}
 			err := server.Send(resp)
 			return errors.WithStack(err)
 		}
@@ -177,7 +171,7 @@
 	if c == nil {
 		// TODO: add metrics
 		log.Info("cluster isn't initialized")
-		return &schedulingpb.StoreHeartbeatResponse{Header: &schedulingpb.ResponseHeader{ClusterId: s.clusterID}}, nil
+		return &schedulingpb.StoreHeartbeatResponse{Header: s.notBootstrappedHeader()}, nil
 	}
 
 	if c.GetStore(request.GetStats().GetStoreId()) == nil {
@@ -191,6 +185,75 @@
 	return &schedulingpb.StoreHeartbeatResponse{Header: &schedulingpb.ResponseHeader{ClusterId: s.clusterID}}, nil
 }
 
+// SplitRegions split regions by the given split keys
+func (s *Service) SplitRegions(ctx context.Context, request *schedulingpb.SplitRegionsRequest) (*schedulingpb.SplitRegionsResponse, error) {
+	c := s.GetCluster()
+	if c == nil {
+		return &schedulingpb.SplitRegionsResponse{Header: s.notBootstrappedHeader()}, nil
+	}
+	finishedPercentage, newRegionIDs := c.GetRegionSplitter().SplitRegions(ctx, request.GetSplitKeys(), int(request.GetRetryLimit()))
+	return &schedulingpb.SplitRegionsResponse{
+		Header:             s.header(),
+		RegionsId:          newRegionIDs,
+		FinishedPercentage: uint64(finishedPercentage),
+	}, nil
+}
+
+// ScatterRegions implements gRPC PDServer.
+func (s *Service) ScatterRegions(ctx context.Context, request *schedulingpb.ScatterRegionsRequest) (*schedulingpb.ScatterRegionsResponse, error) {
+	c := s.GetCluster()
+	if c == nil {
+		return &schedulingpb.ScatterRegionsResponse{Header: s.notBootstrappedHeader()}, nil
+	}
+
+	opsCount, failures, err := c.GetRegionScatterer().ScatterRegionsByID(request.GetRegionsId(), request.GetGroup(), int(request.GetRetryLimit()), request.GetSkipStoreLimit())
+	if err != nil {
+		return nil, err
+	}
+	percentage := 100
+	if len(failures) > 0 {
+		percentage = 100 - 100*len(failures)/(opsCount+len(failures))
+		log.Debug("scatter regions", zap.Errors("failures", func() []error {
+			r := make([]error, 0, len(failures))
+			for _, err := range failures {
+				r = append(r, err)
+			}
+			return r
+		}()))
+	}
+	return &schedulingpb.ScatterRegionsResponse{
+		Header:             s.header(),
+		FinishedPercentage: uint64(percentage),
+	}, nil
+}
+
+// GetOperator gets information about the operator belonging to the specify region.
+func (s *Service) GetOperator(ctx context.Context, request *schedulingpb.GetOperatorRequest) (*schedulingpb.GetOperatorResponse, error) {
+	c := s.GetCluster()
+	if c == nil {
+		return &schedulingpb.GetOperatorResponse{Header: s.notBootstrappedHeader()}, nil
+	}
+
+	opController := c.GetCoordinator().GetOperatorController()
+	requestID := request.GetRegionId()
+	r := opController.GetOperatorStatus(requestID)
+	if r == nil {
+		header := s.errorHeader(&schedulingpb.Error{
+			Type:    schedulingpb.ErrorType_UNKNOWN,
+			Message: "Not Found",
+		})
+		return &schedulingpb.GetOperatorResponse{Header: header}, nil
+	}
+
+	return &schedulingpb.GetOperatorResponse{
+		Header:   s.header(),
+		RegionId: requestID,
+		Desc:     []byte(r.Desc()),
+		Kind:     []byte(r.Kind().String()),
+		Status:   r.Status,
+	}, nil
+}
+
 // RegisterGRPCService registers the service to gRPC server.
 func (s *Service) RegisterGRPCService(g *grpc.Server) {
 	schedulingpb.RegisterSchedulingServer(g, s)
@@ -201,3 +264,29 @@
 	handler, group := SetUpRestHandler(s)
 	apiutil.RegisterUserDefinedHandlers(userDefineHandlers, &group, handler)
 }
+
+func (s *Service) errorHeader(err *schedulingpb.Error) *schedulingpb.ResponseHeader {
+	return &schedulingpb.ResponseHeader{
+		ClusterId: s.clusterID,
+		Error:     err,
+	}
+}
+
+func (s *Service) notBootstrappedHeader() *schedulingpb.ResponseHeader {
+	return s.errorHeader(&schedulingpb.Error{
+		Type:    schedulingpb.ErrorType_NOT_BOOTSTRAPPED,
+		Message: "cluster is not initialized",
+	})
+}
+
+func (s *Service) header() *schedulingpb.ResponseHeader {
+	if s.clusterID == 0 {
+		return s.wrapErrorToHeader(schedulingpb.ErrorType_NOT_BOOTSTRAPPED, "cluster id is not ready")
+	}
+	return &schedulingpb.ResponseHeader{ClusterId: s.clusterID}
+}
+
+func (s *Service) wrapErrorToHeader(
+	errorType schedulingpb.ErrorType, message string) *schedulingpb.ResponseHeader {
+	return s.errorHeader(&schedulingpb.Error{Type: errorType, Message: message})
+}

pkg/tso/allocator_manager.go

@@ -413,7 +413,7 @@
 		if err != nil {
 			log.Warn("get server id and dcLocation from etcd failed, invalid server id",
 				logutil.CondUint32("keyspace-group-id", am.kgID, am.kgID > 0),
-				zap.Any("splitted-serverPath", serverPath),
+				zap.Any("split-serverPath", serverPath),
 				zap.String("dc-location", dcLocation),
 				errs.ZapError(err))
 			continue
@@ -973,8 +973,8 @@
 		if err != nil {
 			return nil, err
 		}
-		splittedKey := strings.Split(string(kv.Key), "/")
-		dcLocation := splittedKey[len(splittedKey)-1]
+		splitKey := strings.Split(string(kv.Key), "/")
+		dcLocation := splitKey[len(splitKey)-1]
 		dcLocationSuffix[dcLocation] = int32(suffix)
 	}
 	return dcLocationSuffix, nil

server/grpc_service.go

@@ -1770,6 +1770,35 @@
 
 // ScatterRegion implements gRPC PDServer.
 func (s *GrpcServer) ScatterRegion(ctx context.Context, request *pdpb.ScatterRegionRequest) (*pdpb.ScatterRegionResponse, error) {
+	if s.IsAPIServiceMode() {
+		s.updateSchedulingClient(ctx)
+		if s.schedulingClient.Load() != nil {
+			regionsID := request.GetRegionsId()
+			if len(regionsID) == 0 {
+				return &pdpb.ScatterRegionResponse{
+					Header: s.invalidValue("regions id is required"),
+				}, nil
+			}
+			req := &schedulingpb.ScatterRegionsRequest{
+				Header: &schedulingpb.RequestHeader{
+					ClusterId: request.GetHeader().GetClusterId(),
+					SenderId:  request.GetHeader().GetSenderId(),
+				},
+				RegionsId:      regionsID,
+				Group:          request.GetGroup(),
+				RetryLimit:     request.GetRetryLimit(),
+				SkipStoreLimit: request.GetSkipStoreLimit(),
+			}
+			resp, err := s.schedulingClient.Load().(*schedulingClient).getClient().ScatterRegions(ctx, req)
+			if err != nil {
+				// reset to let it be updated in the next request
+				s.schedulingClient.Store(&schedulingClient{})
+				return s.convertScatterResponse(resp), err
+			}
+			return s.convertScatterResponse(resp), nil
+		}
+	}
+
 	fn := func(ctx context.Context, client *grpc.ClientConn) (interface{}, error) {
 		return pdpb.NewPDClient(client).ScatterRegion(ctx, request)
 	}
@@ -1960,6 +1989,25 @@
 
 // GetOperator gets information about the operator belonging to the specify region.
 func (s *GrpcServer) GetOperator(ctx context.Context, request *pdpb.GetOperatorRequest) (*pdpb.GetOperatorResponse, error) {
+	if s.IsAPIServiceMode() {
+		s.updateSchedulingClient(ctx)
+		if s.schedulingClient.Load() != nil {
+			req := &schedulingpb.GetOperatorRequest{
+				Header: &schedulingpb.RequestHeader{
+					ClusterId: request.GetHeader().GetClusterId(),
+					SenderId:  request.GetHeader().GetSenderId(),
+				},
+				RegionId: request.GetRegionId(),
+			}
+			resp, err := s.schedulingClient.Load().(*schedulingClient).getClient().GetOperator(ctx, req)
+			if err != nil {
+				// reset to let it be updated in the next request
+				s.schedulingClient.Store(&schedulingClient{})
+				return s.convertOperatorResponse(resp), err
+			}
+			return s.convertOperatorResponse(resp), nil
+		}
+	}
 	fn := func(ctx context.Context, client *grpc.ClientConn) (interface{}, error) {
 		return pdpb.NewPDClient(client).GetOperator(ctx, request)
 	}
@@ -2049,6 +2097,45 @@
 	})
 }
 
+func (s *GrpcServer) convertHeader(header *schedulingpb.ResponseHeader) *pdpb.ResponseHeader {
+	switch header.GetError().GetType() {
+	case schedulingpb.ErrorType_UNKNOWN:
+		return &pdpb.ResponseHeader{
+			ClusterId: header.GetClusterId(),
+			Error: &pdpb.Error{
+				Type:    pdpb.ErrorType_UNKNOWN,
+				Message: header.GetError().GetMessage(),
+			},
+		}
+	default:
+		return &pdpb.ResponseHeader{ClusterId: header.GetClusterId()}
+	}
+}
+
+func (s *GrpcServer) convertSplitResponse(resp *schedulingpb.SplitRegionsResponse) *pdpb.SplitRegionsResponse {
+	return &pdpb.SplitRegionsResponse{
+		Header:             s.convertHeader(resp.GetHeader()),
+		FinishedPercentage: resp.GetFinishedPercentage(),
+	}
+}
+
+func (s *GrpcServer) convertScatterResponse(resp *schedulingpb.ScatterRegionsResponse) *pdpb.ScatterRegionResponse {
+	return &pdpb.ScatterRegionResponse{
+		Header:             s.convertHeader(resp.GetHeader()),
+		FinishedPercentage: resp.GetFinishedPercentage(),
+	}
+}
+
+func (s *GrpcServer) convertOperatorResponse(resp *schedulingpb.GetOperatorResponse) *pdpb.GetOperatorResponse {
+	return &pdpb.GetOperatorResponse{
+		Header:   s.convertHeader(resp.GetHeader()),
+		RegionId: resp.GetRegionId(),
+		Desc:     resp.GetDesc(),
+		Kind:     resp.GetKind(),
+		Status:   resp.GetStatus(),
+	}
+}
+
 // Only used for the TestLocalAllocatorLeaderChange.
 var mockLocalAllocatorLeaderChangeFlag = false
 
@@ -2153,6 +2240,27 @@
 
 // SplitRegions split regions by the given split keys
 func (s *GrpcServer) SplitRegions(ctx context.Context, request *pdpb.SplitRegionsRequest) (*pdpb.SplitRegionsResponse, error) {
+	if s.IsAPIServiceMode() {
+		s.updateSchedulingClient(ctx)
+		if s.schedulingClient.Load() != nil {
+			req := &schedulingpb.SplitRegionsRequest{
+				Header: &schedulingpb.RequestHeader{
+					ClusterId: request.GetHeader().GetClusterId(),
+					SenderId:  request.GetHeader().GetSenderId(),
+				},
+				SplitKeys:  request.GetSplitKeys(),
+				RetryLimit: request.GetRetryLimit(),
+			}
+			resp, err := s.schedulingClient.Load().(*schedulingClient).getClient().SplitRegions(ctx, req)
+			if err != nil {
+				// reset to let it be updated in the next request
+				s.schedulingClient.Store(&schedulingClient{})
+				return s.convertSplitResponse(resp), err
+			}
+			return s.convertSplitResponse(resp), nil
+		}
+	}
+
 	fn := func(ctx context.Context, client *grpc.ClientConn) (interface{}, error) {
 		return pdpb.NewPDClient(client).SplitRegions(ctx, request)
 	}
@@ -2175,7 +2283,7 @@
 }
 
 // SplitAndScatterRegions split regions by the given split keys, and scatter regions.
-// Only regions which splited successfully will be scattered.
+// Only regions which split successfully will be scattered.
 // scatterFinishedPercentage indicates the percentage of successfully splited regions that are scattered.
 func (s *GrpcServer) SplitAndScatterRegions(ctx context.Context, request *pdpb.SplitAndScatterRegionsRequest) (*pdpb.SplitAndScatterRegionsResponse, error) {
 	fn := func(ctx context.Context, client *grpc.ClientConn) (interface{}, error) {

tests/integrations/client/go.mod

@@ -13,7 +13,7 @@ replace google.golang.org/grpc v1.54.0 => google.golang.org/grpc v1.26.0
 require (
 	github.com/docker/go-units v0.4.0
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
-	github.com/pingcap/kvproto v0.0.0-20230920042517-db656f45023b
+	github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/stretchr/testify v1.8.2
 	github.com/tikv/pd v0.0.0-00010101000000-000000000000

tests/integrations/client/go.sum

@@ -402,8 +402,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 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-20230920042517-db656f45023b h1:XwwIxepR+uuSYWhdQtstEdr67XUE7X6lpSIHVh5iWjs=
-github.com/pingcap/kvproto v0.0.0-20230920042517-db656f45023b/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
+github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a h1:xbak0nFiTkinvdtr/EUK0dvdkX32twKRpIgsMRiiVQA=
+github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
 github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 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=

tests/integrations/mcs/go.mod

@@ -13,7 +13,7 @@ replace google.golang.org/grpc v1.54.0 => google.golang.org/grpc v1.26.0
 require (
 	github.com/docker/go-units v0.4.0
 	github.com/pingcap/failpoint v0.0.0-20210918120811-547c13e3eb00
-	github.com/pingcap/kvproto v0.0.0-20230920042517-db656f45023b
+	github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/stretchr/testify v1.8.2
 	github.com/tikv/pd v0.0.0-00010101000000-000000000000

tests/integrations/mcs/go.sum

@@ -407,8 +407,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 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-20230920042517-db656f45023b h1:XwwIxepR+uuSYWhdQtstEdr67XUE7X6lpSIHVh5iWjs=
-github.com/pingcap/kvproto v0.0.0-20230920042517-db656f45023b/go.mod h1:r0q/CFcwvyeRhKtoqzmWMBebrtpIziQQ9vR+JKh1knc=
+github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a h1:xbak0nFiTkinvdtr/EUK0dvdkX32twKRpIgsMRiiVQA=
+github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a/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=

tests/integrations/tso/go.mod

@@ -13,7 +13,7 @@ replace google.golang.org/grpc v1.54.0 => google.golang.org/grpc v1.26.0
 
 require (
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
-	github.com/pingcap/kvproto v0.0.0-20230920042517-db656f45023b
+	github.com/pingcap/kvproto v0.0.0-20231017055627-c06b434f3c3a
 	github.com/stretchr/testify v1.8.4
 	github.com/tikv/pd v0.0.0-00010101000000-000000000000
 	github.com/tikv/pd/client v0.0.0-00010101000000-000000000000