*: follower support to handle GetRegion and other region api

client/client.go

@@ -298,7 +298,7 @@
 type client struct {
 	keyspaceID      uint32
 	svrUrls         []string
-	pdSvcDiscovery  ServiceDiscovery
+	pdSvcDiscovery  *pdServiceDiscovery
 	tokenDispatcher *tokenDispatcher
 
 	// For service mode switching.
@@ -503,7 +503,7 @@
 			return err
 		}
 		// c.keyspaceID is the source of truth for keyspace id.
-		c.pdSvcDiscovery.(*pdServiceDiscovery).SetKeyspaceID(c.keyspaceID)
+		c.pdSvcDiscovery.SetKeyspaceID(c.keyspaceID)
 		return nil
 	}
 
@@ -733,6 +733,23 @@
 	return pdpb.NewPDClient(serviceClient.GetClientConn()), serviceClient.BuildGRPCTargetContext(ctx, true)
 }
 
+// getClientAndContext returns the leader pd client and the original context. If leader is unhealthy, it returns
+// follower pd client and the context which holds forward information.
+func (c *client) getRegionAPIClientAndContext(ctx context.Context, allowFollower bool) (ServiceClient, context.Context) {
+	var serviceClient ServiceClient
+	if allowFollower {
+		serviceClient = c.pdSvcDiscovery.getServiceClientByKind(regionAPIKind)
+		if serviceClient != nil {
+			return serviceClient, serviceClient.BuildGRPCTargetContext(ctx, !allowFollower)
+		}
+	}
+	serviceClient = c.pdSvcDiscovery.GetServiceClient()
+	if serviceClient == nil {
+		return nil, ctx
+	}
+	return serviceClient, serviceClient.BuildGRPCTargetContext(ctx, !allowFollower)
+}
+
 func (c *client) GetTSAsync(ctx context.Context) TSFuture {
 	return c.GetLocalTSAsync(ctx, globalDCLocation)
 }
@@ -885,6 +902,7 @@
 	start := time.Now()
 	defer func() { cmdDurationGetRegion.Observe(time.Since(start).Seconds()) }()
 	ctx, cancel := context.WithTimeout(ctx, c.option.timeout)
+	defer cancel()
 
 	options := &GetRegionOp{}
 	for _, opt := range opts {
@@ -895,13 +913,18 @@
 		RegionKey:   key,
 		NeedBuckets: options.needBuckets,
 	}
-	protoClient, ctx := c.getClientAndContext(ctx)
-	if protoClient == nil {
-		cancel()
+	serviceClient, cctx := c.getRegionAPIClientAndContext(ctx, options.allowFollowerHandle && c.option.getEnableFollowerHandle())
+	if serviceClient == nil {
 		return nil, errs.ErrClientGetProtoClient
 	}
-	resp, err := protoClient.GetRegion(ctx, req)
-	cancel()
+	resp, err := pdpb.NewPDClient(serviceClient.GetClientConn()).GetRegion(cctx, req)
+	if serviceClient.NeedRetry(resp.GetHeader().GetError(), err) {
+		protoClient, cctx := c.getClientAndContext(ctx)
+		if protoClient == nil {
+			return nil, errs.ErrClientGetProtoClient
+		}
+		resp, err = protoClient.GetRegion(cctx, req)
+	}
 
 	if err = c.respForErr(cmdFailDurationGetRegion, start, err, resp.GetHeader()); err != nil {
 		return nil, err
@@ -917,6 +940,7 @@
 	start := time.Now()
 	defer func() { cmdDurationGetPrevRegion.Observe(time.Since(start).Seconds()) }()
 	ctx, cancel := context.WithTimeout(ctx, c.option.timeout)
+	defer cancel()
 
 	options := &GetRegionOp{}
 	for _, opt := range opts {
@@ -927,13 +951,18 @@
 		RegionKey:   key,
 		NeedBuckets: options.needBuckets,
 	}
-	protoClient, ctx := c.getClientAndContext(ctx)
-	if protoClient == nil {
-		cancel()
+	serviceClient, cctx := c.getRegionAPIClientAndContext(ctx, options.allowFollowerHandle && c.option.getEnableFollowerHandle())
+	if serviceClient == nil {
 		return nil, errs.ErrClientGetProtoClient
 	}
-	resp, err := protoClient.GetPrevRegion(ctx, req)
-	cancel()
+	resp, err := pdpb.NewPDClient(serviceClient.GetClientConn()).GetPrevRegion(cctx, req)
+	if serviceClient.NeedRetry(resp.GetHeader().GetError(), err) {
+		protoClient, cctx := c.getClientAndContext(ctx)
+		if protoClient == nil {
+			return nil, errs.ErrClientGetProtoClient
+		}
+		resp, err = protoClient.GetPrevRegion(cctx, req)
+	}
 
 	if err = c.respForErr(cmdFailDurationGetPrevRegion, start, err, resp.GetHeader()); err != nil {
 		return nil, err
@@ -949,6 +978,7 @@
 	start := time.Now()
 	defer func() { cmdDurationGetRegionByID.Observe(time.Since(start).Seconds()) }()
 	ctx, cancel := context.WithTimeout(ctx, c.option.timeout)
+	defer cancel()
 
 	options := &GetRegionOp{}
 	for _, opt := range opts {
@@ -959,13 +989,18 @@
 		RegionId:    regionID,
 		NeedBuckets: options.needBuckets,
 	}
-	protoClient, ctx := c.getClientAndContext(ctx)
-	if protoClient == nil {
-		cancel()
+	serviceClient, cctx := c.getRegionAPIClientAndContext(ctx, options.allowFollowerHandle && c.option.getEnableFollowerHandle())
+	if serviceClient == nil {
 		return nil, errs.ErrClientGetProtoClient
 	}
-	resp, err := protoClient.GetRegionByID(ctx, req)
-	cancel()
+	resp, err := pdpb.NewPDClient(serviceClient.GetClientConn()).GetRegionByID(cctx, req)
+	if serviceClient.NeedRetry(resp.GetHeader().GetError(), err) {
+		protoClient, cctx := c.getClientAndContext(ctx)
+		if protoClient == nil {
+			return nil, errs.ErrClientGetProtoClient
+		}
+		resp, err = protoClient.GetRegionByID(cctx, req)
+	}
 
 	if err = c.respForErr(cmdFailedDurationGetRegionByID, start, err, resp.GetHeader()); err != nil {
 		return nil, err
@@ -987,18 +1022,28 @@
 		scanCtx, cancel = context.WithTimeout(ctx, c.option.timeout)
 		defer cancel()
 	}
+	options := &GetRegionOp{}
+	for _, opt := range opts {
+		opt(options)
+	}
 	req := &pdpb.ScanRegionsRequest{
 		Header:   c.requestHeader(),
 		StartKey: key,
 		EndKey:   endKey,
 		Limit:    int32(limit),
 	}
-	protoClient, scanCtx := c.getClientAndContext(scanCtx)
-	if protoClient == nil {
-		cancel()
+	serviceClient, cctx := c.getRegionAPIClientAndContext(scanCtx, options.allowFollowerHandle && c.option.getEnableFollowerHandle())
+	if serviceClient == nil {
 		return nil, errs.ErrClientGetProtoClient
 	}
-	resp, err := protoClient.ScanRegions(scanCtx, req)
+	resp, err := pdpb.NewPDClient(serviceClient.GetClientConn()).ScanRegions(cctx, req)
+	if !serviceClient.IsConnectedToLeader() && err != nil || resp.Header.GetError() != nil {
+		protoClient, cctx := c.getClientAndContext(scanCtx)
+		if protoClient == nil {
+			return nil, errs.ErrClientGetProtoClient
+		}
+		resp, err = protoClient.ScanRegions(cctx, req)
+	}
 
 	if err = c.respForErr(cmdFailedDurationScanRegions, start, err, resp.GetHeader()); err != nil {
 		return nil, err

client/pd_service_discovery.go

@@ -53,6 +53,7 @@ type apiKind int
 
 const (
 	forwardAPIKind apiKind = iota
+	regionAPIKind
 	apiKindCount
 )
 
@@ -445,7 +446,7 @@ func newPDServiceDiscovery(
 		ctx:                 ctx,
 		cancel:              cancel,
 		wg:                  wg,
-		apiCandidateNodes:   [apiKindCount]*pdServiceBalancer{newPDServiceBalancer(emptyErrorFn)},
+		apiCandidateNodes:   [apiKindCount]*pdServiceBalancer{newPDServiceBalancer(emptyErrorFn), newPDServiceBalancer(regionAPIErrorFn)},
 		serviceModeUpdateCb: serviceModeUpdateCb,
 		updateKeyspaceIDCb:  updateKeyspaceIDCb,
 		keyspaceID:          keyspaceID,
@@ -563,6 +564,7 @@ func (c *pdServiceDiscovery) updateServiceModeLoop() {
 		}
 	}
 }
+
 func (c *pdServiceDiscovery) memberHealthCheckLoop() {
 	defer c.wg.Done()
 

pkg/utils/grpcutil/grpcutil.go

@@ -38,6 +38,8 @@
 const (
 	// ForwardMetadataKey is used to record the forwarded host of PD.
 	ForwardMetadataKey = "pd-forwarded-host"
+	// FollowerHandleMetadataKey is used to mark the permit of follower handle.
+	FollowerHandleMetadataKey = "pd-allow-follower-handle"
 )
 
 // TLSConfig is the configuration for supporting tls.
@@ -173,7 +175,7 @@
 func GetForwardedHost(ctx context.Context) string {
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
-		log.Debug("failed to get forwarding metadata")
+		log.Debug("failed to get gRPC incoming metadata when getting forwarded host")
 		return ""
 	}
 	if t, ok := md[ForwardMetadataKey]; ok {
@@ -182,6 +184,17 @@
 	return ""
 }
 
+// IsFollowerHandleEnabled returns the follower host in metadata.
+func IsFollowerHandleEnabled(ctx context.Context) bool {
+	md, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		log.Debug("failed to get gRPC incoming metadata when checking follower handle is enabled")
+		return false
+	}
+	_, ok = md[FollowerHandleMetadataKey]
+	return ok
+}
+
 func establish(ctx context.Context, addr string, tlsConfig *TLSConfig, do ...grpc.DialOption) (*grpc.ClientConn, error) {
 	tlsCfg, err := tlsConfig.ToTLSConfig()
 	if err != nil {

server/forward.go

@@ -384,16 +384,16 @@ func (s *GrpcServer) getForwardedHost(ctx, streamCtx context.Context, serviceNam
 	return forwardedHost, nil
 }
 
-func (s *GrpcServer) isLocalRequest(forwardedHost string) bool {
+func (s *GrpcServer) isLocalRequest(host string) bool {
 	failpoint.Inject("useForwardRequest", func() {
 		failpoint.Return(false)
 	})
-	if forwardedHost == "" {
+	if host == "" {
 		return true
 	}
 	memberAddrs := s.GetMember().Member().GetClientUrls()
 	for _, addr := range memberAddrs {
-		if addr == forwardedHost {
+		if addr == host {
 			return true
 		}
 	}