Repository: tikv/client-go
Pull request #1022 (2023-10-17 08:12:14)

Fix batch client batchSendLoop panic (#1021)
cherry-pick #1021 ,  #967


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 bf25a6d..918dbd2 100644
--- a/internal/client/client_batch.go
+++ b/internal/client/client_batch.go
@@ -296,14 +296,18 @@ 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 {
 			metrics.TiKVPanicCounter.WithLabelValues(metrics.LabelBatchSendLoop).Inc()
 			logutil.BgLogger().Error("batchSendLoop",
-				zap.Reflect("r", r),
+				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)
 		}
 	}()
@@ -430,7 +434,7 @@ func (s *batchCommandsStream) recv() (resp *tikvpb.BatchCommandsResponse, err er
 		if r := recover(); r != nil {
 			metrics.TiKVPanicCounter.WithLabelValues(metrics.LabelBatchRecvLoop).Inc()
 			logutil.BgLogger().Error("batchCommandsClient.recv panic",
-				zap.Reflect("r", r),
+				zap.Any("r", r),
 				zap.Stack("stack"))
 			err = errors.New("batch conn recv paniced")
 		}
@@ -595,7 +599,7 @@ func (c *batchCommandsClient) batchRecvLoop(cfg config.TiKVClient, tikvTransport
 		if r := recover(); r != nil {
 			metrics.TiKVPanicCounter.WithLabelValues(metrics.LabelBatchRecvLoop).Inc()
 			logutil.BgLogger().Error("batchRecvLoop",
-				zap.Reflect("r", r),
+				zap.Any("r", r),
 				zap.Stack("stack"))
 			logutil.BgLogger().Info("restart batchRecvLoop")
 			go c.batchRecvLoop(cfg, tikvTransportLayerLoad, streamClient)
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 1669fb8..6cdaa68 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -580,7 +580,7 @@ func (c *RegionCache) checkAndResolve(needCheckStores []*Store, needCheck func(*
 		r := recover()
 		if r != nil {
 			logutil.BgLogger().Error("panic in the checkAndResolve goroutine",
-				zap.Reflect("r", r),
+				zap.Any("r", r),
 				zap.Stack("stack trace"))
 		}
 	}()
@@ -2975,7 +2975,7 @@ func (c *RegionCache) checkAndUpdateStoreSlowScores() {
 		r := recover()
 		if r != nil {
 			logutil.BgLogger().Error("panic in the checkAndUpdateStoreSlowScores goroutine",
-				zap.Reflect("r", r),
+				zap.Any("r", r),
 				zap.Stack("stack trace"))
 		}
 	}()
diff --git a/tikvrpc/tikvrpc.go b/tikvrpc/tikvrpc.go
index 76f4e61..568588e 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
diff --git a/util/misc.go b/util/misc.go
index bd3e2b7..e324bf7 100644
--- a/util/misc.go
+++ b/util/misc.go
@@ -89,7 +89,7 @@ func WithRecovery(exec func(), recoverFn func(r interface{})) {
 		}
 		if r != nil {
 			logutil.BgLogger().Error("panic in the recoverable goroutine",
-				zap.Reflect("r", r),
+				zap.Any("r", r),
 				zap.Stack("stack trace"))
 		}
 	}()


Test patch:
diff --git a/internal/client/client_fail_test.go b/internal/client/client_fail_test.go
index 32d976e..c63c722 100644
--- a/internal/client/client_fail_test.go
+++ b/internal/client/client_fail_test.go
@@ -36,7 +36,6 @@ package client
 
 import (
 	"context"
-	"fmt"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -47,6 +46,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/config"
+	"github.com/tikv/client-go/v2/internal/client/mock_server"
 	"github.com/tikv/client-go/v2/tikvrpc"
 )
 
@@ -54,11 +54,11 @@ func TestPanicInRecvLoop(t *testing.T) {
 	require.Nil(t, failpoint.Enable("tikvclient/panicInFailPendingRequests", `panic`))
 	require.Nil(t, failpoint.Enable("tikvclient/gotErrorInRecvLoop", `return("0")`))
 
-	server, port := startMockTikvService()
+	server, port := mock_server.StartMockTikvService()
 	require.True(t, port > 0)
 	defer server.Stop()
 
-	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
+	addr := server.Addr()
 	rpcClient := NewRPCClient()
 	defer rpcClient.Close()
 	rpcClient.option.dialTimeout = time.Second / 3
@@ -81,10 +81,10 @@ func TestPanicInRecvLoop(t *testing.T) {
 }
 
 func TestRecvErrorInMultipleRecvLoops(t *testing.T) {
-	server, port := startMockTikvService()
+	server, port := mock_server.StartMockTikvService()
 	require.True(t, port > 0)
 	defer server.Stop()
-	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
+	addr := server.Addr()
 
 	// Enable batch and limit the connection count to 1 so that
 	// there is only one BatchCommands stream for each host or forwarded host.
diff --git a/internal/client/client_test.go b/internal/client/client_test.go
index d05632a..b589421 100644
--- a/internal/client/client_test.go
+++ b/internal/client/client_test.go
@@ -53,6 +53,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/config"
+	"github.com/tikv/client-go/v2/internal/client/mock_server"
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/metadata"
@@ -114,12 +115,12 @@ func TestCancelTimeoutRetErr(t *testing.T) {
 }
 
 func TestSendWhenReconnect(t *testing.T) {
-	server, port := startMockTikvService()
+	server, port := mock_server.StartMockTikvService()
 	require.True(t, port > 0)
 
 	rpcClient := NewRPCClient()
 	defer rpcClient.Close()
-	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
+	addr := server.Addr()
 	conn, err := rpcClient.getConnArray(addr, true)
 	assert.Nil(t, err)
 
@@ -238,7 +239,7 @@ func TestCollapseResolveLock(t *testing.T) {
 }
 
 func TestForwardMetadataByUnaryCall(t *testing.T) {
-	server, port := startMockTikvService()
+	server, port := mock_server.StartMockTikvService()
 	require.True(t, port > 0)
 	defer server.Stop()
 	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
@@ -253,7 +254,7 @@ func TestForwardMetadataByUnaryCall(t *testing.T) {
 
 	var checkCnt uint64
 	// Check no corresponding metadata if ForwardedHost is empty.
-	server.setMetaChecker(func(ctx context.Context) error {
+	server.SetMetaChecker(func(ctx context.Context) error {
 		atomic.AddUint64(&checkCnt, 1)
 		// gRPC may set some metadata by default, e.g. "context-type".
 		md, ok := metadata.FromIncomingContext(ctx)
@@ -281,7 +282,7 @@ func TestForwardMetadataByUnaryCall(t *testing.T) {
 	checkCnt = 0
 	forwardedHost := "127.0.0.1:6666"
 	// Check the metadata exists.
-	server.setMetaChecker(func(ctx context.Context) error {
+	server.SetMetaChecker(func(ctx context.Context) error {
 		atomic.AddUint64(&checkCnt, 1)
 		// gRPC may set some metadata by default, e.g. "context-type".
 		md, ok := metadata.FromIncomingContext(ctx)
@@ -306,10 +307,10 @@ func TestForwardMetadataByUnaryCall(t *testing.T) {
 }
 
 func TestForwardMetadataByBatchCommands(t *testing.T) {
-	server, port := startMockTikvService()
+	server, port := mock_server.StartMockTikvService()
 	require.True(t, port > 0)
 	defer server.Stop()
-	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
+	addr := server.Addr()
 
 	// Enable batch and limit the connection count to 1 so that
 	// there is only one BatchCommands stream for each host or forwarded host.
@@ -322,7 +323,7 @@ func TestForwardMetadataByBatchCommands(t *testing.T) {
 
 	var checkCnt uint64
 	setCheckHandler := func(forwardedHost string) {
-		server.setMetaChecker(func(ctx context.Context) error {
+		server.SetMetaChecker(func(ctx context.Context) error {
 			atomic.AddUint64(&checkCnt, 1)
 			md, ok := metadata.FromIncomingContext(ctx)
 			if forwardedHost == "" {
diff --git a/internal/client/mock_tikv_service_test.go b/internal/client/mock_server/mock_tikv_service.go
similarity index 63%
rename from internal/client/mock_tikv_service_test.go
rename to internal/client/mock_server/mock_tikv_service.go
index 5bcfcfc..392d3a5 100644
--- a/internal/client/mock_tikv_service_test.go
+++ b/internal/client/mock_server/mock_tikv_service.go
@@ -18,13 +18,14 @@
 // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/client/mock_tikv_service_test.go
 //
 
-package client
+package mock_server
 
 import (
 	"context"
 	"fmt"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/pingcap/kvproto/pkg/coprocessor"
@@ -35,9 +36,11 @@ import (
 	"google.golang.org/grpc"
 )
 
-type server struct {
+type MockServer struct {
 	tikvpb.TikvServer
 	grpcServer *grpc.Server
+	addr       string
+	running    int64 // 0: not running, 1: running
 	// metaChecker check the metadata of each request. Now only requests
 	// which need redirection set it.
 	metaChecker struct {
@@ -46,21 +49,28 @@ type server struct {
 	}
 }
 
-func (s *server) KvPrewrite(ctx context.Context, req *kvrpcpb.PrewriteRequest) (*kvrpcpb.PrewriteResponse, error) {
+func (s *MockServer) KvGet(ctx context.Context, req *kvrpcpb.GetRequest) (*kvrpcpb.GetResponse, error) {
+	if err := s.checkMetadata(ctx); err != nil {
+		return nil, err
+	}
+	return &kvrpcpb.GetResponse{}, nil
+}
+
+func (s *MockServer) KvPrewrite(ctx context.Context, req *kvrpcpb.PrewriteRequest) (*kvrpcpb.PrewriteResponse, error) {
 	if err := s.checkMetadata(ctx); err != nil {
 		return nil, err
 	}
 	return &kvrpcpb.PrewriteResponse{}, nil
 }
 
-func (s *server) CoprocessorStream(req *coprocessor.Request, ss tikvpb.Tikv_CoprocessorStreamServer) error {
+func (s *MockServer) CoprocessorStream(req *coprocessor.Request, ss tikvpb.Tikv_CoprocessorStreamServer) error {
 	if err := s.checkMetadata(ss.Context()); err != nil {
 		return err
 	}
 	return ss.Send(&coprocessor.Response{})
 }
 
-func (s *server) BatchCommands(ss tikvpb.Tikv_BatchCommandsServer) error {
+func (s *MockServer) BatchCommands(ss tikvpb.Tikv_BatchCommandsServer) error {
 	if err := s.checkMetadata(ss.Context()); err != nil {
 		return err
 	}
@@ -91,13 +101,13 @@ func (s *server) BatchCommands(ss tikvpb.Tikv_BatchCommandsServer) error {
 	}
 }
 
-func (s *server) setMetaChecker(check func(context.Context) error) {
+func (s *MockServer) SetMetaChecker(check func(context.Context) error) {
 	s.metaChecker.Lock()
 	s.metaChecker.check = check
 	s.metaChecker.Unlock()
 }
 
-func (s *server) checkMetadata(ctx context.Context) error {
+func (s *MockServer) checkMetadata(ctx context.Context) error {
 	s.metaChecker.Lock()
 	defer s.metaChecker.Unlock()
 	if s.metaChecker.check != nil {
@@ -106,32 +116,52 @@ func (s *server) checkMetadata(ctx context.Context) error {
 	return nil
 }
 
-func (s *server) Stop() {
+func (s *MockServer) IsRunning() bool {
+	return atomic.LoadInt64(&s.running) == 1
+}
+
+func (s *MockServer) Addr() string {
+	return s.addr
+}
+
+func (s *MockServer) Stop() {
 	s.grpcServer.Stop()
+	atomic.StoreInt64(&s.running, 0)
 }
 
-// Try to start a gRPC server and retrun the server instance and binded port.
-func startMockTikvService() (*server, int) {
+func (s *MockServer) Start(addr string) int {
+	if addr == "" {
+		addr = fmt.Sprintf("%s:%d", "127.0.0.1", 0)
+	}
 	port := -1
-	lis, err := net.Listen("tcp", fmt.Sprintf("%s:%d", "127.0.0.1", 0))
+	lis, err := net.Listen("tcp", addr)
 	if err != nil {
 		logutil.BgLogger().Error("can't listen", zap.Error(err))
 		logutil.BgLogger().Error("can't start mock tikv service because no available ports")
-		return nil, port
+		return port
 	}
 	port = lis.Addr().(*net.TCPAddr).Port
 
-	server := &server{}
-	s := grpc.NewServer(grpc.ConnectionTimeout(time.Minute))
-	tikvpb.RegisterTikvServer(s, server)
-	server.grpcServer = s
+	grpcServer := grpc.NewServer(grpc.ConnectionTimeout(time.Minute))
+	tikvpb.RegisterTikvServer(grpcServer, s)
+	s.grpcServer = grpcServer
 	go func() {
-		if err = s.Serve(lis); err != nil {
+		if err = grpcServer.Serve(lis); err != nil {
 			logutil.BgLogger().Error(
 				"can't serve gRPC requests",
 				zap.Error(err),
 			)
 		}
 	}()
+	atomic.StoreInt64(&s.running, 1)
+	s.addr = fmt.Sprintf("%s:%d", "127.0.0.1", port)
+	logutil.BgLogger().Info("mock server started", zap.String("addr", s.addr))
+	return port
+}
+
+// StartMockTikvService try to start a gRPC server and retrun the server instance and binded port.
+func StartMockTikvService() (*MockServer, int) {
+	server := &MockServer{}
+	port := server.Start("")
 	return server, port
 }
diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
index 42c8ceb..9329332 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"
@@ -52,8 +54,10 @@ import (
 	"github.com/pingcap/kvproto/pkg/tikvpb"
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/suite"
+	"github.com/tikv/client-go/v2/config"
 	"github.com/tikv/client-go/v2/internal/apicodec"
 	"github.com/tikv/client-go/v2/internal/client"
+	"github.com/tikv/client-go/v2/internal/client/mock_server"
 	"github.com/tikv/client-go/v2/internal/mockstore/mocktikv"
 	"github.com/tikv/client-go/v2/internal/retry"
 	"github.com/tikv/client-go/v2/tikvrpc"
@@ -693,3 +697,52 @@ func (s *testRegionRequestToSingleStoreSuite) TestStaleReadRetry() {
 	s.Nil(regionErr)
 	s.Equal([]byte("value"), resp.Resp.(*kvrpcpb.GetResponse).Value)
 }
+
+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))
+}