From 39084386b3246ab9e478879dc557a6fbf8da05f6 Mon Sep 17 00:00:00 2001
From: crazycs <crazycs520@gmail.com>
Date: Mon, 18 Sep 2023 11:17:51 +0800
Subject: [PATCH] *: configurable kv_read_timeout should not affect write
 request (#978)

* *: configurable kv_read_timeout should not affect write request

Signed-off-by: crazycs520 <crazycs520@gmail.com>

* refine test

Signed-off-by: crazycs520 <crazycs520@gmail.com>

* refine logic

Signed-off-by: crazycs520 <crazycs520@gmail.com>

---------

Signed-off-by: crazycs520 <crazycs520@gmail.com>
---
 internal/locate/region_request.go       | 38 ++++++++++++++++---------
 internal/locate/region_request3_test.go | 16 +++++++++++
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index f260ad18..9d823e4f 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -1085,13 +1085,25 @@ func (s *replicaSelector) onSendFailure(bo *retry.Backoffer, err error) {
 	s.state.onSendFailure(bo, s, err)
 }
 
-func (s *replicaSelector) onDeadlineExceeded() {
-	if target := s.targetReplica(); target != nil {
-		target.deadlineErrUsingConfTimeout = true
+func (s *replicaSelector) onReadReqConfigurableTimeout(req *tikvrpc.Request) bool {
+	if req.MaxExecutionDurationMs >= uint64(client.ReadTimeoutShort.Milliseconds()) {
+		// Configurable timeout should less than `ReadTimeoutShort`.
+		return false
 	}
-	if accessLeader, ok := s.state.(*accessKnownLeader); ok {
-		// If leader return deadline exceeded error, we should try to access follower next time.
-		s.state = &tryFollower{leaderIdx: accessLeader.leaderIdx, lastIdx: accessLeader.leaderIdx}
+	switch req.Type {
+	case tikvrpc.CmdGet, tikvrpc.CmdBatchGet, tikvrpc.CmdScan,
+		tikvrpc.CmdCop, tikvrpc.CmdBatchCop, tikvrpc.CmdCopStream:
+		if target := s.targetReplica(); target != nil {
+			target.deadlineErrUsingConfTimeout = true
+		}
+		if accessLeader, ok := s.state.(*accessKnownLeader); ok {
+			// If leader return deadline exceeded error, we should try to access follower next time.
+			s.state = &tryFollower{leaderIdx: accessLeader.leaderIdx, lastIdx: accessLeader.leaderIdx}
+		}
+		return true
+	default:
+		// Only work for read requests, return false for non-read requests.
+		return false
 	}
 }
 
@@ -1776,9 +1788,8 @@ func (s *RegionRequestSender) onSendFail(bo *retry.Backoffer, ctx *RPCContext, r
 		return errors.WithStack(err)
 	} else if LoadShuttingDown() > 0 {
 		return errors.WithStack(tikverr.ErrTiDBShuttingDown)
-	} else if isCauseByDeadlineExceeded(err) && req.MaxExecutionDurationMs < uint64(client.ReadTimeoutShort.Milliseconds()) {
-		if s.replicaSelector != nil {
-			s.replicaSelector.onDeadlineExceeded()
+	} else if isCauseByDeadlineExceeded(err) {
+		if s.replicaSelector != nil && s.replicaSelector.onReadReqConfigurableTimeout(req) {
 			return nil
 		}
 	}
@@ -2097,8 +2108,9 @@ func (s *RegionRequestSender) onRegionError(
 
 	if serverIsBusy := regionErr.GetServerIsBusy(); serverIsBusy != nil {
 		if s.replicaSelector != nil && strings.Contains(serverIsBusy.GetReason(), "deadline is exceeded") {
-			s.replicaSelector.onDeadlineExceeded()
-			return true, nil
+			if s.replicaSelector.onReadReqConfigurableTimeout(req) {
+				return true, nil
+			}
 		}
 		if s.replicaSelector != nil {
 			return s.replicaSelector.onServerIsBusy(bo, ctx, req, serverIsBusy)
@@ -2228,8 +2240,8 @@ func (s *RegionRequestSender) onRegionError(
 		return true, nil
 	}
 
-	if isDeadlineExceeded(regionErr) && s.replicaSelector != nil {
-		s.replicaSelector.onDeadlineExceeded()
+	if isDeadlineExceeded(regionErr) && s.replicaSelector != nil && s.replicaSelector.onReadReqConfigurableTimeout(req) {
+		return true, nil
 	}
 
 	if mismatch := regionErr.GetMismatchPeerId(); mismatch != nil {
diff --git a/internal/locate/region_request3_test.go b/internal/locate/region_request3_test.go
index a079e2c6..41dd69b6 100644
--- a/internal/locate/region_request3_test.go
+++ b/internal/locate/region_request3_test.go
@@ -1278,6 +1278,22 @@ func (s *testRegionRequestToThreeStoresSuite) TestSendReqFirstTimeout() {
 			s.Equal(0, bo.GetTotalBackoffTimes())                                // no backoff since fast retry.
 		}
 	}
+
+	// Test for write request.
+	tf := func(s *Store, bo *retry.Backoffer) livenessState {
+		return reachable
+	}
+	s.regionRequestSender.regionCache.testingKnobs.mockRequestLiveness.Store((*livenessFunc)(&tf))
+	resetStats()
+	req := tikvrpc.NewRequest(tikvrpc.CmdPrewrite, &kvrpcpb.PrewriteRequest{}, kvrpcpb.Context{})
+	req.ReplicaReadType = kv.ReplicaReadLeader
+	loc := getLocFn()
+	bo = retry.NewBackoffer(context.Background(), 1000)
+	resp, _, _, err := s.regionRequestSender.SendReqCtx(bo, req, loc.Region, time.Millisecond, tikvrpc.TiKV)
+	s.Nil(resp)
+	s.Equal(context.DeadlineExceeded, err)
+	backoffTimes := bo.GetBackoffTimes()
+	s.True(backoffTimes["tikvRPC"] > 0) // write request timeout won't do fast retry, so backoff times should be more than 0.
 }
 
 func (s *testRegionRequestToThreeStoresSuite) TestStaleReadFallback2Follower() {