tikv: avoid switch peer when batchRequest be cancelled (#10822) (#10850)

Author: lysu

store/tikv/client.go

@@ -35,10 +35,8 @@ import (
 	"github.com/pingcap/tidb/util/logutil"
 	"go.uber.org/zap"
 	"google.golang.org/grpc"
-	gcodes "google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/keepalive"
-	gstatus "google.golang.org/grpc/status"
 )
 
 // MaxSendMsgSize set max gRPC request message size sent to server. If any request message size is larger than
@@ -618,8 +616,9 @@ func sendBatchRequest(
 	select {
 	case connArray.batchCommandsCh <- entry:
 	case <-ctx1.Done():
-		logutil.Logger(context.Background()).Warn("send request is timeout", zap.String("to", addr))
-		return nil, errors.Trace(gstatus.Error(gcodes.DeadlineExceeded, "Canceled or timeout"))
+		logutil.Logger(context.Background()).Warn("send request is cancelled",
+			zap.String("to", addr), zap.String("cause", ctx1.Err().Error()))
+		return nil, errors.Trace(ctx1.Err())
 	}
 
 	select {
@@ -630,8 +629,9 @@ func sendBatchRequest(
 		return tikvrpc.FromBatchCommandsResponse(res), nil
 	case <-ctx1.Done():
 		atomic.StoreInt32(&entry.canceled, 1)
-		logutil.Logger(context.Background()).Warn("send request is canceled", zap.String("to", addr))
-		return nil, errors.Trace(gstatus.Error(gcodes.DeadlineExceeded, "Canceled or timeout"))
+		logutil.Logger(context.Background()).Warn("wait response is cancelled",
+			zap.String("to", addr), zap.String("cause", ctx1.Err().Error()))
+		return nil, errors.Trace(ctx1.Err())
 	}
 }
 

store/tikv/client_test.go

@@ -14,9 +14,12 @@
 package tikv
 
 import (
+	"context"
 	"testing"
+	"time"
 
 	. "github.com/pingcap/check"
+	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/tikvpb"
 	"github.com/pingcap/tidb/config"
 )
@@ -77,3 +80,16 @@ func (s *testClientSuite) TestRemoveCanceledRequests(c *C) {
 	newEntryPtr := &entries[0]
 	c.Assert(entryPtr, Equals, newEntryPtr)
 }
+
+func (s *testClientSuite) TestCancelTimeoutRetErr(c *C) {
+	req := new(tikvpb.BatchCommandsRequest_Request)
+	a := &connArray{batchCommandsCh: make(chan *batchCommandsEntry, 1)}
+
+	ctx, cancel := context.WithCancel(context.TODO())
+	cancel()
+	_, err := sendBatchRequest(ctx, "", a, req, 2*time.Second)
+	c.Assert(errors.Cause(err), Equals, context.Canceled)
+
+	_, err = sendBatchRequest(context.Background(), "", a, req, 0)
+	c.Assert(errors.Cause(err), Equals, context.DeadlineExceeded)
+}