Repository: tikv/client-go
Pull request #1109 (2024-01-10 11:27:00)

make experimental improvements to replica selector
- pick #1100
- [stale read] support increasing timeout when retry leader
- [stale read] support disabling retry follower when leader is unavailable



Patch:
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 110b32a..b10958c 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -576,15 +576,22 @@ func (c *RPCContext) String() string {
 type contextPatcher struct {
 	staleRead   *bool
 	replicaRead *bool
+
+	timeoutFactor float64
 }
 
-func (patcher *contextPatcher) applyTo(pbCtx *kvrpcpb.Context) {
+func (patcher *contextPatcher) applyTo(pbCtx *kvrpcpb.Context, timeout time.Duration) time.Duration {
 	if patcher.staleRead != nil {
 		pbCtx.StaleRead = *patcher.staleRead
 	}
 	if patcher.replicaRead != nil {
 		pbCtx.ReplicaRead = *patcher.replicaRead
 	}
+	if patcher.timeoutFactor > 0 {
+		pbCtx.MaxExecutionDurationMs = uint64(float64(pbCtx.MaxExecutionDurationMs) * patcher.timeoutFactor)
+		timeout = time.Duration(float64(timeout) * patcher.timeoutFactor)
+	}
+	return timeout
 }
 
 type storeSelectorOp struct {
@@ -2448,6 +2455,8 @@ func (s *Store) reResolve(c *RegionCache) (bool, error) {
 	addr = store.GetAddress()
 	if s.addr != addr || !s.IsSameLabels(store.GetLabels()) {
 		newStore := &Store{storeID: s.storeID, addr: addr, saddr: store.GetStatusAddress(), storeType: storeType, labels: store.GetLabels(), state: uint64(resolved)}
+		newStore.livenessState = atomic.LoadUint32(&s.livenessState)
+		newStore.unreachableSince = s.unreachableSince
 		c.storeMu.Lock()
 		c.storeMu.stores[newStore.storeID] = newStore
 		c.storeMu.Unlock()
@@ -2600,14 +2609,27 @@ func (s *Store) startHealthCheckLoopIfNeeded(c *RegionCache, liveness livenessSt
 	// It may be already started by another thread.
 	if atomic.CompareAndSwapUint32(&s.livenessState, uint32(reachable), uint32(liveness)) {
 		s.unreachableSince = time.Now()
-		go s.checkUntilHealth(c)
+		reResolveInterval := 30 * time.Second
+		if val, err := util.EvalFailpoint("injectReResolveInterval"); err == nil {
+			if dur, err := time.ParseDuration(val.(string)); err == nil {
+				reResolveInterval = dur
+			}
+		}
+		go s.checkUntilHealth(c, liveness, reResolveInterval)
 	}
 }
 
-func (s *Store) checkUntilHealth(c *RegionCache) {
-	defer atomic.StoreUint32(&s.livenessState, uint32(reachable))
-
+func (s *Store) checkUntilHealth(c *RegionCache, liveness livenessState, reResolveInterval time.Duration) {
 	ticker := time.NewTicker(time.Second)
+	defer func() {
+		ticker.Stop()
+		if liveness != reachable {
+			logutil.BgLogger().Warn("[health check] store was still not reachable at the end of health check loop",
+				zap.Uint64("storeID", s.storeID),
+				zap.String("state", s.getResolveState().String()),
+				zap.String("liveness", s.getLivenessState().String()))
+		}
+	}()
 	lastCheckPDTime := time.Now()
 
 	for {
@@ -2615,26 +2637,37 @@ func (s *Store) checkUntilHealth(c *RegionCache) {
 		case <-c.ctx.Done():
 			return
 		case <-ticker.C:
-			if time.Since(lastCheckPDTime) > time.Second*30 {
+			if time.Since(lastCheckPDTime) > reResolveInterval {
 				lastCheckPDTime = time.Now()
 
 				valid, err := s.reResolve(c)
 				if err != nil {
 					logutil.BgLogger().Warn("[health check] failed to re-resolve unhealthy store", zap.Error(err))
 				} else if !valid {
-					logutil.BgLogger().Info("[health check] store meta deleted, stop checking", zap.Uint64("storeID", s.storeID), zap.String("addr", s.addr))
+					if s.getResolveState() == deleted {
+						// if the store is deleted, a new store with same id must be inserted (guaranteed by reResolve).
+						c.storeMu.RLock()
+						newStore := c.storeMu.stores[s.storeID]
+						c.storeMu.RUnlock()
+						logutil.BgLogger().Info("[health check] store meta changed",
+							zap.Uint64("storeID", s.storeID),
+							zap.String("oldAddr", s.addr),
+							zap.String("oldLabels", fmt.Sprintf("%v", s.labels)),
+							zap.String("newAddr", newStore.addr),
+							zap.String("newLabels", fmt.Sprintf("%v", newStore.labels)))
+						go newStore.checkUntilHealth(c, liveness, reResolveInterval)
+					}
 					return
 				}
 			}
 
 			bo := retry.NewNoopBackoff(c.ctx)
-			l := s.requestLiveness(bo, c)
-			if l == reachable {
+			liveness = s.requestLiveness(bo, c)
+			atomic.StoreUint32(&s.livenessState, uint32(liveness))
+			if liveness == reachable {
 				logutil.BgLogger().Info("[health check] store became reachable", zap.Uint64("storeID", s.storeID))
-
 				return
 			}
-			atomic.StoreUint32(&s.livenessState, uint32(l))
 		}
 	}
 }
@@ -2642,7 +2675,20 @@ func (s *Store) checkUntilHealth(c *RegionCache) {
 func (s *Store) requestLiveness(bo *retry.Backoffer, c *RegionCache) (l livenessState) {
 	// It's not convenient to mock liveness in integration tests. Use failpoint to achieve that instead.
 	if val, err := util.EvalFailpoint("injectLiveness"); err == nil {
-		switch val.(string) {
+		liveness := val.(string)
+		if strings.Contains(liveness, " ") {
+			for _, item := range strings.Split(liveness, " ") {
+				kv := strings.Split(item, ":")
+				if len(kv) != 2 {
+					continue
+				}
+				if kv[0] == s.addr {
+					liveness = kv[1]
+					break
+				}
+			}
+		}
+		switch liveness {
 		case "unreachable":
 			return unreachable
 		case "reachable":
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index 4618109..8bbce0e 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -39,12 +39,15 @@ import (
 	"context"
 	"fmt"
 	"math/rand"
+	"reflect"
 	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
+	"unsafe"
 
+	uatomic "go.uber.org/atomic"
 	"go.uber.org/zap"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -80,6 +83,30 @@ func LoadShuttingDown() uint32 {
 	return atomic.LoadUint32(&shuttingDown)
 }
 
+// ReplicaSelectorExperimentalOptions defines experimental options of replica selector.
+type ReplicaSelectorExperimentalOptions struct {
+	StaleRead struct {
+		PreventRetryFollower     bool
+		RetryLeaderTimeoutFactor float64
+	}
+}
+
+var selectorExpOpts uatomic.Pointer[ReplicaSelectorExperimentalOptions]
+
+// SetReplicaSelectorExperimentalOptions sets experimental options of replica selector.
+func SetReplicaSelectorExperimentalOptions(opts ReplicaSelectorExperimentalOptions) {
+	selectorExpOpts.Store(&opts)
+}
+
+// GetReplicaSelectorExperimentalOptions gets experimental options of replica selector.
+func GetReplicaSelectorExperimentalOptions() (opts ReplicaSelectorExperimentalOptions) {
+	ptr := selectorExpOpts.Load()
+	if ptr == nil {
+		return ReplicaSelectorExperimentalOptions{}
+	}
+	return *ptr
+}
+
 // RegionRequestSender sends KV/Cop requests to tikv server. It handles network
 // errors and some region errors internally.
 //
@@ -263,6 +290,8 @@ type replicaSelector struct {
 	targetIdx AccessIndex
 	// replicas[proxyIdx] is the store used to redirect requests this time
 	proxyIdx AccessIndex
+
+	ReplicaSelectorExperimentalOptions
 }
 
 // selectorState is the interface of states of the replicaSelector.
@@ -625,22 +654,28 @@ func (state *accessFollower) next(bo *retry.Backoffer, selector *replicaSelector
 		leader := selector.replicas[state.leaderIdx]
 		leaderEpochStale := leader.isEpochStale()
 		leaderUnreachable := leader.store.getLivenessState() != reachable
-		leaderInvalid := leaderEpochStale || leaderUnreachable || state.IsLeaderExhausted(leader)
-		if len(state.option.labels) > 0 {
-			logutil.BgLogger().Warn("unable to find stores with given labels",
+		leaderExhausted := state.IsLeaderExhausted(leader)
+		leaderTimeout := leader.deadlineErrUsingConfTimeout
+		if len(state.option.labels) > 0 && !state.option.leaderOnly {
+			logutil.BgLogger().Warn("unable to find a store with given labels",
 				zap.Uint64("region", selector.region.GetID()),
-				zap.Bool("leader-epoch-stale", leaderEpochStale),
-				zap.Bool("leader-unreachable", leaderUnreachable),
-				zap.Bool("leader-invalid", leaderInvalid),
 				zap.Bool("stale-read", state.isStaleRead),
 				zap.Any("labels", state.option.labels))
 		}
-		if leaderInvalid || leader.deadlineErrUsingConfTimeout {
+		if leaderEpochStale || leaderUnreachable || leaderExhausted || leaderTimeout {
+			logutil.BgLogger().Warn("unable to find a valid leader",
+				zap.Uint64("region", selector.region.GetID()),
+				zap.Bool("stale-read", state.isStaleRead),
+				zap.Bool("epoch-stale", leaderEpochStale),
+				zap.Bool("unreachable", leaderUnreachable),
+				zap.Bool("exhausted", leaderExhausted),
+				zap.Bool("timeout", leaderTimeout))
 			// In stale-read, the request will fallback to leader after the local follower failure.
 			// If the leader is also unavailable, we can fallback to the follower and use replica-read flag again,
 			// The remote follower not tried yet, and the local follower can retry without stale-read flag.
 			// If leader tried and received deadline exceeded error, try follower.
-			if state.isStaleRead || leader.deadlineErrUsingConfTimeout {
+			if (state.isStaleRead && !selector.StaleRead.PreventRetryFollower) ||
+				(!state.isStaleRead && leader.deadlineErrUsingConfTimeout) {
 				selector.state = &tryFollower{
 					leaderIdx: state.leaderIdx,
 					lastIdx:   state.leaderIdx,
@@ -665,6 +700,7 @@ func (state *accessFollower) next(bo *retry.Backoffer, selector *replicaSelector
 	if resetStaleRead {
 		staleRead := false
 		rpcCtx.contextPatcher.staleRead = &staleRead
+		rpcCtx.contextPatcher.timeoutFactor = selector.StaleRead.RetryLeaderTimeoutFactor
 	}
 	return rpcCtx, nil
 }
@@ -770,6 +806,7 @@ func newReplicaSelector(regionCache *RegionCache, regionID RegionVerID, req *tik
 		state,
 		-1,
 		-1,
+		GetReplicaSelectorExperimentalOptions(),
 	}, nil
 }
 
@@ -808,6 +845,17 @@ func (s *replicaSelector) proxyReplica() *replica {
 	return nil
 }
 
+// sliceIdentical checks whether two slices are referencing the same block of memory. Two `nil`s are also considered
+// the same.
+func sliceIdentical[T any](a, b []T) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	aHeader := (*reflect.SliceHeader)(unsafe.Pointer(&a))
+	bHeader := (*reflect.SliceHeader)(unsafe.Pointer(&b))
+	return aHeader.Data == bHeader.Data
+}
+
 func (s *replicaSelector) refreshRegionStore() {
 	oldRegionStore := s.regionStore
 	newRegionStore := s.region.getStore()
@@ -820,7 +868,7 @@ func (s *replicaSelector) refreshRegionStore() {
 	// So we just compare the address here.
 	// When stores change, we mark this replicaSelector as invalid to let the caller
 	// recreate a new replicaSelector.
-	if &oldRegionStore.stores != &newRegionStore.stores {
+	if !sliceIdentical(oldRegionStore.stores, newRegionStore.stores) {
 		s.state = &invalidStore{}
 		return
 	}
@@ -1197,7 +1245,7 @@ func (s *RegionRequestSender) SendReqCtx(
 			}
 		}
 
-		rpcCtx.contextPatcher.applyTo(&req.Context)
+		timeout = rpcCtx.contextPatcher.applyTo(&req.Context, timeout)
 		if req.InputRequestSource != "" && s.replicaSelector != nil {
 			s.replicaSelector.patchRequestSource(req, rpcCtx)
 		}
diff --git a/tikv/region.go b/tikv/region.go
index 4fd0a25..8b95765 100644
--- a/tikv/region.go
+++ b/tikv/region.go
@@ -79,6 +79,9 @@ type RPCCancellerCtxKey = locate.RPCCancellerCtxKey
 // simply return the error to caller.
 type RegionRequestSender = locate.RegionRequestSender
 
+// ReplicaSelectorExperimentalOptions defines experimental options of replica selector.
+type ReplicaSelectorExperimentalOptions = locate.ReplicaSelectorExperimentalOptions
+
 // StoreSelectorOption configures storeSelectorOp.
 type StoreSelectorOption = locate.StoreSelectorOption
 
@@ -152,6 +155,16 @@ func StoreShuttingDown(v uint32) {
 	locate.StoreShuttingDown(v)
 }
 
+// SetReplicaSelectorExperimentalOptions sets experimental options of replica selector.
+func SetReplicaSelectorExperimentalOptions(opts ReplicaSelectorExperimentalOptions) {
+	locate.SetReplicaSelectorExperimentalOptions(opts)
+}
+
+// GetReplicaSelectorExperimentalOptions gets experimental options of replica selector.
+func GetReplicaSelectorExperimentalOptions() (opts ReplicaSelectorExperimentalOptions) {
+	return locate.GetReplicaSelectorExperimentalOptions()
+}
+
 // WithMatchLabels indicates selecting stores with matched labels
 func WithMatchLabels(labels []*metapb.StoreLabel) StoreSelectorOption {
 	return locate.WithMatchLabels(labels)


Test patch:
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index caefb2a..9d79619 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -1665,3 +1665,46 @@ func (s *testRegionCacheSuite) TestBackgroundCacheGC() {
 	}, 3*time.Second, 200*time.Millisecond)
 	s.checkCache(remaining)
 }
+
+func (s *testRegionCacheSuite) TestHealthCheckWithStoreReplace() {
+	// init region cache
+	s.cache.LocateKey(s.bo, []byte("a"))
+
+	store1 := s.cache.getStoreByStoreID(s.store1)
+	s.Require().NotNil(store1)
+	s.Require().Equal(resolved, store1.getResolveState())
+
+	// setup mock liveness func
+	store1Liveness := uint32(unreachable)
+	tf := func(s *Store, bo *retry.Backoffer) livenessState {
+		if s.storeID == store1.storeID {
+			return livenessState(atomic.LoadUint32(&store1Liveness))
+		}
+		return reachable
+	}
+	s.cache.testingKnobs.mockRequestLiveness.Store((*livenessFunc)(&tf))
+
+	// start health check loop
+	atomic.StoreUint32(&store1.livenessState, store1Liveness)
+	go store1.checkUntilHealth(s.cache, livenessState(store1Liveness), time.Second)
+
+	// update store meta
+	s.cluster.UpdateStoreAddr(store1.storeID, store1.addr+"'", store1.labels...)
+
+	// assert that the old store should be deleted and it's not reachable
+	s.Eventually(func() bool {
+		return store1.getResolveState() == deleted && store1.getLivenessState() != reachable
+	}, 3*time.Second, time.Second)
+
+	// assert that the new store should be added and it's also not reachable
+	newStore1 := s.cache.getStoreByStoreID(s.store1)
+	s.Require().NotEqual(reachable, newStore1.getLivenessState())
+
+	// recover store1
+	atomic.StoreUint32(&store1Liveness, uint32(reachable))
+
+	// assert that the new store should be reachable
+	s.Eventually(func() bool {
+		return newStore1.getResolveState() == resolved && newStore1.getLivenessState() == reachable
+	}, 3*time.Second, time.Second)
+}
diff --git a/internal/locate/region_request3_test.go b/internal/locate/region_request3_test.go
index 9d6558b..c674f12 100644
--- a/internal/locate/region_request3_test.go
+++ b/internal/locate/region_request3_test.go
@@ -36,6 +36,7 @@ package locate
 
 import (
 	"context"
+	"fmt"
 	"strconv"
 	"sync/atomic"
 	"testing"
@@ -135,6 +136,82 @@ func (s *testRegionRequestToThreeStoresSuite) TestSwitchPeerWhenNoLeader() {
 	resp, err := s.regionRequestSender.SendReq(bo, req, loc.Region, time.Second)
 	s.Nil(err)
 	s.NotNil(resp)
+	s.Nil(resp.GetRegionError())
+}
+
+func (s *testRegionRequestToThreeStoresSuite) TestSwitchPeerWhenNoLeaderErrorWithNewLeaderInfo() {
+	cnt := 0
+	var location *KeyLocation
+	cli := &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+		cnt++
+		switch cnt {
+		case 1:
+			region := s.cache.GetCachedRegionWithRLock(location.Region)
+			s.NotNil(region)
+			leaderPeerIdx := int(region.getStore().workTiKVIdx)
+			peers := region.meta.Peers
+			// return no leader with new leader info
+			response = &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{
+				RegionError: &errorpb.Error{NotLeader: &errorpb.NotLeader{
+					RegionId: req.RegionId,
+					Leader:   peers[(leaderPeerIdx+1)%len(peers)],
+				}},
+			}}
+		case 2:
+			response = &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{
+				Value: []byte("a"),
+			}}
+		default:
+			return nil, fmt.Errorf("unexpected request")
+		}
+		return response, err
+	}}
+
+	req := tikvrpc.NewRequest(tikvrpc.CmdGet, &kvrpcpb.GetRequest{Key: []byte("a")}, kvrpcpb.Context{})
+	req.ReplicaReadType = kv.ReplicaReadLeader
+	var err error
+	location, err = s.cache.LocateKey(s.bo, []byte("a"))
+	s.Nil(err)
+	s.NotNil(location)
+	bo := retry.NewBackoffer(context.Background(), 1000)
+	resp, _, err := NewRegionRequestSender(s.cache, cli).SendReqCtx(bo, req, location.Region, time.Second, tikvrpc.TiKV)
+	s.Nil(err)
+	s.NotNil(resp)
+	regionErr, err := resp.GetRegionError()
+	s.Nil(err)
+	s.Nil(regionErr)
+	// It's unreasoneable to retry in upper layer, such as cop request, the upper layer will need to rebuild cop request and retry, there are some unnecessary overhead.
+	s.Equal(cnt, 2)
+	r := s.cache.GetCachedRegionWithRLock(location.Region)
+	s.True(r.isValid())
+}
+
+func (s *testRegionRequestToThreeStoresSuite) TestSliceIdentical() {
+	a := make([]int, 0)
+	b := a
+	s.True(sliceIdentical(a, b))
+	// This is not guaranteed.
+	// b = make([]int, 0)
+	// s.False(sliceIdentical(a, b))
+
+	a = append(a, 1, 2, 3)
+	b = a
+	s.True(sliceIdentical(a, b))
+	b = a[:2]
+	s.False(sliceIdentical(a, b))
+	b = a[1:]
+	s.False(sliceIdentical(a, b))
+	a = a[1:]
+	s.True(sliceIdentical(a, b))
+
+	a = nil
+	b = nil
+
+	s.True(sliceIdentical(a, b))
+	a = make([]int, 0)
+	s.False(sliceIdentical(a, b))
+	a = append(a, 1)
+	s.False(sliceIdentical(a, b))
 }
 
 func (s *testRegionRequestToThreeStoresSuite) loadAndGetLeaderStore() (*Store, string) {
@@ -981,7 +1058,7 @@ func (s *testRegionRequestToThreeStoresSuite) TestAccessFollowerAfter1TiKVDown()
 	regionStore := region.getStore()
 	leaderAddr = regionStore.stores[regionStore.workTiKVIdx].addr
 	s.NotEqual(leaderAddr, "")
-	for i := 0; i < 10; i++ {
+	for i := 0; i < 30; i++ {
 		bo := retry.NewBackofferWithVars(context.Background(), 100, nil)
 		resp, _, err := s.regionRequestSender.SendReqCtx(bo, req, loc.Region, client.ReadTimeoutShort, tikvrpc.TiKV)
 		s.Nil(err)
@@ -1442,3 +1519,108 @@ func (s *testRegionRequestToThreeStoresSuite) TestRetryRequestSource() {
 		}
 	}
 }
+
+func (s *testRegionRequestToThreeStoresSuite) TestReplicaSelectorExperimentalOptions() {
+	key := []byte("key")
+	bo := retry.NewBackoffer(context.Background(), -1)
+
+	loc, err := s.cache.LocateKey(bo, key)
+	s.Require().NoError(err)
+
+	region := s.cache.GetCachedRegionWithRLock(loc.Region)
+	leader, _, _, _ := region.WorkStorePeer(region.getStore())
+	follower, _, _, _ := region.FollowerStorePeer(region.getStore(), 0, &storeSelectorOp{})
+
+	newStaleReadReq := func() *tikvrpc.Request {
+		req := tikvrpc.NewReplicaReadRequest(tikvrpc.CmdGet, &kvrpcpb.GetRequest{Key: key}, kv.ReplicaReadMixed, nil)
+		req.ReadReplicaScope = oracle.GlobalTxnScope
+		req.TxnScope = oracle.GlobalTxnScope
+		req.EnableStaleRead()
+		return req
+	}
+
+	errRespTimeout := &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{RegionError: &errorpb.Error{ServerIsBusy: &errorpb.ServerIsBusy{Reason: "mock: deadline is exceeded"}}}}
+	errRespDataNotReady := &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{RegionError: &errorpb.Error{DataIsNotReady: &errorpb.DataIsNotReady{}}}}
+
+	mockLeaderTimeout := func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+		if addr == leader.addr {
+			return errRespTimeout, nil
+		}
+		return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{Value: []byte(addr)}}, nil
+	}
+
+	mockLeaderSlow := func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+		if addr == leader.addr && timeout <= 300*time.Millisecond {
+			return errRespTimeout, nil
+		}
+		if addr == follower.addr {
+			if req.StaleRead {
+				return errRespDataNotReady, nil
+			} else {
+				return errRespTimeout, nil
+			}
+		}
+		return &tikvrpc.Response{Resp: &kvrpcpb.GetResponse{Value: []byte(addr)}}, nil
+	}
+
+	s.Run("PreventRetryFollowerOn", func() {
+		// 1. access local leader -> timeout
+		// 2. cannot retry leader (since timeout), cannot fallback to tryFollower -> pseudo error
+		defer SetReplicaSelectorExperimentalOptions(GetReplicaSelectorExperimentalOptions())
+		var opts ReplicaSelectorExperimentalOptions
+		opts.StaleRead.PreventRetryFollower = true
+		SetReplicaSelectorExperimentalOptions(opts)
+
+		s.regionRequestSender.client = &fnClient{fn: mockLeaderTimeout}
+		s.cache.LocateKey(bo, key)
+
+		resp, _, err := s.regionRequestSender.SendReqCtx(bo, newStaleReadReq(), loc.Region, time.Second, tikvrpc.TiKV, WithMatchLabels(leader.labels))
+		s.Require().NoError(err)
+		regionErr, err := resp.GetRegionError()
+		s.Require().NoError(err)
+		s.Require().NotNil(regionErr)
+		s.Require().True(IsFakeRegionError(regionErr))
+	})
+
+	s.Run("PreventRetryFollowerOff", func() {
+		// 1. access local leader -> timeout
+		// 2. cannot retry leader (since timeout), fallback to tryFollower, access follower -> ok
+		s.regionRequestSender.client = &fnClient{fn: mockLeaderTimeout}
+		s.cache.LocateKey(bo, key)
+
+		resp, rpcCtx, err := s.regionRequestSender.SendReqCtx(bo, newStaleReadReq(), loc.Region, time.Second, tikvrpc.TiKV, WithMatchLabels(leader.labels))
+		s.Require().NoError(err)
+		s.Require().Equal(rpcCtx.Addr, string(resp.Resp.(*kvrpcpb.GetResponse).Value))
+	})
+
+	s.Run("RetryLeaderTimeoutFactorOn", func() {
+		// 1. access local follower -> data is not ready
+		// 2. retry leader with timeout*2 -> ok
+		defer SetReplicaSelectorExperimentalOptions(GetReplicaSelectorExperimentalOptions())
+		var opts ReplicaSelectorExperimentalOptions
+		opts.StaleRead.RetryLeaderTimeoutFactor = 2
+		SetReplicaSelectorExperimentalOptions(opts)
+
+		s.regionRequestSender.client = &fnClient{fn: mockLeaderSlow}
+		s.cache.LocateKey(bo, key)
+
+		resp, _, err := s.regionRequestSender.SendReqCtx(bo, newStaleReadReq(), loc.Region, 200*time.Millisecond, tikvrpc.TiKV, WithMatchLabels(follower.labels))
+		s.Require().NoError(err)
+		s.Require().Equal(leader.addr, string(resp.Resp.(*kvrpcpb.GetResponse).Value))
+	})
+
+	s.Run("RetryLeaderTimeoutFactorOff", func() {
+		// 1. access local follower -> data is not ready
+		// 2. retry leader with timeout -> timeout
+		// 3. fallback to tryFollower, access local follower with replica read -> timeout
+		// 4. access the other follower -> ok
+		s.regionRequestSender.client = &fnClient{fn: mockLeaderSlow}
+		s.cache.LocateKey(bo, key)
+
+		resp, rpcCtx, err := s.regionRequestSender.SendReqCtx(bo, newStaleReadReq(), loc.Region, 200*time.Millisecond, tikvrpc.TiKV, WithMatchLabels(follower.labels))
+		s.Require().NoError(err)
+		s.Require().Equal(rpcCtx.Addr, string(resp.Resp.(*kvrpcpb.GetResponse).Value))
+		s.Require().NotEqual(leader.addr, string(resp.Resp.(*kvrpcpb.GetResponse).Value))
+		s.Require().NotEqual(follower.addr, string(resp.Resp.(*kvrpcpb.GetResponse).Value))
+	})
+}
diff --git a/internal/retry/backoff_test.go b/internal/retry/backoff_test.go
index 1668c5d..895c77e 100644
--- a/internal/retry/backoff_test.go
+++ b/internal/retry/backoff_test.go
@@ -72,7 +72,8 @@ func TestBackoffErrorType(t *testing.T) {
 		err = b.Backoff(BoTxnNotFound, errors.New("txn not found"))
 		if err != nil {
 			// Next backoff should return error of backoff that sleeps for longest time.
-			assert.ErrorIs(t, err, BoTxnNotFound.err)
+			cfg, _ := b.longestSleepCfg()
+			assert.ErrorIs(t, err, cfg.err)
 			return
 		}
 	}
diff --git a/rawkv/rawkv_test.go b/rawkv/rawkv_test.go
index f5eefaa..e9e61db 100644
--- a/rawkv/rawkv_test.go
+++ b/rawkv/rawkv_test.go
@@ -41,14 +41,17 @@ import (
 	"hash/crc64"
 	"testing"
 
+	"github.com/pingcap/failpoint"
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/client-go/v2/internal/locate"
 	"github.com/tikv/client-go/v2/internal/mockstore/mocktikv"
 	"github.com/tikv/client-go/v2/internal/retry"
 	"github.com/tikv/client-go/v2/kv"
+	"github.com/tikv/client-go/v2/tikv"
 )
 
 func TestRawKV(t *testing.T) {
+	tikv.EnableFailpoints()
 	suite.Run(t, new(testRawkvSuite))
 }
 
@@ -77,9 +80,11 @@ func (s *testRawkvSuite) SetupTest() {
 	s.peer1 = peerIDs[0]
 	s.peer2 = peerIDs[1]
 	s.bo = retry.NewBackofferWithVars(context.Background(), 5000, nil)
+	s.Nil(failpoint.Enable("tikvclient/injectReResolveInterval", `return("1s")`))
 }
 
 func (s *testRawkvSuite) TearDownTest() {
+	s.Nil(failpoint.Disable("tikvclient/injectReResolveInterval"))
 	s.mvccStore.Close()
 }
 
@@ -110,6 +115,9 @@ func (s *testRawkvSuite) TestReplaceAddrWithNewStore() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	getVal, err := client.Get(context.Background(), testKey)
 
 	s.Nil(err)
@@ -172,6 +180,9 @@ func (s *testRawkvSuite) TestReplaceNewAddrAndOldOfflineImmediately() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	getVal, err := client.Get(context.Background(), testKey)
 	s.Nil(err)
 	s.Equal(getVal, testValue)
@@ -200,6 +211,9 @@ func (s *testRawkvSuite) TestReplaceStore() {
 	s.cluster.RemovePeer(s.region1, s.peer1)
 	s.cluster.ChangeLeader(s.region1, peer3)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	err = client.Put(context.Background(), testKey, testValue)
 	s.Nil(err)
 }
@@ -234,6 +248,9 @@ func (s *testRawkvSuite) TestColumnFamilyForClient() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	// test get
 	client.SetColumnFamily(cf1)
 	getVal, err := client.Get(context.Background(), testKeyCf1)
@@ -303,6 +320,9 @@ func (s *testRawkvSuite) TestColumnFamilyForOptions() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	// test get
 	getVal, err := client.Get(context.Background(), keyInCf1, SetColumnFamily(cf1))
 	s.Nil(err)
@@ -370,6 +390,9 @@ func (s *testRawkvSuite) TestBatch() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	// test BatchGet
 	returnValues, err := client.BatchGet(context.Background(), keys, SetColumnFamily(cf))
 	s.Nil(err)
@@ -426,6 +449,9 @@ func (s *testRawkvSuite) TestScan() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	// test scan
 	startKey, endKey := []byte("key1"), []byte("keyz")
 	limit := 3
@@ -498,6 +524,9 @@ func (s *testRawkvSuite) TestDeleteRange() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	// test DeleteRange
 	startKey, endKey := []byte("key3"), []byte(nil)
 	err = client.DeleteRange(context.Background(), startKey, endKey, SetColumnFamily(cf))
@@ -535,6 +564,9 @@ func (s *testRawkvSuite) TestCompareAndSwap() {
 	s.cluster.ChangeLeader(s.region1, s.peer2)
 	s.cluster.RemovePeer(s.region1, s.peer1)
 
+	s.Nil(failpoint.Enable("tikvclient/injectLiveness", `return("store1:reachable store2:unreachable")`))
+	defer failpoint.Disable("tikvclient/injectLiveness")
+
 	// test CompareAndSwap for false atomic
 	_, _, err = client.CompareAndSwap(
 		context.Background(),