Repository: tikv/client-go (go)

Issue:
batch client batchSendLoop panic
TiDB Version: v6.5.5, and nightly version also has this problem too.

Panic stack log:
```
[2023/10/13 14:00:33.948 +08:00] [ERROR] [client_batch.go:303] [batchSendLoop] [r={}] [stack="[github.com/tikv/client-go/v2/internal/client.(*batchConn).batchSendLoop.func1](http://github.com/tikv/client-go/v2/internal/client.(*batchConn).batchSendLoop.func1)
	/go/pkg/mod/[github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:305](http://github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:305)
runtime.gopanic
	/usr/local/go/src/runtime/panic.go:884
runtime.goPanicIndex
	/usr/local/go/src/runtime/panic.go:113
[github.com/pingcap/kvproto/pkg/tikvpb.encodeVarintTikvpb](http://github.com/pingcap/kvproto/pkg/tikvpb.encodeVarintTikvpb)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:5438](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:5438)
[github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request_Coprocessor).MarshalToSizedBuffer](http://github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request_Coprocessor).MarshalToSizedBuffer)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:4325](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:4325)
[github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request_Coprocessor).MarshalTo](http://github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request_Coprocessor).MarshalTo)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:4313](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:4313)
[github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request).MarshalToSizedBuffer](http://github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request).MarshalToSizedBuffer)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:3850](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:3850)
[github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest).MarshalToSizedBuffer](http://github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest).MarshalToSizedBuffer)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:3808](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:3808)
[github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest).Marshal](http://github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest).Marshal)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:3766](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:3766)
[google.golang.org/protobuf/internal/impl.legacyMarshal](http://google.golang.org/protobuf/internal/impl.legacyMarshal)
	/go/pkg/mod/[google.golang.org/protobuf@v1.28.1/internal/impl/legacy_message.go:402](http://google.golang.org/protobuf@v1.28.1/internal/impl/legacy_message.go:402)
[google.golang.org/protobuf/proto.MarshalOptions.marshal](http://google.golang.org/protobuf/proto.MarshalOptions.marshal)
	/go/pkg/mod/[google.golang.org/protobuf@v1.28.1/proto/encode.go:166](http://google.golang.org/protobuf@v1.28.1/proto/encode.go:166)
[google.golang.org/protobuf/proto.MarshalOptions.MarshalAppend](http://google.golang.org/protobuf/proto.MarshalOptions.MarshalAppend)
	/go/pkg/mod/[google.golang.org/protobuf@v1.28.1/proto/encode.go:125](http://google.golang.org/protobuf@v1.28.1/proto/encode.go:125)
[github.com/golang/protobuf/proto.marshalAppend](http://github.com/golang/protobuf/proto.marshalAppend)
	/go/pkg/mod/[github.com/golang/protobuf@v1.5.2/proto/wire.go:40](http://github.com/golang/protobuf@v1.5.2/proto/wire.go:40)
[github.com/golang/protobuf/proto.Marshal](http://github.com/golang/protobuf/proto.Marshal)
	/go/pkg/mod/[github.com/golang/protobuf@v1.5.2/proto/wire.go:23](http://github.com/golang/protobuf@v1.5.2/proto/wire.go:23)
[google.golang.org/grpc/encoding/proto.codec.Marshal](http://google.golang.org/grpc/encoding/proto.codec.Marshal)
	/go/pkg/mod/[google.golang.org/grpc@v1.51.0/encoding/proto/proto.go:45](http://google.golang.org/grpc@v1.51.0/encoding/proto/proto.go:45)
[google.golang.org/grpc.encode](http://google.golang.org/grpc.encode)
	/go/pkg/mod/[google.golang.org/grpc@v1.51.0/rpc_util.go:595](http://google.golang.org/grpc@v1.51.0/rpc_util.go:595)
[google.golang.org/grpc.prepareMsg](http://google.golang.org/grpc.prepareMsg)
	/go/pkg/mod/[google.golang.org/grpc@v1.51.0/stream.go:1708](http://google.golang.org/grpc@v1.51.0/stream.go:1708)
[google.golang.org/grpc.(*clientStream).SendMsg](http://google.golang.org/grpc.(*clientStream).SendMsg)
	/go/pkg/mod/[google.golang.org/grpc@v1.51.0/stream.go:846](http://google.golang.org/grpc@v1.51.0/stream.go:846)
[github.com/pingcap/kvproto/pkg/tikvpb.(*tikvBatchCommandsClient).Send](http://github.com/pingcap/kvproto/pkg/tikvpb.(*tikvBatchCommandsClient).Send)
	/go/pkg/mod/[github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:2068](http://github.com/pingcap/kvproto@v0.0.0-20230726063044-73d6d7f3756b/pkg/tikvpb/tikvpb.pb.go:2068)
[github.com/tikv/client-go/v2/internal/client.(*batchCommandsClient).send](http://github.com/tikv/client-go/v2/internal/client.(*batchCommandsClient).send)
	/go/pkg/mod/[github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:519](http://github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:519)
[github.com/tikv/client-go/v2/internal/client.(*batchConn).getClientAndSend](http://github.com/tikv/client-go/v2/internal/client.(*batchConn).getClientAndSend)
	/go/pkg/mod/[github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:381](http://github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:381)
[github.com/tikv/client-go/v2/internal/client.(*batchConn).batchSendLoop](http://github.com/tikv/client-go/v2/internal/client.(*batchConn).batchSendLoop)
	/go/pkg/mod/[github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:344](http://github.com/tikv/client-go/v2@v2.0.4-0.20230912041415-9c163cc8574b/internal/client/client_batch.go:344)"]
```

Pull request description:
Fix batch client batchSendLoop panic
close #1020

panic is cause by data race:

<img width="912" alt="image" src="https://github.com/tikv/client-go/assets/26020263/b043eb26-adc2-4435-b256-53c7c7881e47">

```go
==================
WARNING: DATA RACE
Write at 0x00c0006aabfc by goroutine 1072:
  github.com/tikv/client-go/v2/internal/locate.(*RegionRequestSender).SendReqCtx()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/locate/region_request.go:1462 +0x1ae4
  github.com/tikv/client-go/v2/internal/locate.(*RegionRequestSender).SendReq()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/locate/region_request.go:243 +0x110
  github.com/tikv/client-go/v2/internal/locate.(*testRegionRequestToSingleStoreSuite).TestBatchClientSendLoopPanic.func5()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/locate/region_request_test.go:779 +0xbc

Previous read at 0x00c0006aabfc by goroutine 1267:
  github.com/pingcap/kvproto/pkg/kvrpcpb.(*Context).Size()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/kvrpcpb/kvrpcpb.pb.go:21210 +0x840
  github.com/pingcap/kvproto/pkg/coprocessor.(*Request).Size()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/coprocessor/coprocessor.pb.go:1657 +0x54
  github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request_Coprocessor).Size()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/tikvpb/tikvpb.pb.go:6157 +0x50
  github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest_Request).MarshalToSizedBuffer()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/tikvpb/tikvpb.pb.go:4261 +0x110
  github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest).MarshalToSizedBuffer()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/tikvpb/tikvpb.pb.go:4221 +0x230
  github.com/pingcap/kvproto/pkg/tikvpb.(*BatchCommandsRequest).Marshal()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/tikvpb/tikvpb.pb.go:4179 +0x58
  google.golang.org/protobuf/internal/impl.legacyMarshal()
      /Users/cs/code/goread/pkg/mod/google.golang.org/protobuf@v1.30.0/internal/impl/legacy_message.go:402 +0x88
  google.golang.org/protobuf/proto.MarshalOptions.marshal()
      /Users/cs/code/goread/pkg/mod/google.golang.org/protobuf@v1.30.0/proto/encode.go:166 +0x2b4
  google.golang.org/protobuf/proto.MarshalOptions.MarshalAppend()
      /Users/cs/code/goread/pkg/mod/google.golang.org/protobuf@v1.30.0/proto/encode.go:125 +0x74
  github.com/golang/protobuf/proto.marshalAppend()
      /Users/cs/code/goread/pkg/mod/github.com/golang/protobuf@v1.5.3/proto/wire.go:40 +0x90
  github.com/golang/protobuf/proto.Marshal()
      /Users/cs/code/goread/pkg/mod/github.com/golang/protobuf@v1.5.3/proto/wire.go:23 +0x58
  google.golang.org/grpc/encoding/proto.codec.Marshal()
      /Users/cs/code/goread/pkg/mod/google.golang.org/grpc@v1.54.0/encoding/proto/proto.go:45 +0x5c
  google.golang.org/grpc/encoding/proto.(*codec).Marshal()
      <autogenerated>:1 +0x44
  google.golang.org/grpc.encode()
      /Users/cs/code/goread/pkg/mod/google.golang.org/grpc@v1.54.0/rpc_util.go:632 +0x4c
  google.golang.org/grpc.prepareMsg()
      /Users/cs/code/goread/pkg/mod/google.golang.org/grpc@v1.54.0/stream.go:1753 +0x11c
  google.golang.org/grpc.(*clientStream).SendMsg()
      /Users/cs/code/goread/pkg/mod/google.golang.org/grpc@v1.54.0/stream.go:874 +0x144
  github.com/pingcap/kvproto/pkg/tikvpb.(*tikvBatchCommandsClient).Send()
      /Users/cs/code/goread/pkg/mod/github.com/pingcap/kvproto@v0.0.0-20230904082117-ecdbf1f8c130/pkg/tikvpb/tikvpb.pb.go:2130 +0x54
  github.com/tikv/client-go/v2/internal/client.(*batchCommandsClient).send()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/client/client_batch.go:542 +0x3f0
  github.com/tikv/client-go/v2/internal/client.(*batchConn).getClientAndSend()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/client/client_batch.go:396 +0x64c
  github.com/tikv/client-go/v2/internal/client.(*batchConn).batchSendLoop()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/client/client_batch.go:350 +0x8c
  github.com/tikv/client-go/v2/internal/client.(*connArray).Init.func10()
      /Users/cs/code/goread/src/github.com/pingcap/client-go/internal/client/client.go:323 +0x9c
```

Patch:
diff --git a/internal/apicodec/codec.go b/internal/apicodec/codec.go
index 7945a5a..54b5840 100644
--- a/internal/apicodec/codec.go
+++ b/internal/apicodec/codec.go
@@ -92,28 +92,27 @@ func attachAPICtx(c Codec, req *tikvrpc.Request) *tikvrpc.Request {
 	// Shallow copy the request to avoid concurrent modification.
 	r := *req
 
-	ctx := &r.Context
-	ctx.ApiVersion = c.GetAPIVersion()
-	ctx.KeyspaceId = uint32(c.GetKeyspaceID())
+	r.Context.ApiVersion = c.GetAPIVersion()
+	r.Context.KeyspaceId = uint32(c.GetKeyspaceID())
 
 	switch r.Type {
 	case tikvrpc.CmdMPPTask:
 		mpp := *r.DispatchMPPTask()
 		// Shallow copy the meta to avoid concurrent modification.
 		meta := *mpp.Meta
-		meta.KeyspaceId = ctx.KeyspaceId
-		meta.ApiVersion = ctx.ApiVersion
+		meta.KeyspaceId = r.Context.KeyspaceId
+		meta.ApiVersion = r.Context.ApiVersion
 		mpp.Meta = &meta
 		r.Req = &mpp
 
 	case tikvrpc.CmdCompact:
 		compact := *r.Compact()
-		compact.KeyspaceId = ctx.KeyspaceId
-		compact.ApiVersion = ctx.ApiVersion
+		compact.KeyspaceId = r.Context.KeyspaceId
+		compact.ApiVersion = r.Context.ApiVersion
 		r.Req = &compact
 	}
 
-	tikvrpc.AttachContext(&r, ctx)
+	tikvrpc.AttachContext(&r, r.Context)
 
 	return &r
 }
diff --git a/internal/client/client_batch.go b/internal/client/client_batch.go
index a809109..95aed90 100644
--- a/internal/client/client_batch.go
+++ b/internal/client/client_batch.go
@@ -298,6 +298,9 @@ func (a *batchConn) fetchMorePendingRequests(
 
 const idleTimeout = 3 * time.Minute
 
+// BatchSendLoopPanicCounter is only used for testing.
+var BatchSendLoopPanicCounter int64 = 0
+
 func (a *batchConn) batchSendLoop(cfg config.TiKVClient) {
 	defer func() {
 		if r := recover(); r != nil {
@@ -305,7 +308,8 @@ func (a *batchConn) batchSendLoop(cfg config.TiKVClient) {
 			logutil.BgLogger().Error("batchSendLoop",
 				zap.Any("r", r),
 				zap.Stack("stack"))
-			logutil.BgLogger().Info("restart batchSendLoop")
+			atomic.AddInt64(&BatchSendLoopPanicCounter, 1)
+			logutil.BgLogger().Info("restart batchSendLoop", zap.Int64("count", atomic.LoadInt64(&BatchSendLoopPanicCounter)))
 			go a.batchSendLoop(cfg)
 		}
 	}()
diff --git a/tikvrpc/tikvrpc.go b/tikvrpc/tikvrpc.go
index 83119c3..13e8570 100644
--- a/tikvrpc/tikvrpc.go
+++ b/tikvrpc/tikvrpc.go
@@ -709,7 +709,9 @@ type MPPStreamResponse struct {
 
 // AttachContext sets the request context to the request,
 // return false if encounter unknown request type.
-func AttachContext(req *Request, ctx *kvrpcpb.Context) bool {
+// Parameter `rpcCtx` use `kvrpcpb.Context` instead of `*kvrpcpb.Context` to avoid concurrent modification by shallow copy.
+func AttachContext(req *Request, rpcCtx kvrpcpb.Context) bool {
+	ctx := &rpcCtx
 	switch req.Type {
 	case CmdGet:
 		req.Get().Context = ctx
@@ -807,13 +809,14 @@ func AttachContext(req *Request, ctx *kvrpcpb.Context) bool {
 
 // SetContext set the Context field for the given req to the specified ctx.
 func SetContext(req *Request, region *metapb.Region, peer *metapb.Peer) error {
-	ctx := &req.Context
 	if region != nil {
-		ctx.RegionId = region.Id
-		ctx.RegionEpoch = region.RegionEpoch
+		req.Context.RegionId = region.Id
+		req.Context.RegionEpoch = region.RegionEpoch
 	}
-	ctx.Peer = peer
-	if !AttachContext(req, ctx) {
+	req.Context.Peer = peer
+
+	// Shallow copy the context to avoid concurrent modification.
+	if !AttachContext(req, req.Context) {
 		return errors.Errorf("invalid request type %v", req.Type)
 	}
 	return nil


Test patch:
diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
index 6c97be8..be7ad62 100644
--- a/internal/locate/region_request_test.go
+++ b/internal/locate/region_request_test.go
@@ -37,8 +37,10 @@ package locate
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"net"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 	"unsafe"
@@ -733,3 +735,52 @@ func (s *testRegionRequestToSingleStoreSuite) TestKVReadTimeoutWithDisableBatchC
 	s.True(IsFakeRegionError(regionErr))
 	s.Equal(0, bo.GetTotalBackoffTimes()) // use kv read timeout will do fast retry, so backoff times should be 0.
 }
+
+func (s *testRegionRequestToSingleStoreSuite) TestBatchClientSendLoopPanic() {
+	// This test should use `go test -race` to run.
+	config.UpdateGlobal(func(conf *config.Config) {
+		conf.TiKVClient.MaxBatchSize = 128
+	})()
+
+	server, port := mock_server.StartMockTikvService()
+	s.True(port > 0)
+	rpcClient := client.NewRPCClient()
+	fnClient := &fnClient{fn: func(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (response *tikvrpc.Response, err error) {
+		return rpcClient.SendRequest(ctx, server.Addr(), req, timeout)
+	}}
+	tf := func(s *Store, bo *retry.Backoffer) livenessState {
+		return reachable
+	}
+
+	defer func() {
+		rpcClient.Close()
+		server.Stop()
+	}()
+
+	var wg sync.WaitGroup
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < 100; j++ {
+				ctx, cancel := context.WithCancel(context.Background())
+				bo := retry.NewBackofferWithVars(ctx, int(client.ReadTimeoutShort.Milliseconds()), nil)
+				region, err := s.cache.LocateRegionByID(bo, s.region)
+				s.Nil(err)
+				s.NotNil(region)
+				go func() {
+					// mock for kill query execution or timeout.
+					time.Sleep(time.Millisecond * time.Duration(rand.Intn(5)+1))
+					cancel()
+				}()
+				req := tikvrpc.NewRequest(tikvrpc.CmdCop, &coprocessor.Request{Data: []byte("a"), StartTs: 1})
+				regionRequestSender := NewRegionRequestSender(s.cache, fnClient)
+				regionRequestSender.regionCache.testingKnobs.mockRequestLiveness.Store((*livenessFunc)(&tf))
+				regionRequestSender.SendReq(bo, req, region.Region, client.ReadTimeoutShort)
+			}
+		}()
+	}
+	wg.Wait()
+	// batchSendLoop should not panic.
+	s.Equal(atomic.LoadInt64(&client.BatchSendLoopPanicCounter), int64(0))
+}