locate: retry on ErrClientGetResourceGroup for transient errors

internal/locate/region_request.go

@@ -1865,7 +1865,8 @@ func (s *RegionRequestSender) onSendFail(bo *retry.Backoffer, ctx *RPCContext, r
 		}
 	}
 
-	// don't need to retry for ResourceGroup error
+	// don't need to retry for ResourceGroup error when specific errors are returned from server.
+	// when other error is returned from server (not leader etc.), back off and retry.
 	if errors.Is(err, pderr.ErrClientResourceGroupThrottled) {
 		return err
 	}
@@ -1874,7 +1875,13 @@ func (s *RegionRequestSender) onSendFail(bo *retry.Backoffer, ctx *RPCContext, r
 	}
 	var errGetResourceGroup *pderr.ErrClientGetResourceGroup
 	if errors.As(err, &errGetResourceGroup) {
-		return err
+		if errGetResourceGroup.Cause == "resource group not found" {
+			return err
+		} else {
+			// back off and retry
+			err = bo.Backoff(retry.BoPDRPC, errors.Errorf("send pd request error:%v, ctx: %v, try next peer later", err, ctx))
+			return err
+		}
 	}
 
 	// Retry on send request failure when it's not canceled.

internal/locate/region_request_test.go

@@ -62,6 +62,7 @@ import (
 	"github.com/tikv/client-go/v2/internal/mockstore/mocktikv"
 	"github.com/tikv/client-go/v2/internal/retry"
 	"github.com/tikv/client-go/v2/tikvrpc"
+	pderr "github.com/tikv/pd/client/errs"
 	"google.golang.org/grpc"
 )
 
@@ -151,6 +152,60 @@ func (s *testRegionRequestToSingleStoreSuite) TestOnRegionError() {
 	}()
 }
 
+func (s *testRegionRequestToSingleStoreSuite) TestOnSendFailByResourceGroupError() {
+	req := tikvrpc.NewRequest(tikvrpc.CmdRawPut, &kvrpcpb.RawPutRequest{
+		Key:   []byte("key"),
+		Value: []byte("value"),
+	})
+	region, err := s.cache.LocateRegionByID(s.bo, s.region)
+	s.Nil(err)
+	s.NotNil(region)
+
+	// test ErrClientGetResourceGroup with "resource group not found" cause should not retry
+	func() {
+		oc := s.regionRequestSender.client
+		defer func() {
+			s.regionRequestSender.client = oc
+		}()
+		s.regionRequestSender.client = &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+			return nil, &pderr.ErrClientGetResourceGroup{Cause: "resource group not found"}
+		}}
+		bo := retry.NewBackofferWithVars(context.Background(), 5, nil)
+		_, _, err := s.regionRequestSender.SendReq(bo, req, region.Region, time.Second)
+		s.NotNil(err)
+		var errGetResourceGroup *pderr.ErrClientGetResourceGroup
+		s.True(errors.As(err, &errGetResourceGroup))
+		s.Equal("resource group not found", errGetResourceGroup.Cause)
+		// should not have backoff since error is directly returned
+		s.Equal(0, bo.GetTotalBackoffTimes())
+	}()
+
+	// test ErrClientGetResourceGroup with other causes should retry with backoff
+	func() {
+		oc := s.regionRequestSender.client
+		defer func() {
+			s.regionRequestSender.client = oc
+		}()
+		callCount := 0
+		s.regionRequestSender.client = &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+			callCount++
+			if callCount < 2 {
+				return nil, &pderr.ErrClientGetResourceGroup{Cause: "not leader"}
+			}
+			// return success after first retry
+			return &tikvrpc.Response{Resp: &kvrpcpb.RawPutResponse{}}, nil
+		}}
+		bo := retry.NewBackofferWithVars(context.Background(), 2000, nil)
+		resp, _, err := s.regionRequestSender.SendReq(bo, req, region.Region, time.Second)
+		s.Nil(err)
+		s.NotNil(resp)
+		// should have backoff since error triggered retry
+		s.Greater(bo.GetTotalBackoffTimes(), 0)
+		// should have retried at least once
+		s.GreaterOrEqual(callCount, 2)
+	}()
+}
+
 func (s *testRegionRequestToSingleStoreSuite) TestOnSendFailedWithStoreRestart() {
 	req := tikvrpc.NewRequest(tikvrpc.CmdRawPut, &kvrpcpb.RawPutRequest{
 		Key:   []byte("key"),