Resume the max retry check for stale read processing

internal/locate/region_request.go

@@ -574,7 +574,7 @@ func (state *accessFollower) next(bo *retry.Backoffer, selector *replicaSelector
 			logutil.BgLogger().Warn("unable to find stores with given labels")
 		}
 		leader := selector.replicas[state.leaderIdx]
-		if leader.isEpochStale() || (!state.option.leaderOnly && leader.isExhausted(1)) {
+		if leader.isEpochStale() || state.IsLeaderExhausted(leader) {
 			metrics.TiKVReplicaSelectorFailureCounter.WithLabelValues("exhausted").Inc()
 			selector.invalidateRegion()
 			return nil, nil
@@ -593,6 +593,19 @@ func (state *accessFollower) next(bo *retry.Backoffer, selector *replicaSelector
 	return rpcCtx, nil
 }
 
+func (state *accessFollower) IsLeaderExhausted(leader *replica) bool {
+	// Allow another extra retry for the following case:
+	// 1. The stale read is enabled and leader peer is selected as the target peer at first.
+	// 2. Data is not ready is returned from the leader peer.
+	// 3. Stale read flag is removed and processing falls back to snapshot read on the leader peer.
+	// 4. The leader peer should be retried again using snapshot read.
+	if state.isStaleRead && state.option.leaderOnly {
+		return leader.isExhausted(2)
+	} else {
+		return leader.isExhausted(1)
+	}
+}
+
 func (state *accessFollower) onSendFailure(bo *retry.Backoffer, selector *replicaSelector, cause error) {
 	if selector.checkLiveness(bo, selector.targetReplica()) != reachable {
 		selector.invalidateReplicaStore(selector.targetReplica(), cause)

internal/locate/region_request3_test.go

@@ -36,6 +36,7 @@ package locate
 
 import (
 	"context"
+	"strconv"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -1018,3 +1019,84 @@ func (s *testRegionRequestToThreeStoresSuite) TestAccessFollowerAfter1TiKVDown()
 		s.Equal(0, bo.GetTotalBackoffTimes())
 	}
 }
+
+func (s *testRegionRequestToThreeStoresSuite) TestStaleReadFallback() {
+	leaderStore, _ := s.loadAndGetLeaderStore()
+	leaderLabel := []*metapb.StoreLabel{
+		{
+			Key:   "id",
+			Value: strconv.FormatUint(leaderStore.StoreID(), 10),
+		},
+	}
+	regionLoc, err := s.cache.LocateRegionByID(s.bo, s.regionID)
+	s.Nil(err)
+	s.NotNil(regionLoc)
+	value := []byte("value")
+
+	type testState struct {
+		tryTimes uint8
+		succ     bool
+	}
+
+	state := &testState{}
+	s.regionRequestSender.client = &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+		select {
+		case <-ctx.Done():
+			return nil, errors.New("timeout")
+		default:
+		}
+		// Return `DataIsNotReady` for the first time on leader.
+		if state.tryTimes == 0 {
+			state.tryTimes++
+			return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{RegionError: &errorpb.Error{
+				DataIsNotReady: &errorpb.DataIsNotReady{},
+			}}}, nil
+		} else if state.tryTimes == 1 && state.succ {
+			state.tryTimes++
+			return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{Value: value}}, nil
+		}
+		state.tryTimes++
+		return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{RegionError: &errorpb.Error{
+			DiskFull: &errorpb.DiskFull{},
+		}}}, nil
+	}}
+
+	region := s.cache.getRegionByIDFromCache(regionLoc.Region.GetID())
+	s.True(region.isValid())
+
+	// Test the successful path.
+	state.succ = true
+	req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdGet, &kvrpcpb.GetRequest{Key: []byte("key")}, kv.ReplicaReadLeader, nil)
+	req.ReadReplicaScope = oracle.GlobalTxnScope
+	req.TxnScope = oracle.GlobalTxnScope
+	req.EnableStaleRead()
+	req.ReplicaReadType = kv.ReplicaReadMixed
+	var ops []StoreSelectorOption
+	ops = append(ops, WithMatchLabels(leaderLabel))
+
+	ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
+	bo := retry.NewBackoffer(ctx, -1)
+	s.Nil(err)
+	resp, _, err := s.regionRequestSender.SendReqCtx(bo, req, regionLoc.Region, time.Second, tikvrpc.TiKV, ops...)
+	s.Nil(err)
+
+	regionErr, err := resp.GetRegionError()
+	s.Nil(err)
+	s.Nil(regionErr)
+	getResp, ok := resp.Resp.(*kvrpcpb.GetResponse)
+	s.True(ok)
+	s.Equal(getResp.Value, value)
+
+	// Test the fail path leader retry limit is reached, epoch not match error would be returned.
+	state.tryTimes = 0
+	state.succ = false
+	req.EnableStaleRead()
+	resp, _, err = s.regionRequestSender.SendReqCtx(bo, req, regionLoc.Region, time.Second, tikvrpc.TiKV, ops...)
+	s.Nil(err)
+
+	regionErr, err = resp.GetRegionError()
+	s.Nil(err)
+	s.NotNil(regionErr)
+	s.NotNil(regionErr.GetEpochNotMatch())
+	s.Nil(regionErr.GetDiskFull())
+}